-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
121 lines (97 loc) · 4.83 KB
/
data.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
import cv2
import numpy as np
import os
import random
class DataProcessor(object):
def __init__(self, out_rows, out_cols,
path='./data/',
frame_path="./data/leftImg8bit/",
mask_path="./data/gtFine/",
img_type="png", batch_size=8):
self.out_rows = out_rows
self.out_cols = out_cols
self.path = path
self.frame_path = frame_path
self.mask_path = mask_path
self.img_type = img_type
self.batch_size = batch_size
def reconstruct_folders(self):
print("reconstruct folders")
print('-' * 30)
train_size = int(sum(len(os.listdir(self.frame_path + 'train/' + name))
for name in os.listdir(self.frame_path + 'train/')) * 17 / 18)
os.mkdir(self.path + "train")
os.mkdir(self.path + "val")
os.mkdir(self.path + "test")
os.mkdir(self.path + "frames")
os.mkdir(self.path + "test/test_frames")
os.mkdir(self.path + "test/test_masks")
for name in ['train', 'val']:
os.mkdir(self.path + name + "/" + name + "_frames")
os.mkdir(self.path + name + "/" + name + "_masks")
count = 0
# same city names for frames and masks
city_names = os.listdir(self.frame_path + name)
for city_name in city_names:
image_names = os.listdir(self.frame_path + name + '/' + city_name)
for image_name in image_names:
frame_image = cv2.imread(self.frame_path + name + '/' + city_name + '/' + image_name)
mask_image = cv2.imread(self.mask_path + name + '/' + city_name + '/' + image_name.replace('leftImg8bit', 'gtFine_color'))
is_test = count >= train_size
cv2.imwrite(self.path + f'{"test" if is_test else name}/{"test" if is_test else name}_frames/frame_{count - train_size if is_test else count:04}.{self.img_type}',
cv2.resize(frame_image, (self.out_rows, self.out_cols), interpolation=cv2.INTER_NEAREST))
cv2.imwrite(self.path + f'{"test" if is_test else name}/{"test" if is_test else name}_masks/mask_{count - train_size if is_test else count:04}.{self.img_type}',
cv2.resize(mask_image, (self.out_rows, self.out_cols), interpolation=cv2.INTER_NEAREST))
count += 1
city_names = os.listdir(self.frame_path + "test")
count = 0
for city_name in city_names:
image_names = os.listdir(self.frame_path + 'test/' + city_name)
for image_name in image_names:
frame_image = cv2.imread(self.frame_path + 'test/' + city_name + '/' + image_name)
cv2.imwrite(self.path + 'test/' + f'frames/frame_{count:04}.{self.img_type}',
cv2.resize(frame_image, (self.out_rows, self.out_cols), interpolation=cv2.INTER_NEAREST))
count += 1
@staticmethod
def get_codes():
id_to_rgb = {}
rgb_to_id = {}
codes_path = './id_to_rgb.txt'
codes = open(codes_path, 'r')
lines = codes.readlines()
for line in lines:
numbers = [int(x) for x in line.replace('\n', '').split(',')]
rgb_to_id[tuple(numbers[1:])] = numbers[0]
id_to_rgb[numbers[0]] = tuple(numbers[1:])
return id_to_rgb, rgb_to_id
def get_generator(self, name, shape):
print(f"get {name} generator")
print('-' * 30)
rgb_to_id = self.get_codes()[1]
classes = len(rgb_to_id)
c = 0
frames_folder = f"{self.path}{name}/{name}_frames"
masks_folder = f"{self.path}{name}/{name}_masks"
n = os.listdir(frames_folder)
random.shuffle(n)
while True:
frame = np.zeros((self.batch_size, *shape)).astype('float')
mask = np.zeros((self.batch_size, shape[0], shape[1], classes))
for i in range(c, c + self.batch_size):
train_frame = cv2.imread(f"{frames_folder}/{n[i]}") / 255.
train_frame = cv2.resize(train_frame, (shape[0], shape[1]))
frame[i - c] = train_frame
train_mask = cv2.imread(f"{masks_folder}/{n[i].replace('frame', 'mask')}")
# ret_train_mask = np.zeros((shape[0], shape[1], classes))
for x in range(shape[0]):
for y in range(shape[1]):
mask[i - c, x, y, rgb_to_id[tuple(train_mask[x, y])]] = 1
# mask[i - c] = train_mask
c += self.batch_size
if c + self.batch_size >= len(os.listdir(frames_folder)):
c = 0
random.shuffle(n)
yield frame, mask
if __name__ == "__main__":
processor = DataProcessor(256, 256)
processor.reconstruct_folders()