-
Notifications
You must be signed in to change notification settings - Fork 2
/
utils.py
118 lines (95 loc) · 4.29 KB
/
utils.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
import cv2
import os
import numpy as np
import random
import glob
def make_data_path(train_list, test_list):
train = glob.glob(os.path.join(train_list, '*'))
test = glob.glob(os.path.join(test_list, '*'))
return train, test
class data_augmentation(object):
def __init__(self, train_list, test_list):
self.train_list = train_list
self.test_list = test_list
def open_image(self, path, width, height, isCrop=False, isResize=True, isflip=False, isRotate=False, angle=0):
# print(self.data_list[class_list][subclass_list][start])
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if isCrop:
img = img[20:198, 0:178]
if isResize:
img = cv2.resize(img, (width, height), interpolation=cv2.INTER_LINEAR)
if isflip: # horizontal flip
img = cv2.flip(img, 1)
if isRotate:
img_waf = img
img = cv2.getRotationMatrix2D((width/2, height/2), angle, 1)
img = cv2.warpAffine(img_waf, img,(width, height))
img_lr_2 = cv2.resize(img, (64, 64), interpolation=cv2.INTER_LINEAR)
img_lr_4 = cv2.resize(img_lr_2, (32, 32), interpolation=cv2.INTER_LINEAR)
img_lr = cv2.resize(img_lr_4, (16, 16), interpolation=cv2.INTER_LINEAR)
img_grad = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img_grad = cv2.Canny(img_grad, 50, 255)
img_grad_2 = cv2.cvtColor(img_lr_2, cv2.COLOR_BGR2GRAY)
img_grad_2 = cv2.Canny(img_grad_2, 50, 255)
img_grad_4 = cv2.cvtColor(img_lr_4, cv2.COLOR_BGR2GRAY)
img_grad_4 = cv2.Canny(img_grad_4, 50, 255)
img = img.astype(np.float32)
img_lr = img_lr.astype(np.float32)
img_grad = img_grad.astype(np.float32)
img_grad = np.reshape(img_grad, (128, 128, 1))
img_grad_2 = img_grad_2.astype(np.float32)
img_grad_2 = np.reshape(img_grad_2, (64, 64, 1))
img_grad_4 = img_grad_4.astype(np.float32)
img_grad_4 = np.reshape(img_grad_4, (32, 32, 1))
return img, img_lr, img_grad, img_grad_2, img_grad_4
def next_batch(self, batch_size, width, height):
imgs = []
imgs_lr = []
imgs_grad = []
imgs_grad_2 = []
imgs_grad_4 = []
for i in range(0, batch_size):
train_img = random.choice(self.train_list)
flip = random.choice([True, False])
rotation = random.choice([True, False])
angle = random.choice([90, 180, 270])
img, img_lr, img_grad, img_grad_2, img_grad_4 = self.open_image(train_img, width=width,
height=height, isflip=flip,
isRotate=rotation, angle=angle)
imgs.append(img)
imgs_lr.append(img_lr)
imgs_grad.append(img_grad)
imgs_grad_2.append(img_grad_2)
imgs_grad_4.append(img_grad_4)
imgs = np.array(imgs)
imgs_lr = np.array(imgs_lr)
imgs_grad = np.array(imgs_grad)
imgs_grad_2 = np.array(imgs_grad_2)
imgs_grad_4 = np.array(imgs_grad_4)
return imgs, imgs_lr, imgs_grad, imgs_grad_2, imgs_grad_4
def next_batch_test(self, batch_size, width, height):
imgs = []
imgs_lr = []
imgs_grad = []
imgs_grad_2 = []
imgs_grad_4 = []
for i in range(0, batch_size):
test_img = random.choice(self.test_list)
flip = False
rotation = False
angle = 0
img, img_lr, img_grad, img_grad_2, img_grad_4 = self.open_image(test_img, width=width,
height=height, isflip=flip,
isRotate=rotation, angle=angle)
imgs.append(img)
imgs_lr.append(img_lr)
imgs_grad.append(img_grad)
imgs_grad_2.append(img_grad_2)
imgs_grad_4.append(img_grad_4)
imgs = np.array(imgs)
imgs_lr = np.array(imgs_lr)
imgs_grad = np.array(imgs_grad)
imgs_grad_2 = np.array(imgs_grad_2)
imgs_grad_4 = np.array(imgs_grad_4)
return imgs, imgs_lr, imgs_grad, imgs_grad_2, imgs_grad_4