convert_crowdhuman_to_coco.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import os
  2. import numpy as np
  3. import json
  4. from PIL import Image
  5. DATA_PATH = 'datasets/crowdhuman/'
  6. OUT_PATH = DATA_PATH + 'annotations/'
  7. SPLITS = ['val', 'train']
  8. DEBUG = False
  9. def load_func(fpath):
  10. print('fpath', fpath)
  11. assert os.path.exists(fpath)
  12. with open(fpath,'r') as fid:
  13. lines = fid.readlines()
  14. records =[json.loads(line.strip('\n')) for line in lines]
  15. return records
  16. if __name__ == '__main__':
  17. if not os.path.exists(OUT_PATH):
  18. os.mkdir(OUT_PATH)
  19. for split in SPLITS:
  20. data_path = DATA_PATH + split
  21. out_path = OUT_PATH + '{}.json'.format(split)
  22. out = {'images': [], 'annotations': [], 'categories': [{'id': 1, 'name': 'person'}]}
  23. ann_path = DATA_PATH + 'annotation_{}.odgt'.format(split)
  24. anns_data = load_func(ann_path)
  25. image_cnt = 0
  26. ann_cnt = 0
  27. video_cnt = 0
  28. for ann_data in anns_data:
  29. image_cnt += 1
  30. file_path = DATA_PATH + 'CrowdHuman_{}/'.format(split) + '{}.jpg'.format(ann_data['ID'])
  31. im = Image.open(file_path)
  32. image_info = {'file_name': '{}.jpg'.format(ann_data['ID']),
  33. 'id': image_cnt,
  34. 'height': im.size[1],
  35. 'width': im.size[0]}
  36. out['images'].append(image_info)
  37. if split != 'test':
  38. anns = ann_data['gtboxes']
  39. for i in range(len(anns)):
  40. ann_cnt += 1
  41. fbox = anns[i]['fbox']
  42. ann = {'id': ann_cnt,
  43. 'category_id': 1,
  44. 'image_id': image_cnt,
  45. 'track_id': -1,
  46. 'bbox_vis': anns[i]['vbox'],
  47. 'bbox': fbox,
  48. 'area': fbox[2] * fbox[3],
  49. 'iscrowd': 1 if 'extra' in anns[i] and \
  50. 'ignore' in anns[i]['extra'] and \
  51. anns[i]['extra']['ignore'] == 1 else 0}
  52. out['annotations'].append(ann)
  53. print('loaded {} for {} images and {} samples'.format(split, len(out['images']), len(out['annotations'])))
  54. json.dump(out, open(out_path, 'w'))