diff --git a/dist/js/custom.js b/dist/js/custom.js index 30a96fb..03d2455 100644 --- a/dist/js/custom.js +++ b/dist/js/custom.js @@ -159,6 +159,14 @@ const customLabels = { scale_factor: 'scale factor', }, }, + MatchTemplate: { + output: { + detected_box: 'detected box', + debug_detection: 'debug - detection', + debug_template_mask: 'debug - template (mask)', + debug_template_image: 'debug - template (image)', + }, + }, }; const customLooks = { XYPlotQueue: { diff --git a/requirements.txt b/requirements.txt index 6c3a3c0..66adf04 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,5 @@ json5 python-dotenv pyyaml pyaml +opencv-python +webcolors diff --git a/src/python/collection/collect_nodes.py b/src/python/collection/collect_nodes.py index 51dbe9c..4c78af6 100644 --- a/src/python/collection/collect_nodes.py +++ b/src/python/collection/collect_nodes.py @@ -18,7 +18,7 @@ def import_all(load_env=True): if DEBUG: MODULES.append('debug') if MAD_SCIENTIST: - pass + MODULES.append('experimental') for mod in MODULES: importlib.import_module('...nodes.' + mod, package=__name__) diff --git a/src/python/nodes/experimental.py b/src/python/nodes/experimental.py new file mode 100644 index 0000000..a551c99 --- /dev/null +++ b/src/python/nodes/experimental.py @@ -0,0 +1,147 @@ +import torch +from torchvision.transforms import v2 as transforms +import cv2 +from webcolors import hex_to_rgb + +from ..collection.register_nodes import register_node +from ..shared.utils import tensor_to_cv, cv_to_tensor + +ALL_METHODS = { + # 'ccoeff': cv2.TM_CCOEFF, + 'ccoeff_normed': cv2.TM_CCOEFF_NORMED, + # 'ccorr': cv2.TM_CCORR, + 'ccorr_normed': cv2.TM_CCORR_NORMED, + # 'sqdiff': cv2.TM_SQDIFF, + 'sqdiff_normed': cv2.TM_SQDIFF_NORMED, +} + + +@register_node('Match Template (OpenCV)', 'experimental') +class MatchTemplate: + """Search for a template inside an image, using OpenCV Template matching method. + + - method: detection method, either 'ccoeff_normed', 'ccorr_normed' or 'sqdiff_normed'. + - Recommended: 'ccoeff_normed' + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + 'required': { + 'image': ('IMAGE',), + 'template': ('IMAGE',), + 'method': ( + list(ALL_METHODS.keys()), + { + 'default': 'ccoeff_normed', + 'tooltip': 'detection algorithm, default is ccoeff_normed', + }, + ), + 'threshold': ( + 'FLOAT', + {'min': 0, 'max': 1, 'default': 0.5, 'step': 0.01}, + ), + 'template_fill_hex': ( + 'STRING', + { + 'default': '#999', + 'tooltip': 'background color to apply to template', + }, + ), + } + } + + FUNCTION = 'run' + RETURN_TYPES = ('MASK', 'FLOAT', 'STRING', 'IMAGE', 'MASK', 'IMAGE') + RETURN_NAMES = ( + 'mask', + 'confidence', + 'detected_box', + 'debug_detection', + 'debug_template_mask', + 'debug_template_image', + ) + OUTPUT_TOOLTIPS = ( + 'mask', + 'confidence level', + 'detected box (x, y, w, h)', + 'detection image', + 'detection - template (mask)', + 'detection - template (image)', + ) + + DESCRIPTION = ( + 'Search for a template inside an image, using OpenCV Template matching method.' + ) + + def run( + self, + image: torch.Tensor, + template: torch.Tensor, + method, + threshold: float, + template_fill_hex: str, + ) -> torch.Tensor: + img_cv = tensor_to_cv(image) + tpl_cv = tensor_to_cv(template) + # tpl_cv = cv2.imread( + # '/mnt/nvm/enc/sd/comfyui/tests/20241029 - test logo/create mask/2 - logo_cleaned - cropped - transparent.png', cv2.IMREAD_UNCHANGED) + meth = ALL_METHODS[method] + + # fill the template with color + tpl_cv_mask = tpl_cv[:, :, 3] == 0 + tpl_cv = tpl_cv.copy() + (r, g, b) = hex_to_rgb(template_fill_hex) + tpl_cv[tpl_cv_mask] = [r, g, b, 255] + tpl_cv = cv2.cvtColor(tpl_cv, cv2.COLOR_BGRA2BGR) + + # detect template + result = cv2.matchTemplate(img_cv, tpl_cv, meth) + min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result) + + # if the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum + if meth in [cv2.TM_SQDIFF, cv2.TM_SQDIFF_NORMED]: + (x, y) = min_loc + else: + (x, y) = max_loc + + # if value is less than threshold, exit now + # if (max_val < threshold): + # return (image, max_val, None, None, None, None) + + # draw the rectangle + # w, h = tpl_cv.shape[::-1] # or 2? + w = tpl_cv.shape[1] + h = tpl_cv.shape[0] + + cv2.rectangle(img_cv, (x, y), (x + w, y + h), (255, 0, 0), (w + h) // 100 + 1) + + # create mask from template + template_sum = template.sum(axis=3) + tpl_mask = torch.where(template_sum > 0, 1.0, 0.0) + tpl_mask = 1.0 - tpl_mask + + # pad the mask + img_w = img_cv.shape[1] + img_h = img_cv.shape[0] + mask = transforms.functional.pad( + tpl_mask, + padding=(x, y, img_w - x - w, img_h - y - h), + fill=1, + padding_mode='constant', + ) + + # format box coordinates + box_coords = '{0}, {1}, {2}, {3}'.format(x, y, w, h) + + return ( + mask, + max_val, + box_coords, + cv_to_tensor(img_cv), + tpl_mask, + cv_to_tensor(tpl_cv), + ) diff --git a/src/python/nodes/plot.py b/src/python/nodes/plot.py index 124fd19..cf22247 100644 --- a/src/python/nodes/plot.py +++ b/src/python/nodes/plot.py @@ -90,10 +90,6 @@ class XYPlotQueue: index: int, dim2: list[str] = [''], ): - print('------------------------------------') - print('index = ', index) - print('------------------------------------') - if index < 0: # value has been reset (completion, interrupt, errors) index = 0 diff --git a/src/python/shared/utils.py b/src/python/shared/utils.py index c3fba49..3738495 100644 --- a/src/python/shared/utils.py +++ b/src/python/shared/utils.py @@ -16,16 +16,25 @@ class AnyReturnTypes(tuple): return ANY_TYPE -def tensor_to_pillow(tensor: torch.tensor) -> Image.Image: +def tensor_to_pillow(tensor: torch.Tensor) -> Image.Image: return Image.fromarray( np.clip(255.0 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) ) -def pillow_to_tensor(image: Image.Image) -> torch.tensor: +def pillow_to_tensor(image: Image.Image) -> torch.Tensor: return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) +def tensor_to_cv(tensor: torch.Tensor): + # return np.clip(255 * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) + return (tensor.numpy().squeeze(0) * 255).astype(np.uint8) # <- works + + +def cv_to_tensor(image) -> torch.Tensor: + return torch.from_numpy(image.astype(np.float32) / 255.0).unsqueeze(0) + + class EmptyNode: def __init__(self): pass diff --git a/src/ts/custom.ts b/src/ts/custom.ts index ad5e708..688b4af 100644 --- a/src/ts/custom.ts +++ b/src/ts/custom.ts @@ -188,6 +188,15 @@ const customLabels = { scale_factor: 'scale factor', }, }, + // EXPERIMENTAL + MatchTemplate: { + output: { + detected_box: 'detected box', + debug_detection: 'debug - detection', + debug_template_mask: 'debug - template (mask)', + debug_template_image: 'debug - template (image)', + }, + }, } as const as Record }> const customLooks = {