Skip to content

Commit b85f071

Browse files
committed
add search box
1 parent 954e29f commit b85f071

File tree

1 file changed

+15
-6
lines changed

1 file changed

+15
-6
lines changed

MTM/__init__.py

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,22 +46,30 @@ def _findLocalMin_(corrMap, score_threshold=0.4):
4646

4747

4848

49-
def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5):
49+
def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5, searchBox=None):
5050
'''
5151
Find all possible templates locations provided a list of template to search and an image
5252
- listTemplate : list of tuples [(templateName, templateImage), (templateName2, templateImage2) ]
5353
- method : one of OpenCV template matching method (0 to 5)
5454
- N_object: expected number of object in the image
55-
- score_threshold: if N>1, returns local minima/maxima respectively below/above the score_threshold
55+
- score_threshold: if N>1, returns local minima/maxima respectively below/above the score_threshold
56+
- searchBox : optional search region as a tuple (X, Y, Width, Height)
5657
'''
5758
if N_object!=float("inf") and type(N_object)!=int:
5859
raise TypeError("N_object must be an integer")
5960

6061
elif N_object<1:
6162
raise ValueError("At least one object should be expected in the image")
63+
64+
## Crop image to search region if provided
65+
if searchBox != None:
66+
xOffset, yOffset, searchWidth, searchHeight = searchBox
67+
image = image[yOffset:yOffset+searchHeight, xOffset:xOffset+searchWidth]
68+
else:
69+
xOffset=yOffset=0
6270

6371
## 16-bit image are converted to 32-bit for matchTemplate
64-
if image.dtype == 'uint16': image = np.float32(image)
72+
if image.dtype == 'uint16': image = np.float32(image)
6573

6674
listHit = []
6775
for templateName, template in listTemplates:
@@ -99,9 +107,10 @@ def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=floa
99107
## Create a dictionnary for each hit with {'TemplateName':, 'BBox': (x,y,Width, Height), 'Score':coeff}
100108

101109
height, width = template.shape[0:2] # slicing make sure it works for RGB too
110+
102111
for peak in Peaks :
103112
coeff = corrMap[tuple(peak)]
104-
newHit = {'TemplateName':templateName, 'BBox': [int(peak[1]), int(peak[0]), width, height], 'Score':coeff}
113+
newHit = {'TemplateName':templateName, 'BBox': [int(peak[1])+xOffset, int(peak[0])+yOffset, width, height], 'Score':coeff}
105114

106115
# append to list of potential hit before Non maxima suppression
107116
listHit.append(newHit)
@@ -110,7 +119,7 @@ def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=floa
110119
return listHit # All possible hit before Non-Maxima Supression
111120

112121

113-
def matchTemplates(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5, maxOverlap=0.25):
122+
def matchTemplates(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5, maxOverlap=0.25, searchBox=None):
114123
'''
115124
Search each template in the image, and return the best N_object location which offer the best score and which do not overlap
116125
- listTemplate : list of tuples (templateName, templateImage)
@@ -121,7 +130,7 @@ def matchTemplates(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=f
121130
if maxOverlap<0 or maxOverlap>1:
122131
raise ValueError("Maximal overlap between bounding box is in range [0-1]")
123132

124-
listHit = findMatches(listTemplates, image, method, N_object, score_threshold)
133+
listHit = findMatches(listTemplates, image, method, N_object, score_threshold, searchBox)
125134

126135
if method == 1: bestHits = NMS(listHit, N_object=N_object, maxOverlap=maxOverlap, sortDescending=False)
127136

0 commit comments

Comments
 (0)