2
2
import os
3
3
import warnings
4
4
from concurrent .futures import ThreadPoolExecutor , as_completed
5
+ from typing import Tuple , List , Sequence , Optional
5
6
6
7
import cv2
7
8
import numpy as np
8
- import pandas as pd
9
9
from scipy .signal import find_peaks
10
10
from skimage .feature import peak_local_max
11
11
12
- from .NMS import NMS
12
+ from .NMS import NMS , Hit
13
13
from .version import __version__
14
14
15
15
__all__ = ['NMS' ]
16
16
17
+ # Define custom "types" for type hints
18
+ BBox = Tuple [int , int , int , int ] # bounding box in the form (x,y,width,height) with x,y top left corner
19
+
17
20
def _findLocalMax_ (corrMap , score_threshold = 0.6 ):
18
21
"""Get coordinates of the local maximas with values above a threshold in the image of the correlation map."""
19
- # IF depending on the shape of the correlation map
22
+ # If depending on the shape of the correlation map
20
23
if corrMap .shape == (1 ,1 ): ## Template size = Image size -> Correlation map is a single digit')
21
24
22
- if corrMap [0 ,0 ]>= score_threshold :
25
+ if corrMap [0 ,0 ] >= score_threshold :
23
26
peaks = np .array ([[0 ,0 ]])
24
27
else :
25
28
peaks = []
@@ -48,7 +51,7 @@ def _findLocalMin_(corrMap, score_threshold=0.4):
48
51
return _findLocalMax_ (- corrMap , - score_threshold )
49
52
50
53
51
- def computeScoreMap (template , image , method = cv2 .TM_CCOEFF_NORMED , mask = None ):
54
+ def computeScoreMap (template , image , method : int = cv2 .TM_CCOEFF_NORMED , mask = None ):
52
55
"""
53
56
Compute score map provided numpy array for template and image (automatically converts images if necessary).
54
57
The template must be smaller or as large as the image.
@@ -87,10 +90,9 @@ def computeScoreMap(template, image, method=cv2.TM_CCOEFF_NORMED, mask=None):
87
90
return cv2 .matchTemplate (image , template , method , mask = mask )
88
91
89
92
90
- def findMatches (listTemplates , image , method = cv2 .TM_CCOEFF_NORMED , N_object = float ("inf" ), score_threshold = 0.5 , searchBox = None ):
93
+ def findMatches (listTemplates , image , method : int = cv2 .TM_CCOEFF_NORMED , N_object = float ("inf" ), score_threshold : float = 0.5 , searchBox : Optional [ BBox ] = None ) -> List [ Hit ] :
91
94
"""
92
95
Find all possible templates locations satisfying the score threshold provided a list of templates to search and an image.
93
- Returns a pandas dataframe with one row per detection.
94
96
95
97
Parameters
96
98
----------
@@ -117,7 +119,10 @@ def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=floa
117
119
118
120
Returns
119
121
-------
120
- - Pandas DataFrame with 1 row per hit and column "TemplateName"(string), "BBox":(X, Y, Width, Height), "Score":float
122
+ A list of hit where each hit is a tuple as following ["TemplateName", (x, y, width, height), score]
123
+ where template name is the name (or label) of the matching template
124
+ (x, y, width, height) is a tuple of the bounding box coordinates in pixels, with xy the coordinates for the top left corner
125
+ score (float) for the confidence of the detection
121
126
"""
122
127
if N_object != float ("inf" ) and not isinstance (N_object , int ):
123
128
raise TypeError ("N_object must be an integer" )
@@ -132,9 +137,9 @@ def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=floa
132
137
## Crop image to search region if provided
133
138
if searchBox is not None :
134
139
xOffset , yOffset , searchWidth , searchHeight = searchBox
135
- image = image [yOffset : yOffset + searchHeight , xOffset : xOffset + searchWidth ]
140
+ image = image [yOffset : yOffset + searchHeight , xOffset : xOffset + searchWidth ]
136
141
else :
137
- xOffset = yOffset = 0
142
+ xOffset = yOffset = 0
138
143
139
144
# Check that the template are all smaller are equal to the image (original, or cropped if there is a search region)
140
145
for index , tempTuple in enumerate (listTemplates ):
@@ -160,22 +165,19 @@ def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=floa
160
165
raise ValueError ("Template '{}' at index {} in the list of templates is larger than {}." .format (tempName , index , fitIn ) )
161
166
162
167
listHit = []
163
- ## Use multi-threading to iterate through all templates, using half the number of cpu cores available.
168
+ # Use multi-threading to iterate through all templates, using half the number of cpu cores available.
169
+ # i.e parallelize the search with the individual templates, in the same image
164
170
with ThreadPoolExecutor (max_workers = round (os .cpu_count ()* .5 )) as executor :
165
171
futures = [executor .submit (_multi_compute , tempTuple , image , method , N_object , score_threshold , xOffset , yOffset , listHit ) for tempTuple in listTemplates ]
166
172
for future in as_completed (futures ):
167
173
_ = future .result ()
168
174
169
- if listHit :
170
- return pd .DataFrame (listHit ) # All possible hits before Non-Maxima Supression
171
- else :
172
- return pd .DataFrame (columns = ["TemplateName" , "BBox" , "Score" ])
173
-
175
+ return listHit # All possible hits before Non-Maxima Supression
174
176
175
- def _multi_compute (tempTuple , image , method , N_object , score_threshold , xOffset , yOffset , listHit ):
177
+ def _multi_compute (tempTuple , image , method : int , N_object : int , score_threshold : float , xOffset : int , yOffset : int , listHit : Sequence [ Hit ] ):
176
178
"""
177
179
Find all possible template locations satisfying the score threshold provided a template to search and an image.
178
- Add the hits in the list of hits .
180
+ Add the hits found to the provided listHit, this function is running in parallel each instance for a different templates .
179
181
180
182
Parameters
181
183
----------
@@ -203,8 +205,7 @@ def _multi_compute(tempTuple, image, method, N_object, score_threshold, xOffset,
203
205
- yOffset : int
204
206
optional the y offset if the search area is provided
205
207
206
- - listHit : the list of hits which we want to add the discovered hit
207
- expected array of hits
208
+ - listHit : New hits are added to this list
208
209
"""
209
210
templateName , template = tempTuple [:2 ]
210
211
mask = None
@@ -232,19 +233,16 @@ def _multi_compute(tempTuple, image, method, N_object, score_threshold, xOffset,
232
233
peaks = _findLocalMax_ (corrMap , score_threshold )
233
234
234
235
#print('Initially found',len(peaks),'hit with this template')
235
-
236
- # Once every peak was detected for this given template
237
- ## Create a dictionnary for each hit with {'TemplateName':, 'BBox': (x,y,Width, Height), 'Score':coeff}
238
-
239
236
height , width = template .shape [0 :2 ] # slicing make sure it works for RGB too
240
237
241
- for peak in peaks :
242
- # append to list of potential hit before Non maxima suppression
243
- # no need to lock the list, append is thread-safe
244
- listHit .append ({'TemplateName' :templateName , 'BBox' : ( int (peak [1 ])+ xOffset , int (peak [0 ])+ yOffset , width , height ) , 'Score' :corrMap [tuple (peak )]}) # empty df with correct column header
238
+ # For each peak create a hit as a list [templateName, (x,y,width,height), score] and add this hit into a bigger list
239
+ newHits = [ (templateName , (int (peak [1 ]) + xOffset , int (peak [0 ]) + yOffset , width , height ), corrMap [tuple (peak )] ) for peak in peaks ]
240
+
241
+ # Finally add these new hits to the original list of hits
242
+ listHit .extend (newHits )
245
243
246
244
247
- def matchTemplates (listTemplates , image , method = cv2 .TM_CCOEFF_NORMED , N_object = float ("inf" ), score_threshold = 0.5 , maxOverlap = 0.25 , searchBox = None ):
245
+ def matchTemplates (listTemplates , image , method : int = cv2 .TM_CCOEFF_NORMED , N_object = float ("inf" ), score_threshold : float = 0.5 , maxOverlap : float = 0.25 , searchBox : Optional [ BBox ] = None ) -> List [ Hit ] :
248
246
"""
249
247
Search each template in the image, and return the best N_object locations which offer the best score and which do not overlap above the maxOverlap threshold.
250
248
@@ -278,23 +276,25 @@ def matchTemplates(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=f
278
276
279
277
Returns
280
278
-------
281
- Pandas DataFrame with 1 row per hit and column "TemplateName"(string), "BBox":(X, Y, Width, Height ), "Score":float
279
+ A list of hit, where each hit is a tuple in the form (label, (x, y, width, height ), score)
282
280
if N=1, return the best matches independently of the score_threshold
283
281
if N<inf, returns up to N best matches that passed the NMS
284
282
if N=inf, returns all matches that passed the NMS
285
283
"""
286
- if maxOverlap < 0 or maxOverlap > 1 :
284
+ if maxOverlap < 0 or maxOverlap > 1 :
287
285
raise ValueError ("Maximal overlap between bounding box is in range [0-1]" )
288
286
289
- tableHit = findMatches (listTemplates , image , method , N_object , score_threshold , searchBox )
287
+ listHits = findMatches (listTemplates , image , method , N_object , score_threshold , searchBox )
290
288
291
- if method == 0 : raise ValueError ("The method TM_SQDIFF is not supported. Use TM_SQDIFF_NORMED instead." )
289
+ if method == 0 :
290
+ raise ValueError ("The method TM_SQDIFF is not supported. Use TM_SQDIFF_NORMED instead." )
291
+
292
292
sortAscending = (method == 1 )
293
293
294
- return NMS (tableHit , score_threshold , sortAscending , N_object , maxOverlap )
294
+ return NMS (listHits , score_threshold , sortAscending , N_object , maxOverlap )
295
295
296
296
297
- def drawBoxesOnRGB (image , tableHit , boxThickness = 2 , boxColor = (255 , 255 , 00 ), showLabel = False , labelColor = (255 , 255 , 0 ), labelScale = 0.5 ):
297
+ def drawBoxesOnRGB (image , listHit : Sequence [ Hit ] , boxThickness : int = 2 , boxColor : Tuple [ int , int , int ] = (255 , 255 , 00 ), showLabel : bool = False , labelColor = (255 , 255 , 0 ), labelScale = 0.5 ):
298
298
"""
299
299
Return a copy of the image with predicted template locations as bounding boxes overlaid on the image
300
300
The name of the template can also be displayed on top of the bounding box with showLabel=True
@@ -303,7 +303,7 @@ def drawBoxesOnRGB(image, tableHit, boxThickness=2, boxColor=(255, 255, 00), sho
303
303
----------
304
304
- image : image in which the search was performed
305
305
306
- - tableHit : list of hit as returned by matchTemplates or findMatches
306
+ - listHit : list of hit as returned by matchTemplates or findMatches
307
307
308
308
- boxThickness: int
309
309
thickness of bounding box contour in pixels
@@ -322,13 +322,21 @@ def drawBoxesOnRGB(image, tableHit, boxThickness=2, boxColor=(255, 255, 00), sho
322
322
original image with predicted template locations depicted as bounding boxes
323
323
"""
324
324
# Convert Grayscale to RGB to be able to see the color bboxes
325
- if image .ndim == 2 : outImage = cv2 .cvtColor (image , cv2 .COLOR_GRAY2RGB ) # convert to RGB to be able to show detections as color box on grayscale image
326
- else : outImage = image .copy ()
325
+ outImage = cv2 .cvtColor (image , cv2 .COLOR_GRAY2RGB ) if image .ndim == 2 else image .copy
327
326
328
- for _ , row in tableHit .iterrows ():
329
- x ,y ,w ,h = row ['BBox' ]
330
- cv2 .rectangle (outImage , (x , y ), (x + w , y + h ), color = boxColor , thickness = boxThickness )
331
- if showLabel : cv2 .putText (outImage , text = row ['TemplateName' ], org = (x , y ), fontFace = cv2 .FONT_HERSHEY_SIMPLEX , fontScale = labelScale , color = labelColor , lineType = cv2 .LINE_AA )
327
+ for label , bbox , _ in listHit :
328
+
329
+ x ,y ,w ,h = bbox
330
+ cv2 .rectangle (outImage , (x , y ), (x + w , y + h ), color = boxColor , thickness = boxThickness )
331
+
332
+ if showLabel :
333
+ cv2 .putText (outImage ,
334
+ text = label ,
335
+ org = (x , y ),
336
+ fontFace = cv2 .FONT_HERSHEY_SIMPLEX ,
337
+ fontScale = labelScale ,
338
+ color = labelColor ,
339
+ lineType = cv2 .LINE_AA )
332
340
333
341
return outImage
334
342
0 commit comments