Files
insightface/detection/scrfd/tools/convert_crowdhuman.py
2021-07-26 20:09:57 +08:00

93 lines
3.5 KiB
Python

from __future__ import print_function
import cv2
import argparse
import os
import os.path as osp
import shutil
import numpy as np
import json
def parse_args():
parser = argparse.ArgumentParser(
description='convert crowdhuman dataset to scrfd format')
parser.add_argument('--raw', help='raw dataset dir')
parser.add_argument('--save', default='data/crowdhuman', help='save path')
args = parser.parse_args()
return args
def main():
args = parse_args()
raw_image_dir = osp.join(args.raw, 'Images')
for subset in ['train', 'val']:
save_image_dir = osp.join(args.save, subset, 'images')
if not osp.exists(save_image_dir):
os.makedirs(save_image_dir)
anno_file = osp.join(args.raw, 'annotation_%s.odgt'%subset)
fullbody_anno_file = osp.join(osp.join(args.save, subset, "label_fullbody.txt"))
head_anno_file = osp.join(osp.join(args.save, subset, "label_head.txt"))
fullbody_f = open(fullbody_anno_file, 'w')
head_f = open(head_anno_file, 'w')
for line in open(anno_file, 'r'):
data = json.loads(line)
img_id = data['ID']
img_name = "%s.jpg"%img_id
raw_image_file = osp.join(raw_image_dir, img_name)
target_image_file = osp.join(save_image_dir, img_name)
img = cv2.imread(raw_image_file)
print(raw_image_file, img.shape)
fullbody_f.write("# %s %d %d\n"%(img_name,img.shape[1],img.shape[0]))
head_f.write("# %s %d %d\n"%(img_name,img.shape[1],img.shape[0]))
shutil.copyfile(raw_image_file, target_image_file)
items = data['gtboxes']
for item in items:
fbox = item['fbox']
is_ignore = False
extra = item['extra']
if 'ignore' in extra:
is_ignore = extra['ignore']==1
bbox = np.array(fbox, dtype=np.float32)
bbox[2] += bbox[0]
bbox[3] += bbox[1]
if is_ignore:
fullbody_f.write("%.5f %.5f %.5f %.5f %d\n"%(bbox[0], bbox[1], bbox[2], bbox[3], is_ignore))
else:
vbox = item['vbox']
vbox = np.array(vbox, dtype=np.float32)
vbox[2] += vbox[0]
vbox[3] += vbox[1]
x1, y1, x2, y2 = vbox[0], vbox[1], vbox[2], vbox[3]
cx = (x1+x2)/2
cy = (y1+y2)/2
kps = np.ones( (5,3), dtype=np.float32)
kps[0,0] = x1
kps[0,1] = y1
kps[1,0] = x2
kps[1,1] = y1
kps[2,0] = cx
kps[2,1] = cy
kps[3,0] = x1
kps[3,1] = y2
kps[4,0] = x2
kps[4,1] = y2
kps_str = " ".join(["%.5f"%x for x in kps.flatten()])
fullbody_f.write("%.5f %.5f %.5f %.5f %s\n"%(bbox[0], bbox[1], bbox[2], bbox[3], kps_str))
hbox = item['hbox']
is_ignore = False
extra = item['head_attr']
if 'ignore' in extra:
is_ignore = extra['ignore']==1
bbox = np.array(hbox, dtype=np.float32)
bbox[2] += bbox[0]
bbox[3] += bbox[1]
head_f.write("%.5f %.5f %.5f %.5f %d\n"%(bbox[0], bbox[1], bbox[2], bbox[3], is_ignore))
fullbody_f.close()
head_f.close()
if __name__ == '__main__':
main()