Skip to content

S
stacked_hourglasses
Commit d7d0ffd6 by baihe
Options

updated

parent 322601f1
Showing with 22 additions and 295 deletions
data/MPII/dp.py deleted 100644 → 0
import cv2
import sys
import os
import torch
import numpy as np
import torch.utils.data
import myutils.img
class GenerateHeatmap():
def __init__(self, output_res, num_parts):
self.output_res = output_res
self.num_parts = num_parts
sigma = self.output_res/64
self.sigma = sigma
size = 6*sigma + 3
x = np.arange(0, size, 1, float)
y = x[:, np.newaxis]
x0, y0 = 3*sigma + 1, 3*sigma + 1
self.g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
def __call__(self, keypoints):
hms = np.zeros(shape = (self.num_parts, self.output_res, self.output_res), dtype = np.float32)
sigma = self.sigma
for p in keypoints:
for idx, pt in enumerate(p):
if pt[0] > 0:
x, y = int(pt[0]), int(pt[1])
if x<0 or y<0 or x>=self.output_res or y>=self.output_res:
continue
ul = int(x - 3*sigma - 1), int(y - 3*sigma - 1)
br = int(x + 3*sigma + 2), int(y + 3*sigma + 2)
c,d = max(0, -ul[0]), min(br[0], self.output_res) - ul[0]
a,b = max(0, -ul[1]), min(br[1], self.output_res) - ul[1]
cc,dd = max(0, ul[0]), min(br[0], self.output_res)
aa,bb = max(0, ul[1]), min(br[1], self.output_res)
hms[idx, aa:bb,cc:dd] = np.maximum(hms[idx, aa:bb,cc:dd], self.g[a:b,c:d])
return hms
class Dataset(torch.utils.data.Dataset):
def __init__(self, config, ds, index):
self.input_res = config['train']['input_res']
self.output_res = config['train']['output_res']
self.generateHeatmap = GenerateHeatmap(self.output_res, config['inference']['num_parts'])
self.ds = ds
self.index = index
def __len__(self):
return len(self.index)
def __getitem__(self, idx):
return self.loadImage(self.index[idx % len(self.index)])
def loadImage(self, idx):
ds = self.ds
## load + crop
orig_img = ds.get_img(idx)
path = ds.get_path(idx)
orig_keypoints = ds.get_kps(idx)
kptmp = orig_keypoints.copy()
c = ds.get_center(idx)
s = ds.get_scale(idx)
normalize = ds.get_normalized(idx)
cropped = myutils.img.crop(orig_img, c, s, (self.input_res, self.input_res))
for i in range(np.shape(orig_keypoints)[1]):
if orig_keypoints[0,i,0] > 0:
orig_keypoints[0,i,:2] = myutils.img.transform(orig_keypoints[0,i,:2], c, s, (self.input_res, self.input_res))
keypoints = np.copy(orig_keypoints)
## augmentation -- to be done to cropped image
height, width = cropped.shape[0:2]
center = np.array((width/2, height/2))
scale = max(height, width)/200
aug_rot=0
aug_rot = (np.random.random() * 2 - 1) * 30.
aug_scale = np.random.random() * (1.25 - 0.75) + 0.75
scale *= aug_scale
mat_mask = myutils.img.get_transform(center, scale, (self.output_res, self.output_res), aug_rot)[:2]
mat = myutils.img.get_transform(center, scale, (self.input_res, self.input_res), aug_rot)[:2]
inp = cv2.warpAffine(cropped, mat, (self.input_res, self.input_res)).astype(np.float32)/255
keypoints[:,:,0:2] = myutils.img.kpt_affine(keypoints[:,:,0:2], mat_mask)
if np.random.randint(2) == 0:
inp = self.preprocess(inp)
inp = inp[:, ::-1]
keypoints = keypoints[:, ds.flipped_parts['mpii']]
keypoints[:, :, 0] = self.output_res - keypoints[:, :, 0]
orig_keypoints = orig_keypoints[:, ds.flipped_parts['mpii']]
orig_keypoints[:, :, 0] = self.input_res - orig_keypoints[:, :, 0]
## set keypoints to 0 when were not visible initially (so heatmap all 0s)
for i in range(np.shape(orig_keypoints)[1]):
if kptmp[0,i,0] == 0 and kptmp[0,i,1] == 0:
keypoints[0,i,0] = 0
keypoints[0,i,1] = 0
orig_keypoints[0,i,0] = 0
orig_keypoints[0,i,1] = 0
## generate heatmaps on outres
heatmaps = self.generateHeatmap(keypoints)
return inp.astype(np.float32), heatmaps.astype(np.float32)
def preprocess(self, data):
# random hue and saturation
data = cv2.cvtColor(data, cv2.COLOR_RGB2HSV);
delta = (np.random.random() * 2 - 1) * 0.2
data[:, :, 0] = np.mod(data[:,:,0] + (delta * 360 + 360.), 360.)
delta_sature = np.random.random() + 0.5
data[:, :, 1] *= delta_sature
data[:,:, 1] = np.maximum( np.minimum(data[:,:,1], 1), 0 )
data = cv2.cvtColor(data, cv2.COLOR_HSV2RGB)
# adjust brightness
delta = (np.random.random() * 2 - 1) * 0.3
data += delta
# adjust contrast
mean = data.mean(axis=2, keepdims=True)
data = (data - mean) * (np.random.random() + 0.5) + mean
data = np.minimum(np.maximum(data, 0), 1)
return data
def init(config):
batchsize = config['train']['batchsize']
current_path = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_path)
import ref as ds
ds.init()
train, valid = ds.setup_val_split()
dataset = { key: Dataset(config, ds, data) for key, data in zip( ['train', 'valid'], [train, valid] ) }
use_data_loader = config['train']['use_data_loader']
loaders = {}
for key in dataset:
loaders[key] = torch.utils.data.DataLoader(dataset[key], batch_size=batchsize, shuffle=True, num_workers=config['train']['num_workers'], pin_memory=False)
def gen(phase):
batchsize = config['train']['batchsize']
batchnum = config['train']['{}_iters'.format(phase)]
loader = loaders[phase].__iter__()
for i in range(batchnum):
try:
imgs, heatmaps = next(loader)
except StopIteration:
# to avoid no data provided by dataloader
loader = loaders[phase].__iter__()
imgs, heatmaps = next(loader)
yield {
'imgs': imgs, #cropped and augmented
'heatmaps': heatmaps, #based on keypoints. 0 if not in img for joint
}
return lambda key: gen(key)
data/MPII/ref.py deleted 100644 → 0
import numpy as np
import h5py
from imageio import imread
import os
import time
def _isArrayLike(obj):
return hasattr(obj, '__iter__') and hasattr(obj, '__len__')
annot_dir = 'data/MPII/annot'
img_dir = 'data/MPII/images'
assert os.path.exists(img_dir)
mpii, num_examples_train, num_examples_val = None, None, None
import cv2
class MPII:
def __init__(self):
print('loading data...')
tic = time.time()
train_f = h5py.File(os.path.join(annot_dir, 'train.h5'), 'r')
val_f = h5py.File(os.path.join(annot_dir, 'valid.h5'), 'r')
self.t_center = train_f['center'][()]
t_scale = train_f['scale'][()]
t_part = train_f['part'][()]
t_visible = train_f['visible'][()]
t_normalize = train_f['normalize'][()]
t_imgname = [None] * len(self.t_center)
for i in range(len(self.t_center)):
t_imgname[i] = train_f['imgname'][i].decode('UTF-8')
self.v_center = val_f['center'][()]
v_scale = val_f['scale'][()]
v_part = val_f['part'][()]
v_visible = val_f['visible'][()]
v_normalize = val_f['normalize'][()]
v_imgname = [None] * len(self.v_center)
for i in range(len(self.v_center)):
v_imgname[i] = val_f['imgname'][i].decode('UTF-8')
self.center = np.append(self.t_center, self.v_center, axis=0)
self.scale = np.append(t_scale, v_scale)
self.part = np.append(t_part, v_part, axis=0)
self.visible = np.append(t_visible, v_visible, axis=0)
self.normalize = np.append(t_normalize, v_normalize)
self.imgname = t_imgname + v_imgname
print('Done (t={:0.2f}s)'.format(time.time()- tic))
def getAnnots(self, idx):
'''
returns h5 file for train or val set
'''
return self.imgname[idx], self.part[idx], self.visible[idx], self.center[idx], self.scale[idx], self.normalize[idx]
def getLength(self):
return len(self.t_center), len(self.v_center)
def init():
global mpii, num_examples_train, num_examples_val
mpii = MPII()
num_examples_train, num_examples_val = mpii.getLength()
# Part reference
parts = {'mpii':['rank', 'rkne', 'rhip',
'lhip', 'lkne', 'lank',
'pelv', 'thrx', 'neck', 'head',
'rwri', 'relb', 'rsho',
'lsho', 'lelb', 'lwri']}
flipped_parts = {'mpii':[5, 4, 3, 2, 1, 0, 6, 7, 8, 9, 15, 14, 13, 12, 11, 10]}
part_pairs = {'mpii':[[0, 5], [1, 4], [2, 3], [6], [7], [8], [9], [10, 15], [11, 14], [12, 13]]}
pair_names = {'mpii':['ankle', 'knee', 'hip', 'pelvis', 'thorax', 'neck', 'head', 'wrist', 'elbow', 'shoulder']}
def setup_val_split():
'''
returns index for train and validation imgs
index for validation images starts after that of train images
so that loadImage can tell them apart
'''
valid = [i+num_examples_train for i in range(num_examples_val)]
train = [i for i in range(num_examples_train)]
return np.array(train), np.array(valid)
def get_img(idx):
imgname, __, __, __, __, __ = mpii.getAnnots(idx)
path = os.path.join(img_dir, imgname)
img = imread(path)
return img
def get_path(idx):
imgname, __, __, __, __, __ = mpii.getAnnots(idx)
path = os.path.join(img_dir, imgname)
return path
def get_kps(idx):
__, part, visible, __, __, __ = mpii.getAnnots(idx)
kp2 = np.insert(part, 2, visible, axis=1)
kps = np.zeros((1, 16, 3))
kps[0] = kp2
return kps
def get_normalized(idx):
__, __, __, __, __, n = mpii.getAnnots(idx)
return n
def get_center(idx):
__, __, __, c, __, __ = mpii.getAnnots(idx)
return c
def get_scale(idx):
__, __, __, __, s, __ = mpii.getAnnots(idx)
return s
\ No newline at end of file
handler.py
import cv2 import cv2
import torch import torch
import data.MPII.ref as ds
import myutils.img import myutils.img
from myutils.group import HeatmapParser from myutils.group import HeatmapParser
from myutils.posture import * from myutils.posture import *
...@@ -11,6 +10,8 @@ CROOKED_HEAD_THRE=8 ...@@ -11,6 +10,8 @@ CROOKED_HEAD_THRE=8
# 斜肩的斜率阈值 # 斜肩的斜率阈值
OBLIQUE_SHOULDER_THRE=2 OBLIQUE_SHOULDER_THRE=2
_flipped_parts = {'mpii':[5, 4, 3, 2, 1, 0, 6, 7, 8, 9, 15, 14, 13, 12, 11, 10]}
# 人体骨骼各坐标点的下标对应图 # 人体骨骼各坐标点的下标对应图
class PersonPosture(): class PersonPosture():
...@@ -145,7 +146,7 @@ def inference(img, func, config, c, s): ...@@ -145,7 +146,7 @@ def inference(img, func, config, c, s):
for ii in tmp1: for ii in tmp1:
tmp[ii] = np.concatenate((tmp1[ii], tmp2[ii]), axis=0) tmp[ii] = np.concatenate((tmp1[ii], tmp2[ii]), axis=0)
det = tmp['det'][0, -1] + tmp['det'][1, -1, :, :, ::-1][ds.flipped_parts['mpii']] det = tmp['det'][0, -1] + tmp['det'][1, -1, :, :, ::-1][_flipped_parts['mpii']]
if det is None: if det is None:
return [], [] return [], []
det = det / 2 det = det / 2
...@@ -170,8 +171,21 @@ def main(): ...@@ -170,8 +171,21 @@ def main():
for i in range(ans.shape[0]): for i in range(ans.shape[0]):
pred.append({'keypoints': ans[i,:,:]}) pred.append({'keypoints': ans[i,:,:]})
return pred return pred
from urllib.request import urlopen
def url_to_image(url, readFlag=cv2.IMREAD_COLOR):
# download the image, convert it to a NumPy array, and then read
# it into OpenCV format
resp = urlopen(url)
image = np.asarray(bytearray(resp.read()), dtype="uint8")
image = cv2.imdecode(image, readFlag)
# return the image
return image
if __name__ == '__main__': if __name__ == '__main__':
image_path = "data/custom/O型腿3.jpeg" # image_path = "data/custom/O型腿3.jpeg"
image_path = 'https://gimg2.baidu.com/image_search/src=http%3A%2F%2Fgss0.baidu.com%2F-4o3dSag_xI4khGko9WTAnF6hhy%2Fzhidao%2Fpic%2Fitem%2Fac345982b2b7d0a2c8b09418ccef76094b369a3e.jpg&refer=http%3A%2F%2Fgss0.baidu.com&app=2002&size=f9999,10000&q=a80&n=0&g=0n&fmt=jpeg?sec=1640224763&t=142191b9f29e133d7c8c0d7af15c7879'
from train import init from train import init
func, config = init() func, config = init()
...@@ -188,8 +202,11 @@ if __name__ == '__main__': ...@@ -188,8 +202,11 @@ if __name__ == '__main__':
return pred return pred
input_res = 256 input_res = 256
orig_img = cv2.imread(image_path) # orig_img = cv2.imread(image_path)
orig_img_reverse = cv2.imread(image_path)[:,:,::-1] # orig_img_reverse = cv2.imread(image_path)[:,:,::-1]
orig_img = url_to_image(image_path)
orig_img_reverse = url_to_image(image_path)[:,:,::-1]
shape = orig_img_reverse.shape[0:2] shape = orig_img_reverse.shape[0:2]
......
myutils/constant.py deleted 100644 → 0
# 常量
MIN_VISIBLE=0.3
train.py
...@@ -106,14 +106,11 @@ def init(): ...@@ -106,14 +106,11 @@ def init():
task = importlib.import_module('task.pose') task = importlib.import_module('task.pose')
exp_path = os.path.join('exp', opt.exp) exp_path = os.path.join('exp', opt.exp)
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
config = task.__config__ config = task.__config__
try: os.makedirs(exp_path) try: os.makedirs(exp_path)
except FileExistsError: pass except FileExistsError: pass
config['opt'] = opt config['opt'] = opt
config['data_provider'] = importlib.import_module(config['data_provider'])
# 加载模型，func就是模型预测函数 # 加载模型，func就是模型预测函数
func = task.make_network(config) func = task.make_network(config)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment