Spaces:
Build error
Build error
| import pandas as pd | |
| import numpy as np | |
| import os.path as osp | |
| import os | |
| import json | |
| import tqdm | |
| import argparse | |
| from generate_mots_format_files import save_seqinfo | |
| def parse_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--motsynth-path', help="Directory path containing the 'annotations' directory with .json files") | |
| parser.add_argument('--save-path', help='Root file in which the new annoation files will be stored. If not provided, motsynth-root will be used') | |
| parser.add_argument('--save-dir', default='mot_annotations', help="name of directory within 'save-path'in which MOTS annotation files will be stored") | |
| args = parser.parse_args() | |
| if args.save_path is None: | |
| args.save_path = args.motsynth_path | |
| return args | |
| def main(args): | |
| ann_dir = osp.join(args.motsynth_path, 'annotations') | |
| mot_ann_dir = osp.join(args.save_path, args.save_dir) | |
| seqs = [f'{seq_num:03}' for seq_num in range(768) if seq_num not in (629, 757, 524, 652)] | |
| for seq in tqdm.tqdm(seqs): | |
| ann_path = osp.join(ann_dir, f'{seq}.json') | |
| with open(ann_path) as f: | |
| seq_ann = json.load(f) | |
| rows = [] | |
| img_id2frame = {im['id']: im['frame_n'] for im in seq_ann['images']} | |
| for ann in seq_ann['annotations']: | |
| # We compute the 3D location as the mid point between both feet keypoints in 3D | |
| kps = np.array(ann['keypoints_3d']).reshape(-1, 4) | |
| feet_pos_3d = kps[[-1, -4], :3].mean(axis = 0).round(4) | |
| row = {'frame': img_id2frame[ann['image_id']],# STARTS AT 0!!! | |
| 'id': ann['ped_id'], | |
| 'bb_left': ann['bbox'][0] + 1, # Make it 1-based?? | |
| 'bb_top': ann['bbox'][1] + 1, | |
| 'bb_width': ann['bbox'][2], | |
| 'bb_height': ann['bbox'][3], | |
| 'conf': 1 - ann['iscrowd'], | |
| 'class': 1 if ann['iscrowd'] == 0 else 8, # Class 8 means distractor. It is the one used by Trackeval as 'iscrowd' | |
| # We compute visibility as the proportion of visible keypoints | |
| 'vis': (np.array(ann['keypoints'])[2::3] ==2).mean().round(2), | |
| 'x': feet_pos_3d[0], | |
| 'y': feet_pos_3d[1], | |
| 'z': feet_pos_3d[2]} | |
| rows.append(row) | |
| # Save gt.txt file | |
| # Format in https://github.com/dendorferpatrick/MOTChallengeEvalKit/tree/master/MOT | |
| mot_ann = pd.DataFrame(rows, columns = ['frame', 'id', 'bb_left', 'bb_top', 'bb_width', 'bb_height', 'conf','class', 'vis', 'x', 'y', 'z']) | |
| gt_dir = osp.join(mot_ann_dir, seq, 'gt') | |
| os.makedirs(gt_dir, exist_ok=True) | |
| mot_ann.to_csv(osp.join(gt_dir, 'gt.txt'), header=None, index=None, sep=',') | |
| # Save seqinfo.ini | |
| seqinfo_path = osp.join(mot_ann_dir, seq, 'seqinfo.ini') | |
| save_seqinfo(seqinfo_path, info = seq_ann['info']) | |
| if __name__ =='__main__': | |
| args = parse_args() | |
| main(args) |