-
Notifications
You must be signed in to change notification settings - Fork 0
/
yolo_to_coco.py
147 lines (122 loc) · 5.69 KB
/
yolo_to_coco.py
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
import os
import json
from typing import List
import math
from glob import glob
from tqdm import tqdm
import cv2
import numpy as np
from PIL import Image
rootfolder = 'V:/media/data2/AGC/generated_frames/kang/train/'
image_root = os.path.join(rootfolder, 'images')
txt_path = os.path.join(rootfolder, 'labels')
dest_file = os.path.join(rootfolder, 'train.json')
task = 'train'
# task = 'test'
difficult = '-1'
# NIA_CLASSES = ['배경', '소형 선박', '대형 선박', '민간 항공기', '군용 항공기', '소형 승용차', '버스', '트럭', '기차', '크레인', '다리',
# '정유탱크', '댐', '운동경기장', '헬리패드', '원형 교차로']
# CLASS_NAMES_EN = ('background', 'small_ship', 'large_ship', 'civilian_aircraft', 'military_aircraft', 'small_car',
# 'bus', 'truck', 'train', 'crane', 'bridge', 'oil_tank', 'dam', 'athletic_field', 'helipad',
# 'roundabout')
# CLASS_DICT = {'background':0, 'small_ship':1, 'large_ship':2, 'civilian_aircraft':3, 'military_aircraft':4, 'small_car':5,
# 'bus':6, 'truck':7, 'train':8, 'crane':9, 'bridge':10, 'oil_tank':11, 'dam':12, 'athletic_field':13,
# 'helipad':14, 'roundabout':15}
NIA_CLASSES = ['배경', '사람']
CLASS_NAMES_EN = ('background', 'person')
CLASS_DICT = {'background':0, 'person':1}
# yolo to coco
if task == 'train':
data_dict = {}
data_dict['images'] = []
data_dict['categories'] = []
data_dict['annotations'] = []
for idex, name in enumerate(CLASS_NAMES_EN[1:]):
single_cat = {'id': idex + 1, 'name': name, 'supercategory': name}
data_dict['categories'].append(single_cat)
with open(dest_file, 'w') as f_out:
txts = os.listdir(txt_path)
obj_coords = list()
image_ids = list()
class_indices = list()
class_names = list()
for txt in tqdm(txts, desc='loading txt files'):
with open(os.path.join(txt_path, txt)) as f:
labels = f.readlines()
for label in labels:
label = label.split(' ')
obj_coords.append(label[1:])
image_ids.append(os.path.join(txt_path, txt).replace('.txt', '.jpg').split(os.path.sep)[-1])
class_indices.append(CLASS_DICT['person'])
class_names.append('person')
img_id_map = {img_file: i + 1 for i, img_file in enumerate(list(set(image_ids)))}
image_ids = [img_id_map[img_file] for img_file in image_ids]
# convert_labels_to_objects(coords, class_ids, class_names, image_ids, difficult=0, is_clockwise=False):
objs = list()
inst_count = 1
for coords, cls_id, cls_name, img_id in tqdm(zip(obj_coords, class_indices, class_names, image_ids),
desc="converting labels to objects"):
for i in range(len(coords)):
coords[i] = float(coords[i])
x_center_norm = coords[0]
y_center_norm = coords[1]
w_norm = coords[2]
h_norm = coords[3]
x_center = x_center_norm * 1920
y_center = y_center_norm * 1080
w = w_norm * 1920
h = h_norm * 1080
xmax = ((x_center * 2) + w) / 2
ymax = ((y_center * 2) + h) / 2
xmin = (x_center * 2) - xmax
ymin = (y_center * 2) - ymax
single_obj = {}
single_obj['difficult'] = difficult
single_obj['area'] = w * h
if cls_name in CLASS_NAMES_EN:
single_obj['category_id'] = CLASS_DICT[cls_name]
else:
continue
single_obj['segmentation'] = [[int(p) for p in coords]]
single_obj['iscrowd'] = 0
single_obj['bbox'] = (xmin, ymin, w, h)
single_obj['image_id'] = img_id
single_obj['id'] = inst_count
inst_count += 1
objs.append(single_obj)
data_dict['annotations'].extend(objs)
for imgfile in tqdm(img_id_map, desc='saving img info'):
imagepath = os.path.join(image_root, imgfile)
img_id = img_id_map[imgfile]
img = cv2.imread(imagepath)
height, width, c = img.shape
single_image = {}
single_image['file_name'] = imgfile
single_image['id'] = img_id
single_image['width'] = width
single_image['height'] = height
data_dict['images'].append(single_image)
json.dump(data_dict, f_out)
elif task == 'test':
data_dict = {}
data_dict['images'] = []
data_dict['categories'] = []
for idex, name in enumerate(CLASS_NAMES_EN[1:]):
single_cat = {'id': idex + 1, 'name': name, 'supercategory': name}
data_dict['categories'].append(single_cat)
with open(dest_file, 'w') as f_out:
image_ids = os.listdir(image_root)
img_id_map = {img_file: i + 1 for i, img_file in enumerate(list(set(image_ids)))}
image_ids = [img_id_map[img_file] for img_file in image_ids]
for imgfile in tqdm(img_id_map, desc='saving img info'):
imagepath = os.path.join(image_root, imgfile)
img_id = img_id_map[imgfile]
img = cv2.imread(imagepath)
height, width, c = img.shape
single_image = {}
single_image['file_name'] = imgfile
single_image['id'] = img_id
single_image['width'] = width
single_image['height'] = height
data_dict['images'].append(single_image)
json.dump(data_dict, f_out)