def cropped_image_and_pose_coord(file_path, bbox, joints, config=config()):
file_path = file_path.numpy().decode("utf-8")
img = cv2.imread(os.path.join(config.image_path, file_path),
cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
if img is None:
print('cannot read ' + os.path.join(config.image_path, str(file_path)))
assert 0
x, y, w, h = bbox
aspect_ratio = config.input_shape[1] / config.input_shape[0]
center = np.array([x + w * 0.5, y + h * 0.5])
if w > aspect_ratio * h:
h = w / aspect_ratio
elif w < aspect_ratio * h:
w = h * aspect_ratio
scale = np.array([w, h]) * 1.25
rotation = 0
joints = np.array(joints).reshape(config.num_kps, 3).astype(np.float32)
#data augmentation
scale = scale * np.clip(np.random.randn() * config.scale_factor + 1,
1 - config.scale_factor, 1 + config.scale_factor)
rotation = np.clip(np.random.randn()*config.rotation_factor, -config.rotation_factor*2, config.rotation_factor*2)\
if random.random() <= 0.6 else 0
if random.random() <= 0.5:
img = img[:, ::-1, :]
center[0] = img.shape[1] - 1 - center[0]
joints[:, 0] = img.shape[1] - 1 - joints[:, 0]
for (q, w) in config.kps_symmetry:
joints_q, joints_w = joints[q, :].copy(), joints[w, :].copy()
joints[w, :], joints[q, :] = joints_q, joints_w
trans = get_affine_transform(
center, scale, rotation,
(config.input_shape[1], config.input_shape[0]))
cropped_img = cv2.warpAffine(
img,
trans, (config.input_shape[1], config.input_shape[0]),
flags=cv2.INTER_LINEAR)
for i in range(config.num_kps):
if joints[i, 2] > 0:
joints[i, :2] = affine_transform(joints[i, :2], trans)
joints[i, 2] *= ((joints[i, 0] >= 0) &
(joints[i, 0] < config.input_shape[1]) &
(joints[i, 1] >= 0) &
(joints[i, 1] < config.input_shape[0]))
target_coord = joints[:, :2].astype(np.int16)
target_valid = joints[:, 2]
return [cropped_img[:, :, ::-1], target_coord]
def render_gaussian_heatmap(coord, output_shape, config=config()):
x = tf.constant([i for i in range(output_shape[1])], tf.float32)
y = tf.constant([i for i in range(output_shape[0])], tf.float32)
xx, yy = tf.meshgrid(x, y)
xx = tf.reshape(xx, (*output_shape, 1))
yy = tf.reshape(yy, (*output_shape, 1))
x = tf.floor(
tf.reshape(coord[:, 0], [1, 1, config.num_kps]) /
config.input_shape[1] * output_shape[1] + 0.5)
y = tf.floor(
tf.reshape(coord[:, 1], [1, 1, config.num_kps]) /
config.input_shape[0] * output_shape[0] + 0.5)
heatmap = tf.exp(-(((xx - x) / config.sigma)**2) / 2 -
(((yy - y) / config.sigma)**2) / 2)
return heatmap * 255.
joints = [i['joints'] for i in val_data]
scores = [i['score'] for i in val_data]
ds = tf.data.Dataset.from_tensor_slices(
(image_paths, bbox, joints, scores)).shuffle(100)
root_path = os.path.abspath('.')
# first step
data = []
for sample in ds.take(1):
for i in sample:
data.append(i)
print(i.numpy())
# second step
config = config()
img = plt.imread(dbcfg.img_path + os.sep + data[0].numpy().decode("utf-8"))
cropped_img, target_coord = image_process.cropped_image_and_pose_coord(
data[0], data[1], data[2])
# plt.imshow(img)
# plt.show()
# plt.imshow(cropped_img)
# plt.show()
kps_array = np.array(data[2]).reshape((17, 3))
plt.imshow(dbcfg.vis_keypoints(img, kps_array.T))
plt.show()
imaga_path = dbcfg.img_path
class Dataset(object):
config = config()
dataset_name = 'COCO'
num_kps = 17
kps_names = [
'nose', 'l_eye', 'r_eye', 'l_ear', 'r_ear', 'l_shoulder', 'r_shoulder',
'l_elbow', 'r_elbow', 'l_wrist', 'r_wrist', 'l_hip', 'r_hip', 'l_knee',
'r_knee', 'l_ankle', 'r_ankle'
]
kps_symmetry = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12),
(13, 14), (15, 16)]
kps_lines = [(1, 2), (0, 1), (0, 2), (2, 4), (1, 3),
(6, 8), (8, 10), (5, 7), (7, 9), (12, 14), (14, 16), (11, 13),
(13, 15), (5, 6), (11, 12)]
human_det_path = osp.join(root_path, 'datasets', dataset_name, 'dets',
'human_detection.json') # human detection result
# osp.join(root_path, 'datasets', dataset_name, 'images')
img_path = config.image_path
train_annot_path = osp.join(config.dataset_path, 'annotations',
'person_keypoints_train2017.json')
val_annot_path = osp.join(config.dataset_path, 'annotations',
'person_keypoints_val2017.json')
test_annot_path = osp.join(config.dataset_path, 'annotations',
'image_info_test-dev2017.json')
def load_train_data(self, score=False):
coco = COCO(self.train_annot_path)
train_data = []
for aid in coco.anns.keys():
ann = coco.anns[aid]
imgname = 'train2017/' + coco.imgs[ann['image_id']]['file_name']
joints = ann['keypoints']
if (ann['image_id'] not in coco.imgs) or ann['iscrowd'] or (np.sum(
joints[2::3]) == 0) or (ann['num_keypoints'] == 0):
continue
# sanitize bboxes
x, y, w, h = ann['bbox']
img = coco.loadImgs(ann['image_id'])[0]
width, height = img['width'], img['height']
x1 = np.max((0, x))
y1 = np.max((0, y))
x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
if ann['area'] > 0 and x2 >= x1 and y2 >= y1:
bbox = [x1, y1, x2 - x1, y2 - y1]
else:
continue
if score:
data = dict(image_id=ann['image_id'],
imgpath=imgname,
bbox=bbox,
joints=joints,
score=1)
else:
data = dict(image_id=ann['image_id'],
imgpath=imgname,
bbox=bbox,
joints=joints)
train_data.append(data)
return train_data
def load_val_data_with_annot(self):
coco = COCO(self.val_annot_path)
val_data = []
for aid in coco.anns.keys():
ann = coco.anns[aid]
if ann['image_id'] not in coco.imgs:
continue
imgname = 'val2017/' + coco.imgs[ann['image_id']]['file_name']
bbox = ann['bbox']
joints = ann['keypoints']
data = dict(image_id=ann['image_id'],
imgpath=imgname,
bbox=bbox,
joints=joints,
score=1)
val_data.append(data)
return val_data
def load_annot(self, db_set):
if db_set == 'train':
coco = COCO(self.train_annot_path)
elif db_set == 'val':
coco = COCO(self.val_annot_path)
elif db_set == 'test':
coco = COCO(self.test_annot_path)
else:
print('Unknown db_set')
assert 0
return coco
def load_imgid(self, annot):
return annot.imgs
def imgid_to_imgname(self, annot, imgid, db_set):
imgs = annot.loadImgs(imgid)
imgname = [db_set + '2017/' + i['file_name'] for i in imgs]
return imgname
def evaluation(self, result, gt, result_dir, db_set):
result_path = osp.join(result_dir, 'result.json')
with open(result_path, 'w') as f:
json.dump(result, f)
result = gt.loadRes(result_path)
cocoEval = COCOeval(gt, result, iouType='keypoints')
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
result_path = osp.join(result_dir, 'result.pkl')
with open(result_path, 'wb') as f:
pickle.dump(cocoEval, f, 2)
print("Saved result file to " + result_path)
def vis_keypoints(self, img, kps, kp_thresh=0.4, alpha=1):
# Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(self.kps_lines) + 2)]
colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]
# Perform the drawing on a copy of the image, to allow for blending.
kp_mask = np.copy(img)
# Draw mid shoulder / mid hip first for better visualization.
mid_shoulder = (kps[:2, 5] + kps[:2, 6]) / 2.0
sc_mid_shoulder = np.minimum(kps[2, 5], kps[2, 6])
mid_hip = (kps[:2, 11] + kps[:2, 12]) / 2.0
sc_mid_hip = np.minimum(kps[2, 11], kps[2, 12])
nose_idx = 0
if sc_mid_shoulder > kp_thresh and kps[2, nose_idx] > kp_thresh:
cv2.line(kp_mask,
tuple(mid_shoulder.astype(np.int32)),
tuple(kps[:2, nose_idx].astype(np.int32)),
color=colors[len(self.kps_lines)],
thickness=2,
lineType=cv2.LINE_AA)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
cv2.line(kp_mask,
tuple(mid_shoulder.astype(np.int32)),
tuple(mid_hip.astype(np.int32)),
color=colors[len(self.kps_lines) + 1],
thickness=2,
lineType=cv2.LINE_AA)
# Draw the keypoints.
for l in range(len(self.kps_lines)):
i1 = self.kps_lines[l][0]
i2 = self.kps_lines[l][1]
p1 = kps[0, i1].astype(np.int32), kps[1, i1].astype(np.int32)
p2 = kps[0, i2].astype(np.int32), kps[1, i2].astype(np.int32)
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
cv2.line(kp_mask,
p1,
p2,
color=colors[l],
thickness=2,
lineType=cv2.LINE_AA)
if kps[2, i1] > kp_thresh:
cv2.circle(kp_mask,
p1,
radius=3,
color=colors[l],
thickness=-1,
lineType=cv2.LINE_AA)
if kps[2, i2] > kp_thresh:
cv2.circle(kp_mask,
p2,
radius=3,
color=colors[l],
thickness=-1,
lineType=cv2.LINE_AA)
# Blend the keypoints.
return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)