File size: 4,359 Bytes
e6e7cb5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import os, os.path as osp
import argparse
import numpy as np
import yaml
from tqdm import tqdm
def write_kp_labels(data):
assert not osp.isdir(osp.join(data['path'], data['labels'])), \
'Labels already generated. Remove or choose new name for labels.'
is_coco = 'coco' in data['path']
if is_coco:
from pycocotools.coco import COCO
else:
from crowdposetools.coco import COCO
splits = [osp.splitext(osp.split(data[s])[-1])[0] for s in ['train', 'val', 'test'] if s in data]
annotations = [osp.join(data['path'], data['{}_annotations'.format(s)]) for s in ['train', 'val', 'test'] if s in data]
test_split = [0 if s in ['train', 'val'] else 1 for s in ['train', 'val', 'test'] if s in data]
img_txt_dir = osp.join(data['path'], data['labels'], 'img_txt')
os.makedirs(img_txt_dir, exist_ok=True)
for split, annot, is_test in zip(splits, annotations, test_split):
img_txt_path = osp.join(img_txt_dir, '{}.txt'.format(split))
labels_path = osp.join(data['path'], '{}/{}'.format(data['labels'], split if is_coco else ''))
if not is_test:
os.makedirs(labels_path, exist_ok=True)
coco = COCO(annot)
if not is_test:
pbar = tqdm(coco.anns.keys(), total=len(coco.anns.keys()))
pbar.desc = 'Writing {} labels to {}'.format(split, labels_path)
for id in pbar:
a = coco.anns[id]
if a['image_id'] not in coco.imgs:
continue
if 'train' in split:
if is_coco and a['iscrowd']:
continue
img_info = coco.imgs[a['image_id']]
img_h, img_w = img_info['height'], img_info['width']
x, y, w, h = a['bbox']
xc, yc = x + w / 2, y + h / 2
xc /= img_w
yc /= img_h
w /= img_w
h /= img_h
keypoints = np.array(a['keypoints']).reshape([-1, 3])
# some of crowdpose keypoints are just outside image so clip to image extents
if not is_coco:
keypoints[:, 0] = np.clip(keypoints[:, 0], 0, img_w)
keypoints[:, 1] = np.clip(keypoints[:, 1], 0, img_h)
with open(osp.join(labels_path, '{}.txt'.format(osp.splitext(img_info['file_name'])[0])), 'a') as f:
# write person object
s = '{} {:.6f} {:.6f} {:.6f} {:.6f}'.format(0, xc, yc, w, h)
if data['pose_obj']:
for i, (x, y, v) in enumerate(keypoints):
s += ' {:.6f} {:.6f} {:.6f}'.format(x / img_w, y / img_h, v)
s += '\n'
f.write(s)
# write keypoint objects
for i, (x, y, v) in enumerate(keypoints):
if v:
if isinstance(data['kp_bbox'], list):
kp_bbox = data['kp_bbox'][i]
else:
kp_bbox = data['kp_bbox']
s = '{} {:.6f} {:.6f} {:.6f} {:.6f}'.format(
i + 1, x / img_w, y / img_h,
kp_bbox * max(img_h, img_w) / img_w,
kp_bbox * max(img_h, img_w) / img_h)
if data['pose_obj']:
for _ in range(keypoints.shape[0]):
s += ' {:.6f} {:.6f} {:.6f}'.format(0, 0, 0)
s += '\n'
f.write(s)
pbar.close()
with open(img_txt_path, 'w') as f:
for img_info in coco.imgs.values():
f.write(osp.join(data['path'], 'images',
'{}'.format(split if is_coco else ''),
img_info['file_name']) + '\n')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--data', default='data/coco-kp.yaml')
args = parser.parse_args()
assert osp.isfile(args.data), 'Data config file not found at {}'.format(args.data)
with open(args.data, 'rb') as f:
data = yaml.safe_load(f)
write_kp_labels(data) |