Skip to content

Commit 826ad66

Browse files
committed
add typing for numpy arrays
1 parent 9863c9b commit 826ad66

File tree

2 files changed

+27
-17
lines changed

2 files changed

+27
-17
lines changed

Diff for: MTM/__init__.py

+25-15
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import warnings
44
from concurrent.futures import ThreadPoolExecutor, as_completed
55
from typing import Tuple, List, Sequence, Optional
6+
from numpy.typing import NDArray
67

78
import cv2
89
import numpy as np
@@ -16,8 +17,9 @@
1617

1718
# Define custom "types" for type hints
1819
BBox = Tuple[int, int, int, int] # bounding box in the form (x,y,width,height) with x,y top left corner
20+
TemplateTuple = Tuple[str, NDArray, Optional[NDArray]]
1921

20-
def _findLocalMax_(corrMap, score_threshold=0.6):
22+
def _findLocalMax_(corrMap:NDArray, score_threshold=0.6):
2123
"""Get coordinates of the local maximas with values above a threshold in the image of the correlation map."""
2224
# If depending on the shape of the correlation map
2325
if corrMap.shape == (1,1): ## Template size = Image size -> Correlation map is a single digit')
@@ -46,12 +48,12 @@ def _findLocalMax_(corrMap, score_threshold=0.6):
4648

4749

4850

49-
def _findLocalMin_(corrMap, score_threshold=0.4):
51+
def _findLocalMin_(corrMap:NDArray, score_threshold=0.4):
5052
"""Find coordinates of local minimas with values below a threshold in the image of the correlation map."""
5153
return _findLocalMax_(-corrMap, -score_threshold)
5254

5355

54-
def computeScoreMap(template, image, method:int = cv2.TM_CCOEFF_NORMED, mask=None):
56+
def computeScoreMap(template:NDArray, image:NDArray, method:int = cv2.TM_CCOEFF_NORMED, mask=None):
5557
"""
5658
Compute score map provided numpy array for template and image (automatically converts images if necessary).
5759
The template must be smaller or as large as the image.
@@ -90,7 +92,7 @@ def computeScoreMap(template, image, method:int = cv2.TM_CCOEFF_NORMED, mask=Non
9092
return cv2.matchTemplate(image, template, method, mask=mask)
9193

9294

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]:
95+
def findMatches(listTemplates : Sequence[TemplateTuple], image:NDArray, method:int = cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold:float=0.5, searchBox:Optional[BBox] = None) -> List[Hit]:
9496
"""
9597
Find all possible templates locations satisfying the score threshold provided a list of templates to search and an image.
9698
@@ -174,7 +176,7 @@ def findMatches(listTemplates, image, method:int = cv2.TM_CCOEFF_NORMED, N_objec
174176

175177
return listHit # All possible hits before Non-Maxima Supression
176178

177-
def _multi_compute(tempTuple, image, method:int, N_object:int, score_threshold:float, xOffset:int, yOffset:int, listHit:Sequence[Hit]):
179+
def _multi_compute(tempTuple : Sequence[TemplateTuple], image:NDArray, method:int, N_object:int, score_threshold:float, xOffset:int, yOffset:int, listHit:Sequence[Hit]):
178180
"""
179181
Find all possible template locations satisfying the score threshold provided a template to search and an image.
180182
Add the hits found to the provided listHit, this function is running in parallel each instance for a different templates.
@@ -242,7 +244,7 @@ def _multi_compute(tempTuple, image, method:int, N_object:int, score_threshold:f
242244
listHit.extend(newHits)
243245

244246

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]:
247+
def matchTemplates(listTemplates:List[TemplateTuple], image:NDArray, 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]:
246248
"""
247249
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.
248250
@@ -294,7 +296,7 @@ def matchTemplates(listTemplates, image, method:int = cv2.TM_CCOEFF_NORMED, N_ob
294296
return NMS(listHits, score_threshold, sortAscending, N_object, maxOverlap)
295297

296298

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 ):
299+
def drawBoxesOnRGB(image:NDArray, listHit:Sequence[Hit], boxThickness:int=2, boxColor:Tuple[int,int,int] = (255, 255, 00), showLabel:bool=False, labelColor=(255, 255, 0), labelScale=0.5 ):
298300
"""
299301
Return a copy of the image with predicted template locations as bounding boxes overlaid on the image
300302
The name of the template can also be displayed on top of the bounding box with showLabel=True
@@ -341,7 +343,7 @@ def drawBoxesOnRGB(image, listHit:Sequence[Hit], boxThickness:int=2, boxColor:Tu
341343
return outImage
342344

343345

344-
def drawBoxesOnGray(image, tableHit, boxThickness=2, boxColor=255, showLabel=False, labelColor=255, labelScale=0.5):
346+
def drawBoxesOnGray(image:NDArray, listHit:Sequence[Hit], boxThickness=2, boxColor=255, showLabel=False, labelColor=255, labelScale=0.5):
345347
"""
346348
Same as drawBoxesOnRGB but with Graylevel.
347349
If a RGB image is provided, the output image will be a grayscale image
@@ -350,7 +352,7 @@ def drawBoxesOnGray(image, tableHit, boxThickness=2, boxColor=255, showLabel=Fal
350352
----------
351353
- image : image in which the search was performed
352354
353-
- tableHit: list of hit as returned by matchTemplates or findMatches
355+
- listHit: list of hit as returned by matchTemplates or findMatches
354356
355357
- boxThickness: int
356358
thickness of bounding box contour in pixels
@@ -370,12 +372,20 @@ def drawBoxesOnGray(image, tableHit, boxThickness=2, boxColor=255, showLabel=Fal
370372
original image with predicted template locations depicted as bounding boxes
371373
"""
372374
# Convert RGB to grayscale
373-
if image.ndim == 3: outImage = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) # convert to RGB to be able to show detections as color box on grayscale image
374-
else: outImage = image.copy()
375+
outImage = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) if image.ndim == 3 else image.copy()
375376

376-
for _, row in tableHit.iterrows():
377-
x,y,w,h = row['BBox']
377+
for label, bbox, _ in listHit:
378+
379+
x,y,w,h = bbox
378380
cv2.rectangle(outImage, (x, y), (x+w, y+h), color=boxColor, thickness=boxThickness)
379-
if showLabel: cv2.putText(outImage, text=row['TemplateName'], org=(x, y), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=labelScale, color=labelColor, lineType=cv2.LINE_AA)
381+
382+
if showLabel:
383+
cv2.putText(outImage,
384+
text=label,
385+
org=(x, y),
386+
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
387+
fontScale=labelScale,
388+
color=labelColor,
389+
lineType=cv2.LINE_AA)
380390

381-
return outImage
391+
return outImage

Diff for: setup.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
url="https://github.com/multi-template-matching/MultiTemplateMatching-Python",
2020
packages=["MTM"],
2121
install_requires=[
22-
'numpy',
23-
'opencv-python-headless>=4.5.4',
22+
'numpy >= 1.21',
23+
'opencv-python-headless >= 4.5.4',
2424
'scikit-image',
2525
'scipy',
2626
],

0 commit comments

Comments
 (0)