132 lines
5.5 KiB
Python
132 lines
5.5 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Tuple, List, Dict, Any
|
|
from .img_patch_editor import PatchDetector, PatchMerger
|
|
|
|
|
|
class ImgPatchNode:
|
|
@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."}),
|
|
"threshold": ("INT", {"default": 10, "min": 0, "max": 255, "tooltip": "Sensitivity of difference detection."}),
|
|
"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": False, "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_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,
|
|
threshold: int,
|
|
blur_radius: float,
|
|
mask_expansion: int,
|
|
edit_max_alpha: float,
|
|
allow_rescale: bool,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
|
|
batch_size = original_image.shape[0]
|
|
img2_batch_size = edited_image.shape[0]
|
|
|
|
out_merged: List[torch.Tensor] = []
|
|
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(batch_size):
|
|
# Handle broadcasting for edited_image
|
|
idx2 = i % img2_batch_size
|
|
|
|
t1 = original_image[i]
|
|
t2 = edited_image[idx2]
|
|
|
|
pil_1 = self._tensor_to_pil(t1)
|
|
pil_2 = self._tensor_to_pil(t2)
|
|
|
|
detector = PatchDetector(
|
|
original_img=pil_1,
|
|
edited_img=pil_2,
|
|
region_size=region_size,
|
|
stride=stride,
|
|
allow_rescale=allow_rescale,
|
|
detection_alg=detection_alg,
|
|
)
|
|
merger = PatchMerger(
|
|
detector=detector,
|
|
min_area_size=min_area_size,
|
|
threshold=threshold,
|
|
blur_radius=blur_radius,
|
|
mask_expansion=mask_expansion,
|
|
edit_max_alpha=edit_max_alpha,
|
|
)
|
|
|
|
out_merged.append(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(merger.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 (
|
|
torch.cat(out_merged, dim=0),
|
|
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)
|