|
|
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]) |
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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) |
