Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| # This script converts MOT dataset into ReID dataset. | |
| # Official website of the MOT dataset: https://motchallenge.net/ | |
| # | |
| # Label format of MOT dataset: | |
| # GTs: | |
| # <frame_id> # starts from 1, | |
| # <instance_id>, <x1>, <y1>, <w>, <h>, | |
| # <conf> # conf is annotated as 0 if the object is ignored, | |
| # <class_id>, <visibility> | |
| # | |
| # DETs and Results: | |
| # <frame_id>, <instance_id>, <x1>, <y1>, <w>, <h>, <conf>, | |
| # <x>, <y>, <z> # for 3D objects | |
| # | |
| # Classes in MOT: | |
| # 1: 'pedestrian' | |
| # 2: 'person on vehicle' | |
| # 3: 'car' | |
| # 4: 'bicycle' | |
| # 5: 'motorbike' | |
| # 6: 'non motorized vehicle' | |
| # 7: 'static person' | |
| # 8: 'distractor' | |
| # 9: 'occluder' | |
| # 10: 'occluder on the ground', | |
| # 11: 'occluder full' | |
| # 12: 'reflection' | |
| # | |
| # USELESS classes and IGNORES classes will not be selected | |
| # into the dataset for reid model training. | |
| import argparse | |
| import os | |
| import os.path as osp | |
| import random | |
| import mmcv | |
| import numpy as np | |
| from mmengine.fileio import list_from_file | |
| from tqdm import tqdm | |
| USELESS = [3, 4, 5, 6, 9, 10, 11] | |
| IGNORES = [2, 7, 8, 12, 13] | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='Convert MOT dataset into ReID dataset.') | |
| parser.add_argument('-i', '--input', help='path of MOT data') | |
| parser.add_argument('-o', '--output', help='path to save ReID dataset') | |
| parser.add_argument( | |
| '--val-split', | |
| type=float, | |
| default=0.2, | |
| help='proportion of the validation dataset to the whole ReID dataset') | |
| parser.add_argument( | |
| '--vis-threshold', | |
| type=float, | |
| default=0.3, | |
| help='threshold of visibility for each person') | |
| parser.add_argument( | |
| '--min-per-person', | |
| type=int, | |
| default=8, | |
| help='minimum number of images for each person') | |
| parser.add_argument( | |
| '--max-per-person', | |
| type=int, | |
| default=1000, | |
| help='maxmum number of images for each person') | |
| return parser.parse_args() | |
| def main(): | |
| args = parse_args() | |
| if not osp.isdir(args.output): | |
| os.makedirs(args.output, exist_ok=True) | |
| in_folder = osp.join(args.input, 'train') | |
| video_names = os.listdir(in_folder) | |
| if 'MOT17' in in_folder: | |
| video_names = [ | |
| video_name for video_name in video_names if 'FRCNN' in video_name | |
| ] | |
| is_mot15 = True if 'MOT15' in in_folder else False | |
| for video_name in tqdm(video_names): | |
| # load video infos | |
| video_folder = osp.join(in_folder, video_name) | |
| infos = list_from_file(f'{video_folder}/seqinfo.ini') | |
| # video-level infos | |
| assert video_name == infos[1].strip().split('=')[1] | |
| raw_img_folder = infos[2].strip().split('=')[1] | |
| raw_img_names = os.listdir(f'{video_folder}/{raw_img_folder}') | |
| raw_img_names = sorted(raw_img_names) | |
| num_raw_imgs = int(infos[4].strip().split('=')[1]) | |
| assert num_raw_imgs == len(raw_img_names) | |
| reid_train_folder = osp.join(args.output, 'imgs') | |
| if not osp.exists(reid_train_folder): | |
| os.makedirs(reid_train_folder) | |
| gts = list_from_file(f'{video_folder}/gt/gt.txt') | |
| last_frame_id = -1 | |
| for gt in gts: | |
| gt = gt.strip().split(',') | |
| frame_id, ins_id = map(int, gt[:2]) | |
| ltwh = list(map(float, gt[2:6])) | |
| if is_mot15: | |
| class_id = 1 | |
| visibility = 1. | |
| else: | |
| class_id = int(gt[7]) | |
| visibility = float(gt[8]) | |
| if class_id in USELESS: | |
| continue | |
| elif class_id in IGNORES: | |
| continue | |
| elif visibility < args.vis_threshold: | |
| continue | |
| reid_img_folder = osp.join(reid_train_folder, | |
| f'{video_name}_{ins_id:06d}') | |
| if not osp.exists(reid_img_folder): | |
| os.makedirs(reid_img_folder) | |
| idx = len(os.listdir(reid_img_folder)) | |
| reid_img_name = f'{idx:06d}.jpg' | |
| if frame_id != last_frame_id: | |
| raw_img_name = raw_img_names[frame_id - 1] | |
| raw_img = mmcv.imread( | |
| f'{video_folder}/{raw_img_folder}/{raw_img_name}') | |
| last_frame_id = frame_id | |
| xyxy = np.asarray( | |
| [ltwh[0], ltwh[1], ltwh[0] + ltwh[2], ltwh[1] + ltwh[3]]) | |
| reid_img = mmcv.imcrop(raw_img, xyxy) | |
| mmcv.imwrite(reid_img, f'{reid_img_folder}/{reid_img_name}') | |
| reid_meta_folder = osp.join(args.output, 'meta') | |
| if not osp.exists(reid_meta_folder): | |
| os.makedirs(reid_meta_folder) | |
| reid_train_list = [] | |
| reid_val_list = [] | |
| reid_img_folder_names = sorted(os.listdir(reid_train_folder)) | |
| num_ids = len(reid_img_folder_names) | |
| num_train_ids = int(num_ids * (1 - args.val_split)) | |
| train_label, val_label = 0, 0 | |
| random.seed(0) | |
| for reid_img_folder_name in reid_img_folder_names[:num_train_ids]: | |
| reid_img_names = os.listdir( | |
| f'{reid_train_folder}/{reid_img_folder_name}') | |
| # ignore ids whose number of image is less than min_per_person | |
| if (len(reid_img_names) < args.min_per_person): | |
| continue | |
| # downsampling when there are too many images owned by one id | |
| if (len(reid_img_names) > args.max_per_person): | |
| reid_img_names = random.sample(reid_img_names, args.max_per_person) | |
| # training set | |
| for reid_img_name in reid_img_names: | |
| reid_train_list.append( | |
| f'{reid_img_folder_name}/{reid_img_name} {train_label}\n') | |
| train_label += 1 | |
| reid_entire_dataset_list = reid_train_list.copy() | |
| for reid_img_folder_name in reid_img_folder_names[num_train_ids:]: | |
| reid_img_names = os.listdir( | |
| f'{reid_train_folder}/{reid_img_folder_name}') | |
| # ignore ids whose number of image is less than min_per_person | |
| if (len(reid_img_names) < args.min_per_person): | |
| continue | |
| # downsampling when there are too many images owned by one id | |
| if (len(reid_img_names) > args.max_per_person): | |
| reid_img_names = random.sample(reid_img_names, args.max_per_person) | |
| for reid_img_name in reid_img_names: | |
| # validation set | |
| reid_val_list.append( | |
| f'{reid_img_folder_name}/{reid_img_name} {val_label}\n') | |
| reid_entire_dataset_list.append( | |
| f'{reid_img_folder_name}/{reid_img_name} ' | |
| f'{train_label + val_label}\n') | |
| val_label += 1 | |
| with open( | |
| osp.join(reid_meta_folder, | |
| f'train_{int(100 * (1 - args.val_split))}.txt'), | |
| 'w') as f: | |
| f.writelines(reid_train_list) | |
| with open( | |
| osp.join(reid_meta_folder, f'val_{int(100 * args.val_split)}.txt'), | |
| 'w') as f: | |
| f.writelines(reid_val_list) | |
| with open(osp.join(reid_meta_folder, 'train.txt'), 'w') as f: | |
| f.writelines(reid_entire_dataset_list) | |
| if __name__ == '__main__': | |
| main() | |