Hugging Face's logo

Spaces:
csmithxc
/
YOLO-World3
Sleeping

App Files Community
YOLO-World3
File size: 5,043 Bytes 
186701e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
 
# Copyright (c) OpenMMLab. All rights reserved.
"""Extracting subsets from coco2017 dataset.

This script is mainly used to debug and verify the correctness of the
program quickly.
The root folder format must be in the following format:

├── root
│   ├── annotations
│   ├── train2017
│   ├── val2017
│   ├── test2017

Currently, only support COCO2017. In the future will support user-defined
datasets of standard coco JSON format.

Example:
   python tools/misc/extract_subcoco.py ${ROOT} ${OUT_DIR} --num-img ${NUM_IMG}
"""

import argparse
import os.path as osp
import shutil

import mmengine
import numpy as np
from pycocotools.coco import COCO


# TODO: Currently only supports coco2017
def _process_data(args,
                  in_dataset_type: str,
                  out_dataset_type: str,
                  year: str = '2017'):
    assert in_dataset_type in ('train', 'val')
    assert out_dataset_type in ('train', 'val')

    int_ann_file_name = f'annotations/instances_{in_dataset_type}{year}.json'
    out_ann_file_name = f'annotations/instances_{out_dataset_type}{year}.json'

    ann_path = osp.join(args.root, int_ann_file_name)
    json_data = mmengine.load(ann_path)

    new_json_data = {
        'info': json_data['info'],
        'licenses': json_data['licenses'],
        'categories': json_data['categories'],
        'images': [],
        'annotations': []
    }

    area_dict = {
        'small': [0., 32 * 32],
        'medium': [32 * 32, 96 * 96],
        'large': [96 * 96, float('inf')]
    }

    coco = COCO(ann_path)

    # filter annotations by category ids and area range
    areaRng = area_dict[args.area_size] if args.area_size else []
    catIds = coco.getCatIds(args.classes) if args.classes else []
    ann_ids = coco.getAnnIds(catIds=catIds, areaRng=areaRng)
    ann_info = coco.loadAnns(ann_ids)

    # get image ids by anns set
    filter_img_ids = {ann['image_id'] for ann in ann_info}
    filter_img = coco.loadImgs(filter_img_ids)

    # shuffle
    np.random.shuffle(filter_img)

    num_img = args.num_img if args.num_img > 0 else len(filter_img)
    if num_img > len(filter_img):
        print(
            f'num_img is too big, will be set to {len(filter_img)}, '
            'because of not enough image after filter by classes and area_size'
        )
        num_img = len(filter_img)

    progress_bar = mmengine.ProgressBar(num_img)

    for i in range(num_img):
        file_name = filter_img[i]['file_name']
        image_path = osp.join(args.root, in_dataset_type + year, file_name)

        ann_ids = coco.getAnnIds(
            imgIds=[filter_img[i]['id']], catIds=catIds, areaRng=areaRng)
        img_ann_info = coco.loadAnns(ann_ids)

        new_json_data['images'].append(filter_img[i])
        new_json_data['annotations'].extend(img_ann_info)

        shutil.copy(image_path, osp.join(args.out_dir,
                                         out_dataset_type + year))

        progress_bar.update()

    mmengine.dump(new_json_data, osp.join(args.out_dir, out_ann_file_name))


def _make_dirs(out_dir):
    mmengine.mkdir_or_exist(out_dir)
    mmengine.mkdir_or_exist(osp.join(out_dir, 'annotations'))
    mmengine.mkdir_or_exist(osp.join(out_dir, 'train2017'))
    mmengine.mkdir_or_exist(osp.join(out_dir, 'val2017'))


def parse_args():
    parser = argparse.ArgumentParser(description='Extract coco subset')
    parser.add_argument('root', help='root path')
    parser.add_argument(
        'out_dir', type=str, help='directory where subset coco will be saved.')
    parser.add_argument(
        '--num-img',
        default=50,
        type=int,
        help='num of extract image, -1 means all images')
    parser.add_argument(
        '--area-size',
        choices=['small', 'medium', 'large'],
        help='filter ground-truth info by area size')
    parser.add_argument(
        '--classes', nargs='+', help='filter ground-truth by class name')
    parser.add_argument(
        '--use-training-set',
        action='store_true',
        help='Whether to use the training set when extract the training set. '
        'The training subset is extracted from the validation set by '
        'default which can speed up.')
    parser.add_argument('--seed', default=-1, type=int, help='seed')
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    assert args.out_dir != args.root, \
        'The file will be overwritten in place, ' \
        'so the same folder is not allowed !'

    seed = int(args.seed)
    if seed != -1:
        print(f'Set the global seed: {seed}')
        np.random.seed(int(args.seed))

    _make_dirs(args.out_dir)

    print('====Start processing train dataset====')
    if args.use_training_set:
        _process_data(args, 'train', 'train')
    else:
        _process_data(args, 'val', 'train')
    print('\n====Start processing val dataset====')
    _process_data(args, 'val', 'val')
    print(f'\n Result save to {args.out_dir}')


if __name__ == '__main__':
    main()