-
Notifications
You must be signed in to change notification settings - Fork 5
/
labels_conversion.py
186 lines (135 loc) · 6.04 KB
/
labels_conversion.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
import pandas as pd
import json
from pathlib import Path
import shutil
import numpy as np
def read_coco_labels(json_pth):
with open(str(json_pth)) as f:
coco_annotations = json.load(f)
return coco_annotations
def return_ann_list(image_id,coco_ann_df):
temp_df = coco_ann_df.loc[coco_ann_df['image_id']==image_id].copy()
l=[]
for i,row in temp_df.iterrows():
temp_dct = {}
temp_dct['category_id']=row['category_id']
temp_dct['bbox']=row['bbox']
temp_dct['segmentation']=row['segmentation']
l.append(temp_dct)
return l
def saving_for_src_images(src_df,annotations_df,
src_root,none_or_not,yolo_root,split,increment_by_one=False):
valid_splits={'train','valid','test'}
if split not in valid_splits:
raise ValueError(f"Split must be in {valid_splits}")
cnt=0
for filename in src_df['file_name'].to_list():
if none_or_not==None:
true_f_pth = src_root/'data'/filename
else:
true_f_pth = src_root/filename
img_id = src_df.loc[src_df['file_name']==filename,'id'].iloc[0]
height = src_df.loc[src_df['file_name']==filename,'height'].iloc[0]
width = src_df.loc[src_df['file_name']==filename,'width'].iloc[0]
ann_l = return_ann_list(img_id,annotations_df)
# Checking if exists:
to_check = yolo_root/split/'labels'
to_check.mkdir(exist_ok=True, parents=True)
try:
dest_name = f"{yolo_root}/{split}/labels/{filename.split('.jpg')[0]}.txt"
except:
dest_name = f"{yolo_root}/{split}/labels/{filename.split('.jpeg')[0]}.txt"
file_obj = open(str(dest_name),'a')
for ann in ann_l:
if increment_by_one:
curr_cat = ann['category_id'] - 1
else:
curr_cat = ann['category_id']
file_obj.write(f"{curr_cat} ")
for cnt,one_seg in enumerate(ann['segmentation'][0]):
if cnt%2==0:
one_val = format(one_seg/width,'.6f')
else:
one_val = format(one_seg/height,'.6f')
file_obj.write(f"{one_val} ")
file_obj.write('\n')
file_obj.close()
# For copying images
# Checking if exists:
to_check = yolo_root/split/'images'
to_check.mkdir(exist_ok=True, parents=True)
try:
shutil.copy(true_f_pth,str(f"{yolo_root}/{split}/images"))
except Exception as e:
print(f'image canot be copied due to:::{e}')
pass
print('---------------SUCCESS-------------')
#print(yolo_root)
def convert_to_yolo_seg_labels(coco_labels_pth,
yolo_labels_root_pth,
images_src_pth=None,
train_percent=0.7,
val_percent=0.20,
increment_by_one=False):
'''
Takes COCO format Instance segmentation labels, and images, and save YOLO format labels, in Train, Val, and Test splits, at desired folder
********
Keyword Arguments:
********
coco_labels_pth: type= str | pathlib.Path; required. the labels json file of COCO format instance segmentation labels
yolo_labels_root_pth: type= str | pathlib.Path; required. The Directory, where you want to store YOLO-Seg format labels. It can/can't exist
images_src_pth: Optional , if not passed, it means Src images are stored in "data" folder, at same level where labels json file was present.
Anyway, it is good to pass the path of images, to avoid confusion
train_percent: type: float; required Percentage of data to put in Train Split
val_percent: type: float; required Percentage of data to put in Val Split. Based on Train, and Val, Test split will be calculated
increment_by_one type: bool; If COCO-labels Categories id start with 1, then it should be True, as we need to subtract 1 from CAT-Ids for YOLO format
As they should start from 0
'''
if images_src_pth==None:
images_root= coco_labels_pth.parent
else:
images_root=Path(images_src_pth)
yolo_labels_root_pth=Path(yolo_labels_root_pth)
yolo_labels_root_pth.mkdir(exist_ok=True,parents=True)
yolo_train_pth = yolo_labels_root_pth/'train'
yolo_train_pth.mkdir(exist_ok=True)
yolo_valid_pth = yolo_labels_root_pth/'valid'
yolo_valid_pth.mkdir(exist_ok=True)
yolo_test_pth = yolo_labels_root_pth/'test'
yolo_test_pth.mkdir(exist_ok=True)
coco_annotations = read_coco_labels(coco_labels_pth)
coco_ann_imgs_df = pd.DataFrame(coco_annotations['images'])
coco_train_df,coco_valid_df,coco_test_df = np.split(coco_ann_imgs_df,
[int(train_percent * len(coco_ann_imgs_df)),
int((train_percent + val_percent) * len(coco_ann_imgs_df)),
])
coco_ann_df = pd.DataFrame(coco_annotations['annotations'])
#for split in ['train','valid','test']:
saving_for_src_images(coco_train_df,coco_ann_df,
images_root,images_src_pth,
yolo_labels_root_pth,
'train',increment_by_one)
print('Data and labels for Train Saved')
saving_for_src_images(coco_valid_df,coco_ann_df,
images_root,images_src_pth,
yolo_labels_root_pth,
'valid',increment_by_one)
print('Data and labels for Validation Saved')
saving_for_src_images(coco_test_df,coco_ann_df,
images_root,images_src_pth,
yolo_labels_root_pth,
'test',increment_by_one)
print('Data and labels for Test Saved')
# Now saving data.yaml
selected_labels=[i['name'] for i in coco_annotations['categories']]
yaml_obj = open(str(yolo_labels_root_pth/'data.yaml'),'a')
yaml_obj.write('names:')
yaml_obj.write('\n')
for label in selected_labels:
yaml_obj.write(f'- {label}')
yaml_obj.write('\n')
yaml_obj.write(f'nc: {len(selected_labels)} \n')
yaml_obj.write(f'train: {yolo_labels_root_pth}/train/images \n')
yaml_obj.write(f'val: {yolo_labels_root_pth}/valid/images \n')
yaml_obj.write(f'test: {yolo_labels_root_pth}/test/images \n')
yaml_obj.close()