282 lines
14 KiB
Python
282 lines
14 KiB
Python
# https://github.com/cozheyuanzhangde/Invariant-TemplateMatching
|
|
import cv2
|
|
import numpy as np
|
|
from matplotlib import pyplot as plt
|
|
import matplotlib.patches as patches
|
|
import matplotlib as mpl
|
|
import logging
|
|
import torch
|
|
from PIL import Image
|
|
|
|
logger = logging.getLogger('comfyui_template_matching')
|
|
|
|
def rotate_image(image, angle):
|
|
image_center = tuple(np.array(image.shape[1::-1]) / 2)
|
|
rot_mat = cv2.getRotationMatrix2D(image_center, -angle, 1.0)
|
|
result = cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR)
|
|
return result
|
|
|
|
def scale_image(image, percent, maxwh):
|
|
max_width = maxwh[1]
|
|
max_height = maxwh[0]
|
|
max_percent_width = max_width / image.shape[1] * 100
|
|
max_percent_height = max_height / image.shape[0] * 100
|
|
max_percent = 0
|
|
if max_percent_width < max_percent_height:
|
|
max_percent = max_percent_width
|
|
else:
|
|
max_percent = max_percent_height
|
|
if percent > max_percent:
|
|
percent = max_percent
|
|
width = int(image.shape[1] * percent / 100)
|
|
height = int(image.shape[0] * percent / 100)
|
|
result = cv2.resize(image, (width, height), interpolation = cv2.INTER_AREA)
|
|
return result, percent
|
|
|
|
def invariant_match_template(rgbimage, rgbtemplate, method, matched_thresh, rot_range, rot_interval, scale_range, scale_interval, rm_redundant, minmax, rgbdiff_thresh=float("inf")):
|
|
"""
|
|
rgbimage: RGB image where the search is running.
|
|
rgbtemplate: RGB searched template. It must be not greater than the source image and have the same data type.
|
|
method: [String] Parameter specifying the comparison method
|
|
matched_thresh: [Float] Setting threshold of matched results(0~1).
|
|
rot_range: [Integer] Array of range of rotation angle in degrees. Example: [0,360]
|
|
rot_interval: [Integer] Interval of traversing the range of rotation angle in degrees.
|
|
scale_range: [Integer] Array of range of scaling in percentage. Example: [50,200]
|
|
scale_interval: [Integer] Interval of traversing the range of scaling in percentage.
|
|
rm_redundant: [Boolean] Option for removing redundant matched results based on the width and height of the template.
|
|
minmax:[Boolean] Option for finding points with minimum/maximum value.
|
|
rgbdiff_thresh: [Float] Setting threshold of average RGB difference between template and source image. Default: +inf threshold (no rgbdiff)
|
|
|
|
Returns: List of satisfied matched points in format [[point.x, point.y], angle, scale].
|
|
"""
|
|
img_gray = cv2.cvtColor(rgbimage, cv2.COLOR_RGB2GRAY)
|
|
template_gray = cv2.cvtColor(rgbtemplate, cv2.COLOR_RGB2GRAY)
|
|
image_maxwh = img_gray.shape
|
|
height, width = template_gray.shape
|
|
all_points = []
|
|
if minmax == False:
|
|
for next_angle in range(rot_range[0], rot_range[1], rot_interval):
|
|
for next_scale in range(scale_range[0], scale_range[1], scale_interval):
|
|
scaled_template_gray, actual_scale = scale_image(template_gray, next_scale, image_maxwh)
|
|
if next_angle == 0:
|
|
rotated_template = scaled_template_gray
|
|
else:
|
|
rotated_template = rotate_image(scaled_template_gray, next_angle)
|
|
if method == "TM_CCOEFF":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCOEFF)
|
|
satisfied_points = np.where(matched_points >= matched_thresh)
|
|
elif method == "TM_CCOEFF_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCOEFF_NORMED)
|
|
satisfied_points = np.where(matched_points >= matched_thresh)
|
|
elif method == "TM_CCORR":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCORR)
|
|
satisfied_points = np.where(matched_points >= matched_thresh)
|
|
elif method == "TM_CCORR_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCORR_NORMED)
|
|
satisfied_points = np.where(matched_points >= matched_thresh)
|
|
elif method == "TM_SQDIFF":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_SQDIFF)
|
|
satisfied_points = np.where(matched_points <= matched_thresh)
|
|
elif method == "TM_SQDIFF_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_SQDIFF_NORMED)
|
|
satisfied_points = np.where(matched_points <= matched_thresh)
|
|
else:
|
|
print("There's no such comparison method for template matching.")
|
|
for pt in zip(*satisfied_points[::-1]):
|
|
all_points.append([pt, next_angle, actual_scale])
|
|
else:
|
|
for next_angle in range(rot_range[0], rot_range[1], rot_interval):
|
|
for next_scale in range(scale_range[0], scale_range[1], scale_interval):
|
|
scaled_template_gray, actual_scale = scale_image(template_gray, next_scale, image_maxwh)
|
|
if next_angle == 0:
|
|
rotated_template = scaled_template_gray
|
|
else:
|
|
rotated_template = rotate_image(scaled_template_gray, next_angle)
|
|
if method == "TM_CCOEFF":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCOEFF)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if max_val >= matched_thresh:
|
|
all_points.append([max_loc, next_angle, actual_scale, max_val])
|
|
elif method == "TM_CCOEFF_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCOEFF_NORMED)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if max_val >= matched_thresh:
|
|
all_points.append([max_loc, next_angle, actual_scale, max_val])
|
|
elif method == "TM_CCORR":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCORR)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if max_val >= matched_thresh:
|
|
all_points.append([max_loc, next_angle, actual_scale, max_val])
|
|
elif method == "TM_CCORR_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_CCORR_NORMED)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if max_val >= matched_thresh:
|
|
all_points.append([max_loc, next_angle, actual_scale, max_val])
|
|
elif method == "TM_SQDIFF":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_SQDIFF)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if min_val <= matched_thresh:
|
|
all_points.append([min_loc, next_angle, actual_scale, min_val])
|
|
elif method == "TM_SQDIFF_NORMED":
|
|
matched_points = cv2.matchTemplate(img_gray,rotated_template,cv2.TM_SQDIFF_NORMED)
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(matched_points)
|
|
if min_val <= matched_thresh:
|
|
all_points.append([min_loc, next_angle, actual_scale, min_val])
|
|
else:
|
|
print("There's no such comparison method for template matching.")
|
|
if method == "TM_CCOEFF":
|
|
all_points = sorted(all_points, key=lambda x: -x[3])
|
|
elif method == "TM_CCOEFF_NORMED":
|
|
all_points = sorted(all_points, key=lambda x: -x[3])
|
|
elif method == "TM_CCORR":
|
|
all_points = sorted(all_points, key=lambda x: -x[3])
|
|
elif method == "TM_CCORR_NORMED":
|
|
all_points = sorted(all_points, key=lambda x: -x[3])
|
|
elif method == "TM_SQDIFF":
|
|
all_points = sorted(all_points, key=lambda x: x[3])
|
|
elif method == "TM_SQDIFF_NORMED":
|
|
all_points = sorted(all_points, key=lambda x: x[3])
|
|
if rm_redundant == True:
|
|
lone_points_list = []
|
|
visited_points_list = []
|
|
for point_info in all_points:
|
|
point = point_info[0]
|
|
scale = point_info[2]
|
|
all_visited_points_not_close = True
|
|
if len(visited_points_list) != 0:
|
|
for visited_point in visited_points_list:
|
|
if ((abs(visited_point[0] - point[0]) < (width * scale / 100)) and (abs(visited_point[1] - point[1]) < (height * scale / 100))):
|
|
all_visited_points_not_close = False
|
|
if all_visited_points_not_close == True:
|
|
lone_points_list.append(point_info)
|
|
visited_points_list.append(point)
|
|
else:
|
|
lone_points_list.append(point_info)
|
|
visited_points_list.append(point)
|
|
points_list = lone_points_list
|
|
else:
|
|
points_list = all_points
|
|
if rgbdiff_thresh != float("inf"):
|
|
print(">>>RGBDiff Filtering>>>")
|
|
color_filtered_list = []
|
|
template_channels = cv2.mean(rgbtemplate)
|
|
template_channels = np.array([template_channels[0], template_channels[1], template_channels[2]])
|
|
for point_info in points_list:
|
|
point = point_info[0]
|
|
cropped_img = rgbimage[point[1]:point[1]+height, point[0]:point[0]+width]
|
|
cropped_channels = cv2.mean(cropped_img)
|
|
cropped_channels = np.array([cropped_channels[0], cropped_channels[1], cropped_channels[2]])
|
|
diff_observation = cropped_channels - template_channels
|
|
total_diff = np.sum(np.absolute(diff_observation))
|
|
print(total_diff)
|
|
if total_diff < rgbdiff_thresh:
|
|
color_filtered_list.append([point_info[0],point_info[1],point_info[2]])
|
|
return color_filtered_list
|
|
else:
|
|
return points_list
|
|
|
|
def parse_points_list(points_list,width,height):
|
|
"""
|
|
points_list: List of satisfied matched points in format [[point.x, point.y], angle, scale].
|
|
|
|
Returns: List of satisfied matched points in format [[point.x, point.y], angle, scale].
|
|
"""
|
|
centers_list = []
|
|
box_list = []
|
|
for point_info in points_list:
|
|
point = point_info[0]
|
|
angle = point_info[1]
|
|
scale = point_info[2]
|
|
if angle >=0 and angle <= 5:
|
|
print(f"matched point: {point}, angle: {angle}, scale: {scale}")
|
|
centers_list.append([point, scale])
|
|
box = patches.Rectangle((point[0], point[1]), width*scale/100, height*scale/100, color="green", alpha=0.50, label='Bounding box')
|
|
box_list.append([point[0], point[1], point[0]+width*scale/100, point[1]+height*scale/100])
|
|
return centers_list, box_list
|
|
|
|
def convert_box_to_mask(box_list, image_shape):
|
|
"""
|
|
box_list: List of bounding box in format [x1, y1, x2, y2].
|
|
image_shape: Shape of the image.
|
|
|
|
Returns: Mask of bounding box.
|
|
"""
|
|
res_masks = []
|
|
empty_mask = torch.zeros((image_shape[0], image_shape[1]), dtype=torch.float32)
|
|
if len(box_list) == 0:
|
|
return (empty_mask,)
|
|
for box in box_list:
|
|
x1, y1,x2,y2 = map(int, box)
|
|
empty_mask[y1:y2, x1:x2] = 255
|
|
res_masks.append(empty_mask)
|
|
|
|
return torch.cat(res_masks, dim=0)
|
|
|
|
class TemplateMatching:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
method_list = ["TM_CCOEFF", "TM_CCOEFF_NORMED", "TM_CCORR", "TM_CCORR_NORMED", "TM_SQDIFF", "TM_SQDIFF_NORMED"]
|
|
return {
|
|
"required": {
|
|
"rgbimage": ('IMAGE', {}),
|
|
"rgbtemplate": ('IMAGE', {}),
|
|
"method": (method_list, {"default": "TM_CCOEFF_NORMED"}),
|
|
"matched_thresh": ('FLOAT', {"default": 0.8}),
|
|
"rot_range": ("STRING", {"default": "0,360"}),
|
|
|
|
"rot_interval": ('INT', {"default": 180}),
|
|
"scale_range": ("STRING", {"default": "50,200"}),
|
|
"scale_interval": ('INT', {"default": 10}),
|
|
"rm_redundant": ('BOOLEAN', {"default": True}),
|
|
"minmax": ('BOOLEAN', {"default": True}),
|
|
|
|
}
|
|
}
|
|
CATEGORY = "template_matching"
|
|
FUNCTION = "main"
|
|
RETURN_TYPES = ("MASK",)
|
|
|
|
|
|
|
|
def main(self, rgbimage, rgbtemplate, method, matched_thresh, rot_range, rot_interval, scale_range, scale_interval, rm_redundant, minmax):
|
|
print(f"before: {rgbimage.shape}")
|
|
rgbimage = np.clip(255. * rgbimage.squeeze(0).cpu().numpy(), 0, 255).astype(np.uint8)
|
|
print(f"after: {rgbimage.shape}")
|
|
rgbtemplate = np.clip(255. * rgbtemplate.squeeze(0).cpu().numpy(), 0, 255).astype(np.uint8)
|
|
|
|
if ',' in rot_range:
|
|
rot_range = [int(idx) for idx in rot_range.split(',')]
|
|
if len(rot_range) != 2:
|
|
logger.info("Invalid rot_range. Using default value.")
|
|
rot_range = [0, 360]
|
|
if ',' in scale_range:
|
|
scale_range = [int(idx) for idx in scale_range.split(',')]
|
|
if len(scale_range) != 2:
|
|
logger.info("Invalid scale_range. Using default value.")
|
|
scale_range = [50, 200]
|
|
points_list = invariant_match_template(rgbimage, rgbtemplate, method, matched_thresh, rot_range, rot_interval, scale_range, scale_interval, rm_redundant, minmax)
|
|
mask = torch.zeros((rgbimage.shape[0], rgbimage.shape[1]), dtype=torch.float32)
|
|
if len(points_list) == 0:
|
|
return (mask,)
|
|
centers_list, box_list = parse_points_list(points_list, rgbtemplate.shape[1], rgbtemplate.shape[0])
|
|
print(f"box_list: {box_list}")
|
|
mask = convert_box_to_mask(box_list, [rgbimage.shape[0], rgbimage.shape[1]])
|
|
return (mask,)
|
|
|
|
|
|
class IsMaskEmptyNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
RETURN_TYPES = ["NUMBER"]
|
|
RETURN_NAMES = ["boolean_number"]
|
|
|
|
FUNCTION = "main"
|
|
CATEGORY = "template_matching"
|
|
|
|
def main(self, mask):
|
|
return (torch.all(mask == 0).int().item(), ) |