-
Notifications
You must be signed in to change notification settings - Fork 20
/
Copy pathconvert2voc.py
134 lines (108 loc) · 4.49 KB
/
convert2voc.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
"""convert VOC format
+ VOC2012
+ JPEGImages
+ SegmentationClass
"""
import os, sys
import glob
import cv2
import random
import shutil
import numpy as np
import pandas as pd
import concurrent.futures
import pdb
random.seed(100)
userhome = os.path.expanduser('~')
src_datadir = os.path.join(userhome, 'data/dfsign')
src_traindir = src_datadir + '/train'
src_testdir = src_datadir + '/test'
src_annotation = src_datadir + '/train_label_fix.csv'
dest_datadir = os.path.join(userhome, 'data/dfsign/dfsign_region_voc')
image_dir = dest_datadir + '/JPEGImages'
segmentation_dir = dest_datadir + '/SegmentationClass'
list_folder = dest_datadir + '/ImageSets'
# copy train and test images
def _copy(src_image, dest_path):
shutil.copy(src_image, dest_path)
def _resize(src_image, dest_path):
img = cv2.imread(src_image)
height, width = img.shape[:2]
size = (int(width), int(height))
img = cv2.resize(img, size, interpolation=cv2.INTER_AREA)
name = os.path.basename(src_image)
cv2.imwrite(os.path.join(dest_path, name), img)
def get_box(label_df, im_name):
boxes = []
im_label = label_df[label_df.filename == im_name]
for index, row in im_label.iterrows():
xmin = min(row['X1'], row['X2'], row['X3'], row['X4'])
ymin = min(row['Y1'], row['Y2'], row['Y3'], row['Y4'])
xmax = max(row['X1'], row['X2'], row['X3'], row['X4'])
ymax = max(row['Y1'], row['Y2'], row['Y3'], row['Y4'])
boxes.append([xmin, ymin, xmax, ymax])
return np.array(boxes)
def _generate_mask(img_path, label_df):
try:
# image mask
img_name = os.path.split(img_path)[-1]
width, height = 3200, 1800
# chip mask 40x23, model input size 640x320
mask_w, mask_h = 50, 29
region_mask = np.zeros((mask_h, mask_w), dtype=np.uint8)
boxes = get_box(label_df, img_name)
for box in boxes:
xmin = np.floor(1.0 * box[0] / width * mask_w).astype(np.int32)
ymin = np.floor(1.0 * box[1] / height * mask_h).astype(np.int32)
xmax = np.floor(1.0 * box[2] / width * mask_w).astype(np.int32)
ymax = np.floor(1.0 * box[3] / height * mask_h).astype(np.int32)
ignore_xmin = xmin - 1 if xmin - 1 >= 0 else 0
ignore_ymin = ymin - 1 if ymin - 1 >= 0 else 0
ignore_xmax = xmax + 1 if xmax + 1 < mask_w else mask_w - 1
ignore_ymax = ymax + 1 if ymax + 1 < mask_h else mask_h - 1
region_mask[ignore_ymin : ignore_ymax+1, ignore_xmin : ignore_xmax+1] = 255
region_mask[ymin : ymax+1, xmin : xmax+1] = 1
maskname = os.path.join(segmentation_dir, img_name[:-4] + '_region.png')
cv2.imwrite(maskname, region_mask)
except Exception as e:
print(e)
if __name__ == "__main__":
if not os.path.exists(dest_datadir):
os.mkdir(dest_datadir)
os.mkdir(image_dir)
os.mkdir(segmentation_dir)
os.mkdir(list_folder)
train_list = glob.glob(src_traindir + '/*.jpg')
test_list = glob.glob(src_testdir + '/*.jpg')
random.shuffle(train_list)
train_list, val_list = train_list[:-1000], train_list[-1000:]
all_list = train_list + val_list
# read label
df = pd.read_csv(src_annotation)
trainval = True
test = False
if trainval:
with open(os.path.join(list_folder, 'train.txt'), 'w') as f:
temp = [os.path.basename(x)[:-4]+'\n' for x in train_list]
f.writelines(temp)
with open(os.path.join(list_folder, 'val.txt'), 'w') as f:
temp = [os.path.basename(x)[:-4]+'\n' for x in val_list]
f.writelines(temp)
print('copy train image....')
with concurrent.futures.ThreadPoolExecutor() as exector:
exector.map(_copy, train_list, [image_dir]*len(train_list))
print('copy val image....')
with concurrent.futures.ThreadPoolExecutor() as exector:
exector.map(_copy, val_list, [image_dir]*len(val_list))
print('generate mask...')
with concurrent.futures.ThreadPoolExecutor() as exector:
exector.map(_generate_mask, all_list, [df]*len(all_list))
print('done.')
if test:
with open(os.path.join(list_folder, 'test.txt'), 'w') as f:
temp = [os.path.basename(x)[:-4]+'\n' for x in test_list]
f.writelines(temp)
print('copy test image....')
with concurrent.futures.ThreadPoolExecutor() as exector:
exector.map(_copy, test_list, [image_dir]*len(test_list))
print('done.')