nodes.py

import torch
import numpy as np
from PIL import Image
from typing import Tuple, List, Dict, Any
from .patch_editor import PatchEditor  # type: ignore


class ImgPatchNode:
    """
    A node for patching images by detecting differences and merging them.

    This node takes an original and an edited image, detects the changed areas,
    and merges the changes back into the original image. It provides several
    outputs for visualizing the process, including the final merged image,
    the mask used for merging, a heatmap of the differences, and a debug overlay.
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Any]:
        return {
            "required": {
                "original_image": ("IMAGE",),
                "edited_image": ("IMAGE",),
                "detection_alg": (["difference", "l2_norm", "ssim"],),
                "region_size": ("INT", {"default": 8, "min": 1, "max": 128, "tooltip": "Size of the blocks used for difference detection."}),
                "stride": (
                    "INT",
                    {
                        "default": 1,
                        "min": 1,
                        "max": 64,
                        "tooltip": "Step size for the sliding window. 1 means full overlap (slowest, most accurate), equal to region_size means no overlap (fastest).",
                    },
                ),
                "min_area_size": ("INT", {"default": 4, "min": 0, "max": 128, "tooltip": "Minimum size of the changed area to keep."}),
                "tolerance": (
                    "FLOAT",
                    {
                        "default": 0.5,
                        "min": 0.0,
                        "max": 5.0,
                        "step": 0.01,
                        "tooltip": "Sensitivity of difference detection. Higher values will detect more changes.",
                    },
                ),
                "blur_radius": (
                    "FLOAT",
                    {"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "Radius of the gaussian blur applied to the mask."},
                ),
                "mask_expansion": (
                    "INT",
                    {"default": 0, "min": -128, "max": 128, "tooltip": "Expand or shrink the mask. Positive values expand, negative values shrink."},
                ),
                "edit_max_alpha": (
                    "FLOAT",
                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Maximum alpha for the edited image overlay."},
                ),
                "allow_rescale": ("BOOLEAN", {"default": True, "tooltip": "If true, rescales the edited image to match original size."}),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "MASK", "IMAGE")
    RETURN_NAMES = ("merged_image", "final_mask", "heatmap", "raw_mask", "debug_overlay")
    # OUTPUT_IS_LIST = (False, True, True, True, True)
    OUTPUT_TOOLTIPS = (
        "The original image with the patched areas from the edited image.",
        "The final mask used for merging. White indicates changes (taken from edited image), black indicates unchanged areas.",
        "Visualization of differences. Brighter areas indicate larger differences between original and edited images.",
        "Initial binary mask before post-processing. White indicates detected changes, black indicates no change.",
        "Original image with a red overlay highlighting the detected changes.",
    )
    FUNCTION = "process"
    CATEGORY = "ImgPatch"

    def process(
        self,
        original_image: torch.Tensor,
        edited_image: torch.Tensor,
        detection_alg: str,
        region_size: int,
        stride: int,
        min_area_size: int,
        tolerance: float,
        blur_radius: float,
        mask_expansion: int,
        edit_max_alpha: float,
        allow_rescale: bool,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        print(f"original_image shape: {original_image.shape}")
        print(f"edited_image shape: {edited_image.shape}")
        if original_image.shape[0] > 1:
            raise ValueError("Only single image is supported for original_image")

        current_image_tensor = original_image
        out_final_masks: List[torch.Tensor] = []
        out_heatmaps: List[torch.Tensor] = []
        out_raw_masks: List[torch.Tensor] = []
        out_overlays: List[torch.Tensor] = []

        for i in range(edited_image.shape[0]):
            pil_1 = self._tensor_to_pil(current_image_tensor[0])
            pil_2 = self._tensor_to_pil(edited_image[i])

            editor = PatchEditor()
            editor.ctx.original_img = pil_1
            editor.ctx.edited_img = pil_2
            editor.ctx.detection_alg = detection_alg
            editor.ctx.region_size = region_size
            editor.ctx.stride = stride
            editor.ctx.min_area_size = min_area_size
            editor.ctx.tolerance = tolerance
            editor.ctx.blur_radius = blur_radius
            editor.ctx.mask_expansion = mask_expansion
            editor.ctx.edit_max_alpha = edit_max_alpha
            editor.ctx.allow_rescale = allow_rescale

            merger = editor.patch_merger
            current_image_tensor = self._pil_to_tensor(merger.merged_image)

            out_final_masks.append(self._mask_to_tensor(merger.final_mask))
            out_heatmaps.append(self._pil_to_tensor(editor.patch_detector.heatmap.convert("RGB")))
            out_raw_masks.append(self._mask_to_tensor(merger.raw_mask))
            out_overlays.append(self._pil_to_tensor(merger.debug_overlay))

        return (
            current_image_tensor,
            torch.cat(out_final_masks, dim=0),
            torch.cat(out_heatmaps, dim=0),
            torch.cat(out_raw_masks, dim=0),
            torch.cat(out_overlays, dim=0),
        )

    def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
        # tensor is [H, W, C]
        return Image.fromarray(np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8))

    def _pil_to_tensor(self, image: Image.Image) -> torch.Tensor:
        # image is PIL
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)

    def _mask_to_tensor(self, image: Image.Image) -> torch.Tensor:
        # image is PIL (L) -> Tensor [1, H, W]
        return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)