-
Notifications
You must be signed in to change notification settings - Fork 0
/
buildDatabase.py
145 lines (114 loc) · 5.17 KB
/
buildDatabase.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
import time # For time recording
start_time = time.time() # Start time recording
import pywt # Get access to wavl
import cv2 # OpenCV to open and load images
import numpy as np # For effectively handle arrays
from os import listdir # List local dir
from os.path import isfile, join # Check if file and join paths cross-platform
from sklearn import decomposition # Access PCA and KD-tree
from multiprocessing import Pool # Allowing parallelization
import sys # Access arguments from command line
import _pickle as cPickle # Save and load KD-tree
# Read arguments from user
args = [arg for arg in sys.argv]
# Read images
images_path = [f for f in listdir(args[1]) if isfile(join(args[1], f))]
# Set size for resizing of images
HEIGHT=128
WIDTH=128
# Set number of wavelet transforms
MAX_LEVEL=3
# Main class for image processing
class DWTRootTree():
def __init__(self, img, maxLevel):
self.C1_cA, self.C1_cV, self.C1_cD, self.C1_cH = None, None, None, None
self.C2_cA, self.C2_cV, self.C2_cD, self.C2_cH = None, None, None, None
self.C3_cA, self.C3_cV, self.C3_cD, self.C3_cH = None, None, None, None
self.image = self.resizeImage(img)
self.level = 0
self.maxLevel = maxLevel
self.pca = decomposition.PCA()
# Add color components keys
self.C1 = np.empty(shape=(HEIGHT, WIDTH))
self.C2 = np.empty(shape=(HEIGHT, WIDTH))
self.C3 = np.empty(shape=(HEIGHT, WIDTH))
def run(self):
# Split image's color components
self.getRGBvalues()
# PCA
if ('--pca' in args):
self.components_C1 = self.pca.fit_transform(self.C1)[0:10]
self.components_C2 = self.pca.fit_transform(self.C2)[0:10]
self.components_C3 = self.pca.fit_transform(self.C3)[0:10]
self.C1_cA, self.C1_cV, self.C1_cD, self.C1_cH = self.dwt2(self.components_C1)
self.C2_cA, self.C2_cV, self.C2_cD, self.C2_cH = self.dwt2(self.components_C2)
self.C3_cA, self.C3_cV, self.C3_cD, self.C3_cH = self.dwt2(self.components_C3)
else:
self.C1_cA, self.C1_cV, self.C1_cD, self.C1_cH = self.dwt2(self.C1)
self.C2_cA, self.C2_cV, self.C2_cD, self.C2_cH = self.dwt2(self.C2)
self.C3_cA, self.C3_cV, self.C3_cD, self.C3_cH = self.dwt2(self.C3)
# Display image
def showImage(self, image, window='window'):
cv2.imshow('window', image)
cv2.waitKey()
# DWT Algorithm
def dwt2(self, img, wavelet='db1'):
coeffs = pywt.dwt2(img, wavelet)
cA, (cH, cV, cD) = coeffs
# Return cA for next wavelet transform
return DWTSubTree(cA, level=self.level, maxLevel=self.maxLevel), cH, cV, cD
# Resize
def resizeImage(self, image, inter=cv2.INTER_LINEAR):
dim = (HEIGHT, WIDTH)
resized = cv2.resize(image, dim, interpolation=inter)
return resized
def getRGBvalues(self):
self.C1, self.C2, self.C3 = cv2.split(self.image)
class DWTSubTree():
def __init__(self, img, level, maxLevel):
self.cA, self.cV, self.cD, self.cH = None, None, None, None
self.level = level + 1
self.maxLevel = maxLevel
self.image = img
self.run()
def run(self):
self.cA, self.cH, self.cV, self.cD = self.dwt2(self.image)
# DWT Algorithm
def dwt2(self, img, wavelet='db1',):
coeffs = pywt.dwt2(img, wavelet)
cA, (cH, cV, cD) = coeffs
if (self.level == self.maxLevel):
return cA, cH, cV, cD
return DWTSubTree(cA, level=self.level, maxLevel=self.maxLevel), cV, cD, cH
def getLeafNodes(tree):
if (tree.maxLevel == tree.level):
return tree.cA, tree.cH, tree.cV, tree.cD
return getLeafNodes(tree.cA)
def main(path):
dataset_img = np.uint8(cv2.imread(args[1]+'/%s'%path))
dataset_img = cv2.cvtColor(dataset_img, cv2.COLOR_BGR2HSV)
DatasetImageClass = DWTRootTree(dataset_img, maxLevel=MAX_LEVEL)
DatasetImageClass.run()
feature_vector = []
# Structure feature vectors for KD-tree
[feature_vector.append(np.ravel(value)) for index, value in enumerate([ getLeafNodes(DatasetImageClass.C1_cA),\
getLeafNodes(DatasetImageClass.C2_cA),\
getLeafNodes(DatasetImageClass.C3_cA)])]
return np.ravel(feature_vector)
if __name__ == "__main__":
if (len(args) < 2):
print("Syntax: filename.py ./database --pca(flag)")
sys.exit(0)
# Open pool
pool=Pool()
m = []
result = []
for path in images_path:
r = pool.apply_async(main, (path,))
result.append(r)
for r in result:
m.append(r.get())
from sklearn.neighbors import KDTree
tree = KDTree(m)
print("--- Time to build KD-tree: %.2f seconds ---"% (time.time() - start_time))
cPickle.dump(tree, open('tree_pca.p' if'--pca' in args else 'tree.p','wb'))