|
| 1 | +# This script converts MOT labels into COCO style. |
| 2 | +# Offical website of the MOT dataset: https://motchallenge.net/ |
| 3 | +# |
| 4 | +# Label format of MOT dataset: |
| 5 | +# GTs: |
| 6 | +# <frame_id> # starts from 1 but COCO style starts from 0, |
| 7 | +# <instance_id>, <x1>, <y1>, <w>, <h>, |
| 8 | +# <conf> # conf is annotated as 0 if the object is ignored, |
| 9 | +# <class_id>, <visibility> |
| 10 | +# |
| 11 | +# DETs and Results: |
| 12 | +# <frame_id>, <instance_id>, <x1>, <y1>, <w>, <h>, <conf>, |
| 13 | +# <x>, <y>, <z> # for 3D objects |
| 14 | +# |
| 15 | +# Classes in MOT: |
| 16 | +# 1: 'pedestrian' |
| 17 | +# 2: 'person on vehicle' |
| 18 | +# 3: 'car' |
| 19 | +# 4: 'bicycle' |
| 20 | +# 5: 'motorbike' |
| 21 | +# 6: 'non motorized vehicle' |
| 22 | +# 7: 'static person' |
| 23 | +# 8: 'distractor' |
| 24 | +# 9: 'occluder' |
| 25 | +# 10: 'occluder on the ground', |
| 26 | +# 11: 'occluder full' |
| 27 | +# 12: 'reflection' |
| 28 | +# |
| 29 | +# USELESS classes are not included into the json file. |
| 30 | +# IGNORES classes are included with `ignore=True`. |
| 31 | +import argparse |
| 32 | +import os |
| 33 | +import os.path as osp |
| 34 | +from collections import defaultdict |
| 35 | + |
| 36 | +import mmcv |
| 37 | +import numpy as np |
| 38 | +from tqdm import tqdm |
| 39 | + |
| 40 | +USELESS = [3, 4, 5, 6, 9, 10, 11] |
| 41 | +IGNORES = [2, 7, 8, 12, 13] |
| 42 | + |
| 43 | + |
| 44 | +def parse_args(): |
| 45 | + parser = argparse.ArgumentParser( |
| 46 | + description='Convert MOT label and detections to COCO-VID format.') |
| 47 | + parser.add_argument('-i', '--input', help='path of MOT data') |
| 48 | + parser.add_argument( |
| 49 | + '-o', '--output', help='path to save coco formatted label file') |
| 50 | + parser.add_argument( |
| 51 | + '--convert-det', |
| 52 | + action='store_true', |
| 53 | + help='convert offical detection results.') |
| 54 | + parser.add_argument( |
| 55 | + '--split-train', |
| 56 | + action='store_true', |
| 57 | + help='split the train set into half-train and half-validate.') |
| 58 | + return parser.parse_args() |
| 59 | + |
| 60 | + |
| 61 | +def parse_gts(gts, is_mot15): |
| 62 | + outputs = defaultdict(list) |
| 63 | + for gt in gts: |
| 64 | + gt = gt.strip().split(',') |
| 65 | + frame_id, ins_id = map(int, gt[:2]) |
| 66 | + bbox = list(map(float, gt[2:6])) |
| 67 | + if is_mot15: |
| 68 | + conf = 1. |
| 69 | + class_id = 1 |
| 70 | + visibility = 1. |
| 71 | + else: |
| 72 | + conf = float(gt[6]) |
| 73 | + class_id = int(gt[7]) |
| 74 | + visibility = float(gt[8]) |
| 75 | + if class_id in USELESS: |
| 76 | + continue |
| 77 | + elif class_id in IGNORES: |
| 78 | + continue |
| 79 | + anns = dict( |
| 80 | + category_id=1, |
| 81 | + bbox=bbox, |
| 82 | + area=bbox[2] * bbox[3], |
| 83 | + iscrowd=False, |
| 84 | + visibility=visibility, |
| 85 | + mot_instance_id=ins_id, |
| 86 | + mot_conf=conf, |
| 87 | + mot_class_id=class_id) |
| 88 | + outputs[frame_id].append(anns) |
| 89 | + return outputs |
| 90 | + |
| 91 | + |
| 92 | +def parse_dets(dets): |
| 93 | + outputs = defaultdict(list) |
| 94 | + for det in dets: |
| 95 | + det = det.strip().split(',') |
| 96 | + frame_id, ins_id = map(int, det[:2]) |
| 97 | + assert ins_id == -1 |
| 98 | + bbox = list(map(float, det[2:7])) |
| 99 | + # [x1, y1, x2, y2] to be consistent with mmdet |
| 100 | + bbox = [ |
| 101 | + bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3], bbox[4] |
| 102 | + ] |
| 103 | + outputs[frame_id].append(bbox) |
| 104 | + |
| 105 | + return outputs |
| 106 | + |
| 107 | + |
| 108 | +def main(): |
| 109 | + args = parse_args() |
| 110 | + if not osp.exists(args.output): |
| 111 | + os.makedirs(args.output) |
| 112 | + |
| 113 | + sets = ['train', 'test'] |
| 114 | + if args.split_train: |
| 115 | + sets += ['half-train', 'half-val'] |
| 116 | + vid_id, img_id, ann_id = 1, 1, 1 |
| 117 | + |
| 118 | + for subset in sets: |
| 119 | + ins_id = 0 |
| 120 | + print(f'Converting {subset} set to COCO format') |
| 121 | + if 'half' in subset: |
| 122 | + in_folder = osp.join(args.input, 'train') |
| 123 | + else: |
| 124 | + in_folder = osp.join(args.input, subset) |
| 125 | + out_file = osp.join(args.output, f'{subset}_cocoformat.json') |
| 126 | + outputs = defaultdict(list) |
| 127 | + outputs['categories'] = [dict(id=1, name='pedestrian')] |
| 128 | + if args.convert_det: |
| 129 | + det_file = osp.join(args.output, f'{subset}_detections.pkl') |
| 130 | + detections = dict(bbox_results=dict()) |
| 131 | + video_names = os.listdir(in_folder) |
| 132 | + for video_name in tqdm(video_names): |
| 133 | + # basic params |
| 134 | + parse_gt = 'test' not in subset |
| 135 | + ins_maps = dict() |
| 136 | + # load video infos |
| 137 | + video_folder = osp.join(in_folder, video_name) |
| 138 | + infos = mmcv.list_from_file(f'{video_folder}/seqinfo.ini') |
| 139 | + # video-level infos |
| 140 | + assert video_name == infos[1].strip().split('=')[1] |
| 141 | + img_folder = infos[2].strip().split('=')[1] |
| 142 | + img_names = os.listdir(f'{video_folder}/{img_folder}') |
| 143 | + img_names = sorted(img_names) |
| 144 | + fps = int(infos[3].strip().split('=')[1]) |
| 145 | + num_imgs = int(infos[4].strip().split('=')[1]) |
| 146 | + assert num_imgs == len(img_names) |
| 147 | + width = int(infos[5].strip().split('=')[1]) |
| 148 | + height = int(infos[6].strip().split('=')[1]) |
| 149 | + video = dict( |
| 150 | + id=vid_id, |
| 151 | + name=video_name, |
| 152 | + fps=fps, |
| 153 | + width=width, |
| 154 | + height=height) |
| 155 | + # parse annotations |
| 156 | + if parse_gt: |
| 157 | + gts = mmcv.list_from_file(f'{video_folder}/gt/gt.txt') |
| 158 | + if 'MOT15' in video_folder: |
| 159 | + img2gts = parse_gts(gts, True) |
| 160 | + else: |
| 161 | + img2gts = parse_gts(gts, False) |
| 162 | + if args.convert_det: |
| 163 | + dets = mmcv.list_from_file(f'{video_folder}/det/det.txt') |
| 164 | + img2dets = parse_dets(dets) |
| 165 | + # make half sets |
| 166 | + if 'half' in subset: |
| 167 | + split_frame = num_imgs // 2 + 1 |
| 168 | + if 'train' in subset: |
| 169 | + img_names = img_names[:split_frame] |
| 170 | + elif 'val' in subset: |
| 171 | + img_names = img_names[split_frame:] |
| 172 | + else: |
| 173 | + raise ValueError( |
| 174 | + 'subset must be named with `train` or `val`') |
| 175 | + mot_frame_ids = [str(int(_.split('.')[0])) for _ in img_names] |
| 176 | + with open(f'{video_folder}/gt/gt_{subset}.txt', 'wt') as f: |
| 177 | + for gt in gts: |
| 178 | + if gt.split(',')[0] in mot_frame_ids: |
| 179 | + f.writelines(f'{gt}\n') |
| 180 | + # image and box level infos |
| 181 | + for frame_id, name in enumerate(img_names): |
| 182 | + img_name = osp.join(video_name, img_folder, name) |
| 183 | + mot_frame_id = int(name.split('.')[0]) |
| 184 | + image = dict( |
| 185 | + id=img_id, |
| 186 | + video_id=vid_id, |
| 187 | + file_name=img_name, |
| 188 | + height=height, |
| 189 | + width=width, |
| 190 | + frame_id=frame_id, |
| 191 | + mot_frame_id=mot_frame_id) |
| 192 | + if parse_gt: |
| 193 | + gts = img2gts[mot_frame_id] |
| 194 | + for gt in gts: |
| 195 | + gt.update(id=ann_id, image_id=img_id) |
| 196 | + mot_ins_id = gt['mot_instance_id'] |
| 197 | + if mot_ins_id in ins_maps: |
| 198 | + gt['instance_id'] = ins_maps[mot_ins_id] |
| 199 | + else: |
| 200 | + gt['instance_id'] = ins_id |
| 201 | + ins_maps[mot_ins_id] = ins_id |
| 202 | + ins_id += 1 |
| 203 | + outputs['annotations'].append(gt) |
| 204 | + ann_id += 1 |
| 205 | + if args.convert_det: |
| 206 | + dets = np.array(img2dets[mot_frame_id]) |
| 207 | + if dets.ndim == 1: |
| 208 | + assert len(dets) == 0 |
| 209 | + dets = np.zeros((0, 5)) |
| 210 | + detections['bbox_results'][img_name] = [dets] |
| 211 | + outputs['images'].append(image) |
| 212 | + img_id += 1 |
| 213 | + outputs['videos'].append(video) |
| 214 | + vid_id += 1 |
| 215 | + outputs['num_instances'] = ins_id |
| 216 | + print(f'{subset} has {ins_id} instances.') |
| 217 | + mmcv.dump(outputs, out_file) |
| 218 | + if args.convert_det: |
| 219 | + mmcv.dump(detections, det_file) |
| 220 | + print(f'Done! Saved as {out_file} and {det_file}') |
| 221 | + else: |
| 222 | + print(f'Done! Saved as {out_file}') |
| 223 | + |
| 224 | + |
| 225 | +if __name__ == '__main__': |
| 226 | + main() |
0 commit comments