-
Notifications
You must be signed in to change notification settings - Fork 2
/
STANDARD2COCO_test.py
259 lines (210 loc) · 10.2 KB
/
STANDARD2COCO_test.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
import os
import json
import shutil
import datetime
from utils.utils import *
class Standard2COCO:
def __init__(self, src, dest, val_part):
if os.path.exists(src):
self.merged_standard_json = self.merge_standard_labels(src)
# self.json_path = src
else:
raise Exception('Given path doesn\'t exist.')
self.image_src = os.path.abspath(os.path.join(src, '..', 'images'))
# if not os.path.exists(self.image_src):
# raise Exception('Given images input directory doesn\'t exist.')
self.dest = dest
self.train_dest = os.path.join(dest, 'coco_annotations')
if not os.path.exists(self.train_dest):
os.makedirs(self.train_dest)
self.coco_file_train = os.path.join(self.train_dest, 'coco_file.json')
self.category_dict = dict()
self.category_item_id = 19
self.annotation_id = 7081902000000
self.coco = dict()
self.categories_dict = {1 : 'uav', 2: 'airplane', 3: 'bicycle', 4: 'bird', 5: 'boat',
6: 'bus', 7: 'car', 8: 'cat', 9: 'cow', 10: 'dog', 11: 'horse',
12: 'motorcycle', 13: 'person', 14: 'traffic_light', 15: 'train',
16: 'truck', 17: 'ufo', 18: 'helicopter'}
self.cat_names_dict = {'uav': 1, 'airplane': 2, 'bicycle': 3, 'bird': 4, 'boat': 5,
'bus': 6, 'car': 7, 'cat': 8, 'cow': 9, 'dog': 10, 'horse': 11,
'motorcycle': 12, 'person': 13, 'traffic_light': 14, 'train': 15,
'truck': 16, 'ufo': 17, 'helicopter': 18}
self.coco['categories'] = []
self.get_coco_dset('train')
self.val_part = int(val_part/10)
if self.val_part:
self.coco_file_val = os.path.join(dest, 'val', 'uav_instances_coco_val.json')
self.val_dest = os.path.join(dest, 'coco_annotations_val')
if not os.path.exists(self.val_dest):
os.makedirs(self.val_dest)
self.get_coco_dset('val')
@staticmethod
def correct_category(self, category_name):
""" Method info
"""
if category_name == 'duck':
category_name = 'bird'
elif category_name == 'auv' or category_name == 'drone':
category_name = 'uav'
elif category_name == 'fo':
category_name = 'ufo'
return category_name
def merge_standard_labels(self, src):
""" Method info
"""
all_jsons_under_src = get_all_files(src, 'json')
merged_standard_json = merge_json_files(all_jsons_under_src)
return merged_standard_json
def get_coco_dset(self, part):
self.coco[part] = dict()
self.coco[part]['images'] = []
self.coco[part]['type'] = 'instances'
self.coco[part]['annotations'] = []
self.coco[part]['categories'] = []
def convert_2_coco(self):
self.convert_2_coco_from_standard(self.json_path)
print('[PROGRESS] File conversion is successful.')
self.coco['train']['categories'] = self.coco['categories']
json.dump(self.coco['train'], open(self.coco_file_train, 'w'))
print('-'*80)
print('[INFO] COCO annotations have been created.')
print(' - Number of images in the dataset: {}'.format(len(self.coco['train']['images'])))
print(' - COCO annotation file has been saved as {}'.format(self.coco_file_train))
print('-'*80)
if self.val_part:
self.coco['val']['categories'] = self.coco['categories']
json.dump(self.coco['val'], open(self.coco_file_val, 'w'))
print('[INFO] Validation set has been created. Annotations and images have been saved.')
print(' - Number of images placed in validation set: {}'.format(len(self.coco['val']['images'])))
print(' - Validation images have been saved under {}'.format(self.val_dest))
print(' - Validation-set annotation file has been saved as {}'.format(self.coco_file_val))
print('-' * 80)
print('[INFO] Categories have been created.')
print(' - Number of categories created: {}'.format(len(self.coco['train']['categories'])))
print(' - Names and IDs of created categories: ' + ''.join(['{0}:{1} '.format(k, v) for k,
v in self.cat_names_dict.items()]))
print('-' * 80)
def create_cat_dict(self):
for key, value in self.cat_names_dict.items():
category_item = dict()
category_item['supercategory'] = 'none'
category_item['name'] = key
category_item['id'] = value
self.coco['categories'].append(category_item)
# print(self.coco['categories'])
def add_category_item(self, name):
category_item = dict()
category_item['supercategory'] = 'none'
category_item['id'] = self.cat_names_dict[name]
print(category_item['id'])
category_item['name'] = name
self.coco['categories'].append(category_item)
self.cat_names_dict[name] = self.cat_names_dict[name]
self.category_item_id += 1
return category_item['id']
def get_image_id(self):
now = datetime.datetime.now()
unique_id = '%s' % now.strftime("%Y%m%d%H%M%S%f")
return int(unique_id)
def add_image_item(self, file_name, img_obj, part):
if file_name is None:
raise Exception('Could not find filename info in .json file.')
if img_obj['width'] is None:
raise Exception('Could not find width info in .json file.')
if img_obj['height'] is None:
raise Exception('Could not find height info in .json file.')
image_id = self.get_image_id()
image_item = dict()
image_item['id'] = image_id
file_name_list = os.path.split(file_name)
print(img_obj.keys())
# image_item['file_name'] = os.path.join(file_name_list[-2], file_name_list[-1])
image_item['file_name'] = img_obj['image_path']
image_item['width'] = img_obj['width']
image_item['height'] = img_obj['height']
try:
image_item['daytime'] = img_obj['daytime']
image_item['illuminated'] = img_obj['illuminated']
image_item['fog'] = img_obj['fog']
image_item['rain'] = img_obj['rain']
image_item['sky'] = img_obj['sky']
image_item['snow'] = img_obj['snow']
image_item['spectrum'] = img_obj['spectrum']
image_item['video_id'] = img_obj['video_id']
image_item['resolution'] = img_obj['resolution']
image_item['depth'] = img_obj['depth']
except:
pass
self.coco[part]['images'].append(image_item)
return image_id
def add_annotation_item(self, image_id, category_id, bbox, part, type):
annotation_item = dict()
#annotation_item['segmentation'] = []
seg = []
# bbox[] is x,y,w,h
# left_top
seg.append(bbox[0])
seg.append(bbox[1])
# left_bottom
seg.append(bbox[0])
seg.append(bbox[1] + bbox[3])
# right_bottom
seg.append(bbox[0] + bbox[2])
seg.append(bbox[1] + bbox[3])
# right_top
seg.append(bbox[0] + bbox[2])
seg.append(bbox[1])
# annotation_item['segmentation'].append(seg)
annotation_item['area'] = bbox[2] * bbox[3]
annotation_item['iscrowd'] = 0
annotation_item['ignore'] = 0
annotation_item['image_id'] = image_id
annotation_item['bbox'] = bbox
annotation_item['category_id'] = category_id
annotation_item['type'] = type
self.annotation_id += 1
annotation_item['id'] = self.annotation_id
self.coco[part]['annotations'].append(annotation_item)
def convert_2_coco_from_standard(self, json_path):
with open(json_path) as f:
obj = json.load(f)
self.create_cat_dict()
for name in obj:
if len(obj[name]['labels']) == 0:
continue
img_path_list = obj[name]['image_path'].split('/')
img_name = os.path.join(img_path_list[-3], img_path_list[-2], img_path_list[-1])
part = 'train'
# Add image-item dict to uav_coco_dset['images'] list
current_img_id = self.add_image_item(img_name, obj[name], part)
# Assign the dict that contains bounding box list and category id to label_element and create COCO bbox list
for label_element in obj[name]['labels']:
print(current_img_id, label_element)
bbox = label_element['bbox']
category_name = label_element['category_name'].lower()
bbox_height = bbox[3] - bbox[1]
bbox_width = bbox[2] - bbox[0]
coco_bbox = [bbox[0], bbox[1], bbox_width, bbox_height]
# shutil.copyfile(os.path.join(self.image_src, name), os.path.join(self.dest, part, 'images',name))
# Correct category_name conflicts and check if the category name and id is already registered
category_name = self.correct_category(category_name)
# if category_name == 'duck':
# category_name = 'bird'
# elif category_name == 'auv' or category_name == 'drone':
# category_name = 'uav'
# elif category_name == 'fo':
# category_name = 'ufo'
# print(category_name)
if category_name not in self.cat_names_dict:
current_category_id = self.add_category_item(category_name)
else:
current_category_id = self.cat_names_dict[category_name]
# print(current_category_id)
# If the annotation belongs to a UAV add drone type
try:
drone_type = label_element['type']
except:
drone_type = None
# Add annotation-item dict to uav_coco_dset['annotations'] list
self.add_annotation_item(current_img_id, current_category_id, coco_bbox, part, drone_type)