@@ -4,6 +4,7 @@
from src .benchmark import Benchmark
from copy import deepcopy
import matplotlib .pyplot as plt
class TemplateMatching :
"""
@@ -12,32 +13,67 @@ class TemplateMatching:
methods = ['cv2.TM_CCOEFF' , 'cv2.TM_CCOEFF_NORMED' , 'cv2.TM_CCORR' ,
'cv2.TM_CCORR_NORMED' , 'cv2.TM_SQDIFF' , 'cv2.TM_SQDIFF_NORMED' ]
def __init__ (self , method = "cv2.TM_CCORR_NORMED" , threshold = 0.95 ):
# self.method = eval(method)
# self.threshold = threshold
def __init__ (self ):
self .images = None
self .corners = None
self .w = None
self .h = None
self .transformation = None
self .pts1 = None
def load_corners (self , path ):
if self .corners is None :
self .corners = {}
corner_normal = cv2 .imread (os .path .join (path , "template_corner.png" ), 0 )
tl = cv2 .imread (os .path .join (path , "top_left_4.png" ), 0 )
tr = cv2 .imread (os .path .join (path , "top_right_2.png" ), 0 )
corner = cv2 .imread (os .path .join (path , "template_orange.png" ))
example = [tl , tr , np .flip (tr , axis = 0 ), np .flip (tl , axis = 0 )]
corner_normal = cv2 .cvtColor (corner , cv2 .COLOR_BGR2GRAY )
self .w , self .h = corner_normal .shape
# tl = cv2.imread(os.path.join(path, "top_left_4.png"), 0)
# tr = cv2.imread(os.path.join(path, "top_right_2.png"), 0)
# example = [tl, tr, np.flip(tr, axis=0), np.flip(tl, axis=0)]
self .corners ["nominal" ] = [np .rot90 (corner_normal )] # TODO: small mistake
self .corners ["top_left" ] = [example ]
# self.corners["top_left"] = [example]
for i in range (3 ):
self .corners ["nominal" ].append (np .rot90 (self .corners ["nominal" ][i ]))
self .corners ["top_left" ].append ([np .rot90 (self .corners ["top_left" ][i ][j - 1 ], axes = (1 , 0 )) for j in range (4 )])
# self.corners["top_left"].append([np.rot90(self.corners["top_left"][i][j-1], axes=(1, 0)) for j in range(4)])
# if corner_normal.shape == tl.shape == tr.shape:
# self.w, self.h = corner_normal.shape
# else:
# raise IndexError("Corners of templates have different sizes")
center_x , center_y , r = (8 , 11 , 7 )
pts1 = np .float32 ([[center_x , center_y - r ], [center_x , center_y ], [center_x + r , center_y ]])
transformation = lambda alpha : np .float32 ([[center_x , center_y - r ],
[center_x , center_y ],
[center_x + r * np .cos (np .pi / 180 * alpha ), center_y - r * np .sin (np .pi / 180 * alpha )]])
# generate new rotated templates
for new_angle in np .arange (- 60 , 61 , 5 ):
pts2 = transformation (new_angle )
M = cv2 .getAffineTransform (pts1 , pts2 )
dst = cv2 .warpAffine (corner_normal , M , (16 ,16 ))
# averages = [np.mean(dst[:, i][dst[:, i] > 25], axis=0) for i in range(dst.shape[0])]
if corner_normal .shape == tl .shape == tr .shape :
self .w , self .h = corner_normal .shape
else :
raise IndexError ("Corners of templates have different sizes" )
# fill empty space
for i in range (dst .shape [0 ]):
# col = dst[:, i]
# col[np.where(col < 25)[0]] = averages[i]
# dst[:, i] = col
self .corners [new_angle ] = dst
# self.corners[new_angle] = cv2.dilate(dst[int(7*np.sin(abs(new_angle)*np.pi/180)/2):int(19-7*np.sin(abs(new_angle)*np.pi/180) /2),
# int(2+7*np.sin(abs(new_angle)*np.pi/180)/2):int(19-7*np.sin(abs(new_angle)*np.pi/180) /2)], (3,3), iterations=2)
# self.corners[new_angle] = cv2.cvtColor(dst[int(7*np.sin(abs(new_angle)*np.pi/180)/2):int(19-7*np.sin(abs(new_angle)*np.pi/180) /2),
# int(2+7*np.sin(abs(new_angle)*np.pi/180)/2):int(19-7*np.sin(abs(new_angle)*np.pi/180) /2)],
# cv2.COLOR_BGR2GRAY)
# plt.imshow(self.corners[new_angle])
# plt.show()
# print(self.)
for i , angle in enumerate ([0 ,30 ,- 30 ]):
plt .subplot (131 + i )
plt .imshow (self .corners [angle ], 'gray' )
# plt.show()
def test_single (self , path , method , threshold ):
if method not in self .methods :
@@ -94,33 +130,50 @@ def find_gate(self, img_name, *args):
# TODO: Does 1 extra template matching for curved corners (1st time)
# TODO: Smarter use of TM (if 2 points are far away, or not feasible, stop doing TM)
if not gate_detected :
corner_selection = self .find_corners (img , map (lambda x : x [0 ], self .corners ["top_left" ]), method , threshold )
selected = sorted (corner_selection , key = lambda x : x ["max_val" ])[- 1 ] # selected template by maximum value
# selected_index = [i for i in corner_selection if i == selected]
if selected ["max_val" ] > threshold :
# continue with other 3 corners
corners = self .find_corners (img , self .corners ["top_left" ][selected ["id" ]], method , threshold )
# Change order of corners, so it matches order of benchmark
corners = [corners [2 ], corners [1 ], corners [0 ], corners [3 ]]
# Re-check obtained corners to gate
gate_detected = self .find_gate_detected (corners )
# if not gate_detected:
# # corner_selection = self.find_corners(img, map(lambda x: x[0], self.corners["top_left"]), method, threshold)
# curved_corners = [self.corners[i] for i in np.arange(0, 60, 5)]
# corner_selection = self.find_corners(img, curved_corners, method, threshold, True)
# selected = sorted(corner_selection, key=lambda x: x["max_val"])[-1] # selected template by maximum value
# # selected_index = [i for i in corner_selection if i == selected]
# if selected["max_val"] > threshold:
# # continue with other 3 corners
# # corners = self.find_corners(img, self.corners["top_left"][selected["id"]], method, threshold)
# f_angle = selected["angle"]
# print(f_angle)
# other_3 = [np.flip(self.corners[f_angle], axis=0),
# np.flip(self.corners[-abs(f_angle)], axis=1),
# np.flip(np.flip(self.corners[-abs(f_angle)], axis=1), axis=0)]
# # for i in range(4):
# # if i == 0:
# # plt.subplot(221)
# # plt.imshow(self.corners[f_angle], cmap="gray")
# # else:
# # plt.subplot(221+i)
# # plt.imshow(other_3[i-1], cmap="gray")
# # plt.show()
# corners = self.find_corners(img, other_3, method, threshold)
# # Change order of corners, so it matches order of benchmark
# corners = [corners[1], corners[2], corners[0], selected]
# # Re-check obtained corners to gate
# gate_detected = self.find_gate_detected(corners)
return gate_detected , corners [0 ]["max_loc" ], corners [1 ]["max_loc" ], corners [2 ]["max_loc" ], corners [3 ]["max_loc" ]
def find_corners (self , img , templates , method , threshold ):
def find_corners (self , img , templates , method , threshold , angle = False ):
corners = []
for i , template in enumerate (templates ):
try :
res = cv2 .matchTemplate (img , template , method )
except TypeError :
print ("Invalid img" , img )
return False , ( 0 , 0 ), ( 0 , 0 ), ( 0 , 0 ), ( 0 , 0 )
raise TypeError ("Invalid img " , img )
min_val , max_val , min_loc , max_loc = cv2 .minMaxLoc (res )
top_left = max_loc if max_val > threshold else (0 , 0 )
bottom_right = (top_left [0 ] + self .w , top_left [1 ] + self .h )
corners .append ({"id" : i , "det" : top_left != (0 , 0 ), "res" : res , "max_val" : max_val , "max_loc" : max_loc ,
"top_left" : top_left , "bot_right" : bottom_right })
"top_left" : top_left , "bot_right" : bottom_right , "angle" : 0 if not angle else 0 + i * 5 }) # TODO: change hardcode
# print(i)
return corners
def find_gate_detected (self , corners ):
@@ -129,7 +182,8 @@ def find_gate_detected(self, corners):
if self .corners_ordered (corners ):
if self .corners_dont_cross (corners ):
if self .check_shape (corners ):
return True
if self .check_scale (corners ):
return True
return False
@staticmethod
@@ -154,25 +208,35 @@ def corners_dont_cross(corners):
@staticmethod
def check_shape (corners , threshold = 20 ):
if abs (corners [0 ]["max_loc" ][0 ]- corners [1 ]["max_loc" ][0 ]) < threshold :
if abs (corners [2 ]["max_loc" ][0 ] - corners [3 ]["max_loc" ][0 ]) < threshold :
if abs (corners [0 ]["max_loc" ][1 ] - corners [3 ]["max_loc" ][1 ]) < threshold :
return True
elif abs (corners [1 ]["max_loc" ][1 ] - corners [2 ]["max_loc" ][1 ]) < threshold :
return True
if abs (corners [2 ]["max_loc" ][0 ] - corners [3 ]["max_loc" ][0 ]) < threshold :
if abs (corners [2 ]["max_loc" ][1 ] - corners [1 ]["max_loc" ][1 ]) < threshold :
return True
elif abs (corners [3 ]["max_loc" ][1 ] - corners [0 ]["max_loc" ][1 ]) < threshold :
return True
return False
def animation (self , path , progress = False ):
@staticmethod
def check_scale (corners ):
d = [np .linalg .norm (np .array (corners [i % 4 ]["max_loc" ]) - np .array (corners [(i + 1 ) % 4 ]["max_loc" ])) for i in range (4 )]
f = 0
for i in range (4 ):
if not (0.5 < d [i % 4 ]/ d [(i + 1 ) % 4 ] < 1.5 ):
f += 1
return True if f < 2 else False
def animation (self , path ):
print ("Rendering images" )
images = []
if progress :
import progressbar
import time
bar = progressbar .ProgressBar (max_value = len (tuple (filter (lambda x : 'jpg' in x , os .listdir (path )))))
i = 0
for filename in sorted (os .listdir (path )):
if '.jpg' in filename :
path_to_file = os .path .join (path , filename )
img = cv2 .imread (path_to_file )
gd , c0 , c1 , c2 , c3 = tm .find_gate (path_to_file , "cv2.TM_CCORR_NORMED" , 0.95 )
gd , c0 , c1 , c2 , c3 = self .find_gate (path_to_file , "cv2.TM_CCORR_NORMED" , 0.95 )
cv2 .rectangle (img , (c0 [0 ]+ self .w ,) + (c0 [1 ]+ self .h ,),
(c0 [0 ]- self .w ,) + (c0 [1 ]- self .h ,), (255 , 0 , 0 ), thickness = 3 )
@@ -191,9 +255,6 @@ def animation(self, path, progress=False):
cv2 .putText (img , str (gd ), (50 , 100 ), font , 4 , (255 , 0 , 0 ), 2 , cv2 .LINE_AA )
images .append (img )
if progress :
i += 1
bar .update (i )
print ("Rendering Done" )
i = 0
@@ -210,10 +271,65 @@ def __str__(self):
if __name__ == "__main__" :
path = "../data/cad_renders2_test"
path = "../data/cad_renders3_new"
path = "/home/kit/projects/find-me/data/20June2/GateRenders360Rotation/GateRenders360Rotation_animation"
path_corners = "../data/corners"
tm = TemplateMatching ()
tm .load_corners (path_corners )
times = []
import time
import random
files = random .choices (os .listdir ("../data/validation/5-3" ), k = 10 )
# 0.05964167339843698 -> 16.766799840096485 Hz
for i in range (10 ):
t1 = time .perf_counter ()
print (files [i ])
tm .find_gate (os .path .join ("../data/validation/5-3" , files [i ]), "cv2.TM_CCORR_NORMED" , 0.95 )
t = time .perf_counter () - t1
times .append (t )
print (times )
print (np .mean (times ))
# tm.animation(path)
tm .test_many (path , methods = ["cv2.TM_CCOEFF_NORMED" ])
# # tm.test_many(path, methods=["cv2.TM_CCOEFF_NORMED"])
# tm.animation(path, True)
# cv2.warpPerspective()
# euler = lambda point, alpha: [point[0]*np.cos(alpha), point[1]*np.sin(alpha)]
# import matplotlib.pyplot as plt
# # func = lambda alpha: np.float32([[8-7*np.sin(np.pi/180*alpha), 11-7*np.cos(np.pi/180*alpha)],
# # [8, 11],
# # [8+7*np.cos(np.pi/180*alpha), 11-7*np.sin(np.pi/180*alpha)]])
#
# func = lambda alpha: np.float32([[8-0*np.sin(np.pi/180*alpha), 4-0*np.cos(np.pi/180*alpha)],
# [8, 11],
# [8+7*np.cos(np.pi/180*alpha), 11-7*np.sin(np.pi/180*alpha)]])
# #
# img = cv2.imread(os.path.join(path_corners, "template_corner.png"))
# rows, cols, ch = img.shape
#
# pts1 = np.float32([[8, 4], [8, 11], [15, 11]])
# pts2 = func(45)
#
# M = cv2.getAffineTransform(pts1, pts2)
#
# dst = cv2.warpAffine(img, M, (20, 20))
#
# averages = [np.mean(dst[:, i, :], axis=0) for i in range(dst.shape[0])]
#
# for i in range(dst.shape[0]):
# col = dst[:, i,:]
# col[np.where(col < (25,25,25))[0]] = averages[i]
# dst[:, i, :] = col
# # print(col == (0,0,0))
#
# plt.subplot(133), plt.imshow(dst[:, 1:18, :]) # TODO : figure out
#
# # dst = cv2.morphologyEx(dst, cv2.MORPH_OPEN, (15, 15), iterations=3)
# # dst = cv2.GaussianBlur(dst, (5,5), sigmaX=10, sigmaY=10)
# # dst = cv2.dilate(dst, (5, 5), iterations=2)
# plt.subplot(131), plt.imshow(img), plt.plot([i[0] for i in pts1], [i[1] for i in pts1]), plt.title('Input')
# plt.subplot(132), plt.imshow(dst), plt.plot([i[0] for i in pts2], [i[1] for i in pts2]), plt.title('Output')
# # print(dst)
# plt.show()