Spaces:
Build error
Build error
| import os | |
| import os.path as osp | |
| import numpy as np | |
| import json | |
| import argparse | |
| import configparser | |
| import datetime | |
| import tqdm | |
| def parse_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| '--data-root', help="Path containing the dataset in a folder") | |
| parser.add_argument('--dataset', default='MOT17', | |
| help='Name of the dataset to be used. Should be either MOT17 or MOT20') | |
| parser.add_argument( | |
| '--save-dir', help='Root file in which the new annoation files will be stored. If not provided, data-root will be used') | |
| parser.add_argument('--split', default='train', | |
| help="Split processed within the dataset. Should be either 'train' or 'test'") | |
| parser.add_argument('--save-combined', default=True, action='store_true', | |
| help="Determines whether a separate .json file containing all sequence annotations will be created") | |
| parser.add_argument('--subsample', default=20, type=int, | |
| help="Frame subsampling rate. If e.g. 10 is selected, then we will select 1 in 20 frames") | |
| args = parser.parse_args() | |
| if args.save_dir is None: | |
| args.save_dir = osp.join(args.data_root, 'motcha_coco_annotations') | |
| return args | |
| def get_img_id(dataset, seq, fname): | |
| # Dataset num, seq num, frame num | |
| return int(f"{dataset[3:5]}{seq.split('-')[1]}{int(fname.split('.')[0]):06}") | |
| def read_seqinfo(path): | |
| cp = configparser.ConfigParser() | |
| cp.read(path) | |
| return {'height': int(cp.get('Sequence', 'imHeight')), | |
| 'width': int(cp.get('Sequence', 'imWidth')), | |
| 'fps': int(cp.get('Sequence', 'frameRate'))} | |
| def main(args): | |
| data_path = osp.join(args.data_root, args.dataset, args.split) | |
| seqs = os.listdir(data_path) | |
| if args.save_combined: | |
| comb_data = {'info': {'dataset': args.dataset, | |
| 'split': args.split, | |
| 'creation_date': datetime.datetime.today().strftime('%Y-%m-%d-%H-%M')}, | |
| 'images': [], | |
| 'annotations': [], | |
| 'categories': [{'id': 1, 'name': 'person', 'supercategory': 'person'}]} | |
| for seq in tqdm.tqdm(seqs): | |
| if args.dataset.lower() == 'mot17': | |
| # Choose an arbitrary set of detections for MOT17, annotations are the same | |
| if not seq.endswith('FRCNN'): | |
| continue | |
| print(f"Processing sequence {seq} in dataset {args.dataset}") | |
| seq_path = osp.join(data_path, seq) | |
| seqinfo_path = osp.join(seq_path, 'seqinfo.ini') | |
| gt_path = osp.join(seq_path, 'gt/gt.txt') | |
| im_dir = osp.join(seq_path, 'img1') | |
| if args.dataset.lower() == 'mot17': | |
| seq_ = '-'.join(seq.split('-')[:-1]) # Get rid of detector string | |
| else: | |
| seq_ = seq.copy() | |
| seqinfo = read_seqinfo(seqinfo_path) | |
| data = {'info': {'sequence': seq_, | |
| 'dataset': args.dataset, | |
| 'split': args.split, | |
| 'creation_date': datetime.datetime.today().strftime('%Y-%m-%d-%H-%M'), | |
| **seqinfo}, | |
| 'images': [], | |
| 'annotations': [], | |
| 'categories': [{'id': 1, 'name': 'person', 'supercategory': 'person'}]} | |
| # Load Bounding Box annotations | |
| gt = np.loadtxt(gt_path, dtype=np.float32, delimiter=',') | |
| keep_classes = [1, 2, 7, 8, 12] | |
| mask = np.isin(gt[:, 7], keep_classes) | |
| gt = gt[mask] | |
| anns = [{'ped_id': row[1], | |
| 'frame_n': row[0], | |
| 'category_id': 1, | |
| 'id': f"{get_img_id(args.dataset, seq, f'{int(row[0]):06}.jpg')}{int(row_i):010}", | |
| 'image_id': get_img_id(args.dataset, seq, f'{int(row[0]):06}.jpg'), | |
| # 'bbox': row[2:6].tolist(), | |
| # MOTCha annotations are 1-based | |
| 'bbox': [row[2] - 1, row[3] - 1, row[4], row[5]], | |
| 'area': row[4]*row[5], | |
| 'vis': row[8], | |
| 'iscrowd': 1 - row[6]} | |
| for row_i, row in enumerate(gt.astype(float)) if row[0] % args.subsample == 0] | |
| # Load Image information | |
| all_img_ids = list(set([aa['image_id'] for aa in anns])) | |
| imgs = [{'file_name': osp.join(args.dataset, args.split, seq, 'img1', fname), | |
| 'height': seqinfo['height'], | |
| 'width': seqinfo['width'], | |
| 'frame_n': int(fname.split('.')[0]), | |
| 'id': get_img_id(args.dataset, seq, fname)} | |
| for fname in os.listdir(im_dir) if get_img_id(args.dataset, seq, fname) in all_img_ids] | |
| assert len(set([im['id'] for im in imgs])) == len(imgs) | |
| data['images'].extend(imgs) | |
| assert len(str(imgs[0]['id'])) == len(str(anns[0]['image_id'])) | |
| data['annotations'].extend(anns) | |
| os.makedirs(args.save_dir, exist_ok=True) | |
| fname = f"{args.dataset}_{seq_}.json" if args.dataset not in seq_ else f"{seq_}.json" | |
| save_path = osp.join(args.save_dir, fname) | |
| with open(save_path, 'w') as f: | |
| json.dump(data, f) | |
| print(f"Saved result at {save_path}") | |
| if args.save_combined: | |
| comb_data['annotations'].extend(anns) | |
| comb_data['images'].extend(imgs) | |
| if args.save_combined: | |
| save_path = osp.join( | |
| args.save_dir, f"{args.dataset}_{args.split}.json") | |
| with open(save_path, 'w') as f: | |
| json.dump(comb_data, f) | |
| print(f"Saved combined result at {save_path}") | |
| if __name__ == '__main__': | |
| args = parse_args() | |
| main(args) | |