-
Notifications
You must be signed in to change notification settings - Fork 52
/
Copy pathprepare_dataset.py
147 lines (125 loc) · 6.11 KB
/
prepare_dataset.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
"""Perform data augmentation and preprocessing."""
import argparse
import json
import math
import os
import random
import cv2 as cv
import numpy as np
def get_parser():
"""Return argument parser for generating dataset."""
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True,
choices=['trainval', 'test'],
help="Generate trainval or test dataset.")
parser.add_argument('--val_prop', type=float, default=0.1,
help="The proportion of val sample in trainval.")
parser.add_argument('--label_directory', required=True,
help="The location of label directory.")
parser.add_argument('--image_directory', required=True,
help="The location of image directory.")
parser.add_argument('--output_directory', required=True,
help="The location of output directory.")
return parser
def boundary_check(centralied_marks):
"""Check situation that marking point appears too near to border."""
for mark in centralied_marks:
if mark[0] < -260 or mark[0] > 260 or mark[1] < -260 or mark[1] > 260:
return False
return True
def overlap_check(centralied_marks):
"""Check situation that multiple marking points appear in same cell."""
for i in range(len(centralied_marks) - 1):
i_x = centralied_marks[i, 0]
i_y = centralied_marks[i, 1]
for j in range(i + 1, len(centralied_marks)):
j_x = centralied_marks[j, 0]
j_y = centralied_marks[j, 1]
if abs(j_x - i_x) < 600 / 16 and abs(j_y - i_y) < 600 / 16:
return False
return True
def generalize_marks(centralied_marks):
"""Convert coordinate to [0, 1] and calculate direction label."""
generalized_marks = []
for mark in centralied_marks:
xval = (mark[0] + 300) / 600
yval = (mark[1] + 300) / 600
direction = math.atan2(mark[3] - mark[1], mark[2] - mark[0])
generalized_marks.append([xval, yval, direction, mark[4]])
return generalized_marks
def write_image_and_label(name, image, centralied_marks, name_list):
"""Write image and label with given name."""
name_list.append(os.path.basename(name))
print("Processing NO.%d samples: %s..." % (len(name_list), name_list[-1]))
image = cv.resize(image, (512, 512))
cv.imwrite(name + '.jpg', image, [int(cv.IMWRITE_JPEG_QUALITY), 100])
with open(name + '.json', 'w') as file:
json.dump(generalize_marks(centralied_marks), file)
def rotate_vector(vector, angle_degree):
"""Rotate a vector with given angle in degree."""
angle_rad = math.pi * angle_degree / 180
xval = vector[0]*math.cos(angle_rad) + vector[1]*math.sin(angle_rad)
yval = -vector[0]*math.sin(angle_rad) + vector[1]*math.cos(angle_rad)
return xval, yval
def rotate_centralized_marks(centralied_marks, angle_degree):
"""Rotate centralized marks with given angle in degree."""
rotated_marks = centralied_marks.copy()
for i in range(centralied_marks.shape[0]):
mark = centralied_marks[i]
rotated_marks[i, 0:2] = rotate_vector(mark[0:2], angle_degree)
rotated_marks[i, 2:4] = rotate_vector(mark[2:4], angle_degree)
return rotated_marks
def rotate_image(image, angle_degree):
"""Rotate image with given angle in degree."""
rows, cols, _ = image.shape
rotation_matrix = cv.getRotationMatrix2D((rows/2, cols/2), angle_degree, 1)
return cv.warpAffine(image, rotation_matrix, (rows, cols))
def generate_dataset(args):
"""Generate dataset according to arguments."""
if args.dataset == 'trainval':
val_directory = os.path.join(args.output_directory, 'val')
args.output_directory = os.path.join(args.output_directory, 'train')
elif args.dataset == 'test':
args.output_directory = os.path.join(args.output_directory, 'test')
os.makedirs(args.output_directory, exist_ok=True)
name_list = []
for label_file in os.listdir(args.label_directory):
name = os.path.splitext(label_file)[0]
image = cv.imread(os.path.join(args.image_directory, name + '.jpg'))
with open(os.path.join(args.label_directory, label_file), 'r') as file:
label = json.load(file)
centralied_marks = np.array(label['marks'])
if len(centralied_marks.shape) < 2:
centralied_marks = np.expand_dims(centralied_marks, axis=0)
centralied_marks[:, 0:4] -= 300.5
if boundary_check(centralied_marks) or args.dataset == 'test':
output_name = os.path.join(args.output_directory, name)
write_image_and_label(output_name, image,
centralied_marks, name_list)
if args.dataset == 'test':
continue
for angle in range(5, 360, 5):
rotated_marks = rotate_centralized_marks(centralied_marks, angle)
if boundary_check(rotated_marks) and overlap_check(rotated_marks):
rotated_image = rotate_image(image, angle)
output_name = os.path.join(
args.output_directory, name + '_' + str(angle))
write_image_and_label(
output_name, rotated_image, rotated_marks, name_list)
if args.dataset == 'trainval':
print("Dividing training set and validation set...")
val_idx = random.sample(list(range(len(name_list))),
int(round(len(name_list)*args.val_prop)))
val_samples = [name_list[idx] for idx in val_idx]
os.makedirs(val_directory, exist_ok=True)
for val_sample in val_samples:
train_directory = args.output_directory
image_src = os.path.join(train_directory, val_sample + '.jpg')
label_src = os.path.join(train_directory, val_sample + '.json')
image_dst = os.path.join(val_directory, val_sample + '.jpg')
label_dst = os.path.join(val_directory, val_sample + '.json')
os.rename(image_src, image_dst)
os.rename(label_src, label_dst)
print("Done.")
if __name__ == '__main__':
generate_dataset(get_parser().parse_args())