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. | |
| import argparse | |
| import json | |
| import os | |
| import os.path as osp | |
| from collections import defaultdict | |
| import mmengine | |
| from PIL import Image | |
| from tqdm import tqdm | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='CrowdHuman to COCO Video format') | |
| parser.add_argument( | |
| '-i', | |
| '--input', | |
| help='root directory of CrowdHuman annotations', | |
| ) | |
| parser.add_argument( | |
| '-o', | |
| '--output', | |
| help='directory to save coco formatted label file', | |
| ) | |
| return parser.parse_args() | |
| def load_odgt(filename): | |
| with open(filename, 'r') as f: | |
| lines = f.readlines() | |
| data_infos = [json.loads(line.strip('\n')) for line in lines] | |
| return data_infos | |
| def convert_crowdhuman(ann_dir, save_dir, mode='train'): | |
| """Convert CrowdHuman dataset in COCO style. | |
| Args: | |
| ann_dir (str): The path of CrowdHuman dataset. | |
| save_dir (str): The path to save annotation files. | |
| mode (str): Convert train dataset or validation dataset. Options are | |
| 'train', 'val'. Default: 'train'. | |
| """ | |
| assert mode in ['train', 'val'] | |
| records = dict(img_id=1, ann_id=1) | |
| outputs = defaultdict(list) | |
| outputs['categories'] = [dict(id=1, name='pedestrian')] | |
| data_infos = load_odgt(osp.join(ann_dir, f'annotation_{mode}.odgt')) | |
| for data_info in tqdm(data_infos): | |
| img_name = osp.join('Images', f"{data_info['ID']}.jpg") | |
| img = Image.open(osp.join(ann_dir, mode, img_name)) | |
| width, height = img.size[:2] | |
| image = dict( | |
| file_name=img_name, | |
| height=height, | |
| width=width, | |
| id=records['img_id']) | |
| outputs['images'].append(image) | |
| if mode != 'test': | |
| for ann_info in data_info['gtboxes']: | |
| bbox = ann_info['fbox'] | |
| if 'extra' in ann_info and 'ignore' in ann_info[ | |
| 'extra'] and ann_info['extra']['ignore'] == 1: | |
| iscrowd = True | |
| else: | |
| iscrowd = False | |
| ann = dict( | |
| id=records['ann_id'], | |
| image_id=records['img_id'], | |
| category_id=outputs['categories'][0]['id'], | |
| vis_bbox=ann_info['vbox'], | |
| bbox=bbox, | |
| area=bbox[2] * bbox[3], | |
| iscrowd=iscrowd) | |
| outputs['annotations'].append(ann) | |
| records['ann_id'] += 1 | |
| records['img_id'] += 1 | |
| if not osp.isdir(save_dir): | |
| os.makedirs(save_dir) | |
| mmengine.dump(outputs, osp.join(save_dir, f'crowdhuman_{mode}.json')) | |
| print(f'-----CrowdHuman {mode} set------') | |
| print(f'total {records["img_id"] - 1} images') | |
| if mode != 'test': | |
| print(f'{records["ann_id"] - 1} pedestrians are annotated.') | |
| print('-----------------------') | |
| def main(): | |
| args = parse_args() | |
| convert_crowdhuman(args.input, args.output, mode='train') | |
| convert_crowdhuman(args.input, args.output, mode='val') | |
| if __name__ == '__main__': | |
| main() | |