Skip to content

Commit a7f24b3

Browse files
LauLauThomLauLauThom
committed
Update function docstrings + fix drawBoxes with RGB images
1 parent 72c2fe4 commit a7f24b3

File tree

1 file changed

+197
-145
lines changed

1 file changed

+197
-145
lines changed

‎MTM/__init__.py

Lines changed: 197 additions & 145 deletions
Original file line numberDiff line numberDiff line change
@@ -1,177 +1,229 @@
11
import cv2
22
import numpy as np
33
from skimage.feature import peak_local_max
4-
from scipy.signal import find_peaks
4+
from scipy.signal import find_peaks
55

66
from MTM.NMS import NMS
77

88
__all__ = ['NMS']
99

1010

1111
def _findLocalMax_(corrMap, score_threshold=0.6):
12-
'''
13-
Get coordinates of the local maximas with values above a threshold in the image of the correlation map
14-
'''
15-
16-
# IF depending on the shape of the correlation map
17-
if corrMap.shape == (1,1): ## Template size = Image size -> Correlation map is a single digit')
18-
19-
if corrMap[0,0]>=score_threshold:
20-
Peaks = np.array([[0,0]])
21-
else:
22-
Peaks = []
12+
'''
13+
Get coordinates of the local maximas with values above a threshold in the image of the correlation map
14+
'''
15+
16+
# IF depending on the shape of the correlation map
17+
if corrMap.shape == (1,1): ## Template size = Image size -> Correlation map is a single digit')
18+
19+
if corrMap[0,0]>=score_threshold:
20+
Peaks = np.array([[0,0]])
21+
else:
22+
Peaks = []
2323

24-
# use scipy findpeaks for the 1D cases (would allow to specify the relative threshold for the score directly here rather than in the NMS
25-
elif corrMap.shape[0] == 1: ## Template is as high as the image, the correlation map is a 1D-array
26-
Peaks = find_peaks(corrMap[0], height=score_threshold) # corrMap[0] to have a proper 1D-array
27-
Peaks = [[0,i] for i in Peaks[0]] # 0,i since one coordinate is fixed (the one for which Template = Image)
28-
24+
# use scipy findpeaks for the 1D cases (would allow to specify the relative threshold for the score directly here rather than in the NMS
25+
elif corrMap.shape[0] == 1: ## Template is as high as the image, the correlation map is a 1D-array
26+
Peaks = find_peaks(corrMap[0], height=score_threshold) # corrMap[0] to have a proper 1D-array
27+
Peaks = [[0,i] for i in Peaks[0]] # 0,i since one coordinate is fixed (the one for which Template = Image)
28+
2929

30-
elif corrMap.shape[1] == 1: ## Template is as wide as the image, the correlation map is a 1D-array
31-
#Peaks = argrelmax(corrMap, mode="wrap")
32-
Peaks = find_peaks(corrMap[:,0], height=score_threshold)
33-
Peaks = [[i,0] for i in Peaks[0]]
30+
elif corrMap.shape[1] == 1: ## Template is as wide as the image, the correlation map is a 1D-array
31+
#Peaks = argrelmax(corrMap, mode="wrap")
32+
Peaks = find_peaks(corrMap[:,0], height=score_threshold)
33+
Peaks = [[i,0] for i in Peaks[0]]
3434

3535

36-
else: # Correlatin map is 2D
37-
Peaks = peak_local_max(corrMap, threshold_abs=score_threshold, exclude_border=False).tolist()
36+
else: # Correlatin map is 2D
37+
Peaks = peak_local_max(corrMap, threshold_abs=score_threshold, exclude_border=False).tolist()
3838

39-
return Peaks
39+
return Peaks
4040

4141

4242

4343
def _findLocalMin_(corrMap, score_threshold=0.4):
44-
'''Find coordinates of local minimas with values below a threshold in the image of the correlation map'''
45-
return _findLocalMax_(-corrMap, -score_threshold)
44+
'''Find coordinates of local minimas with values below a threshold in the image of the correlation map'''
45+
return _findLocalMax_(-corrMap, -score_threshold)
4646

4747

4848

4949
def findMatches(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5, searchBox=None):
50-
'''
51-
Find all possible templates locations provided a list of template to search and an image
52-
- listTemplate : list of tuples [(templateName, templateImage), (templateName2, templateImage2) ]
53-
- method : one of OpenCV template matching method (0 to 5)
54-
- 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
56-
- searchBox : optional search region as a tuple (X, Y, Width, Height)
57-
'''
58-
if N_object!=float("inf") and type(N_object)!=int:
59-
raise TypeError("N_object must be an integer")
60-
61-
elif N_object<1:
62-
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
70-
71-
## 16-bit image are converted to 32-bit for matchTemplate
72-
if image.dtype == 'uint16': image = np.float32(image)
73-
74-
listHit = []
75-
for templateName, template in listTemplates:
76-
77-
#print('\nSearch with template : ',templateName)
78-
## 16-bit image are converted to 32-bit for matchTemplate
79-
if template.dtype == 'uint16': template = np.float32(template)
80-
81-
## Compute correlation map
82-
corrMap = cv2.matchTemplate(template, image, method)
83-
84-
## Find possible location of the object
85-
if N_object==1: # Detect global Min/Max
86-
minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(corrMap)
87-
88-
if method==1:
89-
Peaks = [minLoc[::-1]] # opposite sorting than in the multiple detection
90-
91-
elif method in (3,5):
92-
Peaks = [maxLoc[::-1]]
93-
94-
95-
else:# Detect local max or min
96-
if method==1: # Difference => look for local minima
97-
Peaks = _findLocalMin_(corrMap, score_threshold)
98-
99-
elif method in (3,5):
100-
Peaks = _findLocalMax_(corrMap, score_threshold)
101-
102-
103-
#print('Initially found',len(Peaks),'hit with this template')
104-
105-
106-
# Once every peak was detected for this given template
107-
## Create a dictionnary for each hit with {'TemplateName':, 'BBox': (x,y,Width, Height), 'Score':coeff}
108-
109-
height, width = template.shape[0:2] # slicing make sure it works for RGB too
110-
111-
for peak in Peaks :
112-
coeff = corrMap[tuple(peak)]
113-
newHit = {'TemplateName':templateName, 'BBox': [int(peak[1])+xOffset, int(peak[0])+yOffset, width, height], 'Score':coeff}
114-
115-
# append to list of potential hit before Non maxima suppression
116-
listHit.append(newHit)
117-
118-
119-
return listHit # All possible hit before Non-Maxima Supression
120-
50+
'''
51+
Find all possible templates locations provided a list of template to search and an image
52+
Parameters
53+
----------
54+
- listTemplates : list of tuples (LabelString, Grayscale or RGB numpy array)
55+
templates to search in each image, associated to a label
56+
- image : Grayscale or RGB numpy array
57+
image in which to perform the search, it should be the same bitDepth and number of channels than the templates
58+
- method : int
59+
one of OpenCV template matching method (0 to 5), default 5=0-mean cross-correlation
60+
- N_object: int
61+
expected number of objects in the image
62+
- score_threshold: float in range [0,1]
63+
if N>1, returns local minima/maxima respectively below/above the score_threshold
64+
- searchBox : tuple (X, Y, Width, Height) in pixel unit
65+
optional rectangular search region as a tuple
66+
67+
Returns
68+
-------
69+
- listHit: list of match as dictionaries {"TemplateName":string, "BBox":(X, Y, Width, Height), "Score":float}
70+
'''
71+
if N_object!=float("inf") and type(N_object)!=int:
72+
raise TypeError("N_object must be an integer")
73+
74+
elif N_object<1:
75+
raise ValueError("At least one object should be expected in the image")
76+
77+
## Crop image to search region if provided
78+
if searchBox != None:
79+
xOffset, yOffset, searchWidth, searchHeight = searchBox
80+
image = image[yOffset:yOffset+searchHeight, xOffset:xOffset+searchWidth]
81+
else:
82+
xOffset=yOffset=0
83+
84+
## 16-bit image are converted to 32-bit for matchTemplate
85+
if image.dtype == 'uint16': image = np.float32(image)
86+
87+
listHit = []
88+
for templateName, template in listTemplates:
89+
90+
#print('\nSearch with template : ',templateName)
91+
## 16-bit image are converted to 32-bit for matchTemplate
92+
if template.dtype == 'uint16': template = np.float32(template)
93+
94+
## Compute correlation map
95+
corrMap = cv2.matchTemplate(template, image, method)
96+
97+
## Find possible location of the object
98+
if N_object==1: # Detect global Min/Max
99+
minVal, maxVal, minLoc, maxLoc = cv2.minMaxLoc(corrMap)
100+
101+
if method==1:
102+
Peaks = [minLoc[::-1]] # opposite sorting than in the multiple detection
103+
104+
elif method in (3,5):
105+
Peaks = [maxLoc[::-1]]
106+
107+
108+
else:# Detect local max or min
109+
if method==1: # Difference => look for local minima
110+
Peaks = _findLocalMin_(corrMap, score_threshold)
111+
112+
elif method in (3,5):
113+
Peaks = _findLocalMax_(corrMap, score_threshold)
114+
115+
116+
#print('Initially found',len(Peaks),'hit with this template')
117+
118+
119+
# Once every peak was detected for this given template
120+
## Create a dictionnary for each hit with {'TemplateName':, 'BBox': (x,y,Width, Height), 'Score':coeff}
121+
122+
height, width = template.shape[0:2] # slicing make sure it works for RGB too
123+
124+
for peak in Peaks :
125+
coeff = corrMap[tuple(peak)]
126+
newHit = {'TemplateName':templateName, 'BBox': [int(peak[1])+xOffset, int(peak[0])+yOffset, width, height], 'Score':coeff}
127+
128+
# append to list of potential hit before Non maxima suppression
129+
listHit.append(newHit)
130+
131+
return listHit # All possible hit before Non-Maxima Supression
132+
121133

122134
def matchTemplates(listTemplates, image, method=cv2.TM_CCOEFF_NORMED, N_object=float("inf"), score_threshold=0.5, maxOverlap=0.25, searchBox=None):
123-
'''
124-
Search each template in the image, and return the best N_object location which offer the best score and which do not overlap
125-
- listTemplate : list of tuples (templateName, templateImage)
126-
- method : one of OpenCV template matching method (0 to 5)
127-
- N_object: expected number of object in the image
128-
- score_threshold: if N>1, returns local minima/maxima respectively below/above the score_threshold
129-
'''
130-
if maxOverlap<0 or maxOverlap>1:
131-
raise ValueError("Maximal overlap between bounding box is in range [0-1]")
132-
133-
listHit = findMatches(listTemplates, image, method, N_object, score_threshold, searchBox)
134-
135-
if method == 1: bestHits = NMS(listHit, N_object=N_object, maxOverlap=maxOverlap, sortDescending=False)
136-
137-
elif method in (3,5): bestHits = NMS(listHit, N_object=N_object, maxOverlap=maxOverlap, sortDescending=True)
138-
139-
return bestHits
135+
'''
136+
Search each template in the image, and return the best N_object location which offer the best score and which do not overlap
137+
Parameters
138+
----------
139+
- listTemplates : list of tuples (LabelString, Grayscale or RGB numpy array)
140+
templates to search in each image, associated to a label
141+
- image : Grayscale or RGB numpy array
142+
image in which to perform the search, it should be the same bitDepth and number of channels than the templates
143+
- method : int
144+
one of OpenCV template matching method (0 to 5), default 5=0-mean cross-correlation
145+
- N_object: int
146+
expected number of objects in the image
147+
- score_threshold: float in range [0,1]
148+
if N>1, returns local minima/maxima respectively below/above the score_threshold
149+
- maxOverlap: float in range [0,1]
150+
This is the maximal value for the ratio of the Intersection Over Union (IoU) area between a pair of bounding boxes.
151+
If the ratio is over the maxOverlap, the lower score bounding box is discarded.
152+
- searchBox : tuple (X, Y, Width, Height) in pixel unit
153+
optional rectangular search region as a tuple
154+
155+
Returns
156+
-------
157+
- bestHits: list of match as dictionaries {"TemplateName":string, "BBox":(X, Y, Width, Height), "Score":float}
158+
if N=1, return the best matches independently of the score_threshold
159+
if N<inf, returns up to N best matches that passed the score_threshold
160+
if N=inf, returns all matches that passed the score_threshold
161+
'''
162+
if maxOverlap<0 or maxOverlap>1:
163+
raise ValueError("Maximal overlap between bounding box is in range [0-1]")
164+
165+
listHit = findMatches(listTemplates, image, method, N_object, score_threshold, searchBox)
166+
167+
if method == 1: bestHits = NMS(listHit, N_object=N_object, maxOverlap=maxOverlap, sortDescending=False)
168+
169+
elif method in (3,5): bestHits = NMS(listHit, N_object=N_object, maxOverlap=maxOverlap, sortDescending=True)
170+
171+
return bestHits
140172

141173

142174
def drawBoxes(image, listHit, boxThickness=2, boxColor=(255, 255, 00), showLabel=True, labelColor=(255, 255, 0) ):
143-
"""
144-
Return a copy of the image with results of template matching drawn as yellow rectangle and name of the template on top
145-
"""
146-
outImage = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) # convert to RGB to be able to show detections as color box on grayscale image
147-
148-
for hit in listHit:
149-
x,y,w,h = hit['BBox']
150-
cv2.rectangle(outImage, (x, y), (x+w, y+h), color=boxColor, thickness=boxThickness)
151-
if showLabel: cv2.putText(outImage, text=hit['TemplateName'], org=(x, y), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=labelColor, lineType=cv2.LINE_AA)
152-
153-
return outImage
175+
'''
176+
Return a copy of the image with predicted template locations as bounding boxes overlaid on the image
177+
The name of the template can also be displayed on top of the bounding box with showLabel=True
178+
Parameters
179+
----------
180+
- image : image in which the search was performed
181+
- listHit: list of hit as returned by matchTemplates or findMatches
182+
- boxThickness: int
183+
thickness of bounding box contour in pixels
184+
- boxColor: (int, int, int)
185+
RGB color for the bounding box
186+
- showLabel: Boolean
187+
Display label of the bounding box (field TemplateName)
188+
- labelColor: (int, int, int)
189+
RGB color for the label
190+
191+
Returns
192+
-------
193+
outImage: RGB image
194+
original image with predicted template locations depicted as bounding boxes
195+
'''
196+
# Convert Grayscale to RGB to be able to see the color bboxes
197+
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
198+
else: outImage = image.copy()
199+
200+
for hit in listHit:
201+
x,y,w,h = hit['BBox']
202+
cv2.rectangle(outImage, (x, y), (x+w, y+h), color=boxColor, thickness=boxThickness)
203+
if showLabel: cv2.putText(outImage, text=hit['TemplateName'], org=(x, y), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=labelColor, lineType=cv2.LINE_AA)
204+
205+
return outImage
154206

155207

156208
if __name__ == '__main__':
157-
158-
from skimage.data import coins
159-
import matplotlib.pyplot as plt
160-
161-
## Get image and template
162-
smallCoin = coins()[37:37+38, 80:80+41]
163-
bigCoin = coins()[14:14+59,302:302+65]
164-
image = coins()
165-
166-
## Perform matching
167-
listHit = matchTemplates([('small', smallCoin), ('big', bigCoin)], image, score_threshold=0.4, method=cv2.TM_CCOEFF_NORMED, maxOverlap=0)
168-
#listHit = matchTemplates([('small', smallCoin), ('big', bigCoin)], image, N_object=1, score_threshold=0.4, method=cv2.TM_CCOEFF_NORMED, maxOverlap=0)
209+
210+
from skimage.data import coins
211+
import matplotlib.pyplot as plt
212+
213+
## Get image and template
214+
smallCoin = coins()[37:37+38, 80:80+41]
215+
bigCoin = coins()[14:14+59,302:302+65]
216+
image = coins()
217+
218+
## Perform matching
219+
listHit = matchTemplates([('small', smallCoin), ('big', bigCoin)], image, score_threshold=0.4, method=cv2.TM_CCOEFF_NORMED, maxOverlap=0)
220+
#listHit = matchTemplates([('small', smallCoin), ('big', bigCoin)], image, N_object=1, score_threshold=0.4, method=cv2.TM_CCOEFF_NORMED, maxOverlap=0)
169221

170-
print("Found {} coins".format(len(listHit)))
171-
172-
for hit in listHit:
173-
print(hit)
174-
175-
## Display matches
176-
Overlay = drawBoxes(image, listHit)
177-
plt.imshow(Overlay)
222+
print("Found {} coins".format(len(listHit)))
223+
224+
for hit in listHit:
225+
print(hit)
226+
227+
## Display matches
228+
Overlay = drawBoxes(image, listHit)
229+
plt.imshow(Overlay)

0 commit comments

Comments
 (0)