diff --git a/img_patch_editor.py b/img_patch_editor.py index 0b7493d..0c53f5b 100644 --- a/img_patch_editor.py +++ b/img_patch_editor.py @@ -4,6 +4,7 @@ from functools import cached_property from pathlib import Path from typing import Tuple, Optional import numpy as np +import scipy.ndimage from skimage.metrics import structural_similarity @@ -17,6 +18,7 @@ class PatchDetector: original_img: Optional[Image.Image] = None edited_img: Optional[Image.Image] = None region_size: int = 8 + stride: int = 1 allow_rescale: bool = False detection_alg: str = "difference" @@ -51,73 +53,80 @@ class PatchDetector: img1, img2 = self.images if self.detection_alg == "difference": - heatmap = self._heatmap_difference(img1, img2) + heatmap_arr = self._heatmap_difference(img1, img2) elif self.detection_alg == "l2_norm": - heatmap = self._heatmap_l2_norm(img1, img2) + heatmap_arr = self._heatmap_l2_norm(img1, img2) elif self.detection_alg == "ssim": - heatmap = self._heatmap_ssim(img1, img2) + heatmap_arr = self._heatmap_ssim(img1, img2) else: raise ValueError(f"Unknown detection algorithm: {self.detection_alg}") - # Enforce region size (simple block approach) - if self.region_size > 1: - orig_size = heatmap.size - # Downscale to average out blocks - # Using BOX resampling for averaging - small = heatmap.resize((max(1, orig_size[0] // self.region_size), max(1, orig_size[1] // self.region_size)), resample=Image.Resampling.BOX) - # Upscale to restore size but with blocky look - heatmap = small.resize(orig_size, resample=Image.Resampling.NEAREST) + # Handle Stride + if self.stride > 1: + # Subsample + small = heatmap_arr[:: self.stride, :: self.stride] + + # Upscale + h, w = heatmap_arr.shape + small_img = Image.fromarray(np.clip(small, 0, 255).astype(np.uint8), mode='L') + heatmap = small_img.resize((w, h), resample=Image.Resampling.NEAREST) + else: + heatmap = Image.fromarray(np.clip(heatmap_arr, 0, 255).astype(np.uint8), mode='L') return heatmap - def _heatmap_difference(self, img1: Image.Image, img2: Image.Image) -> Image.Image: + def _heatmap_difference(self, img1: Image.Image, img2: Image.Image) -> np.ndarray: # Calculate difference diff = ImageChops.difference(img1, img2) # Convert to grayscale to get intensity of difference - return diff.convert('L') + diff_l = diff.convert('L') + arr = np.array(diff_l, dtype=np.float32) - def _heatmap_l2_norm(self, img1: Image.Image, img2: Image.Image) -> Image.Image: + if self.region_size > 1: + arr = scipy.ndimage.uniform_filter(arr, size=self.region_size, mode='reflect') + + return arr + + def _heatmap_l2_norm(self, img1: Image.Image, img2: Image.Image) -> np.ndarray: # Convert to numpy - a1 = np.array(img1, dtype=np.float32) - a2 = np.array(img2, dtype=np.float32) + a1 = np.array(img1.convert('RGB'), dtype=np.float32) + a2 = np.array(img2.convert('RGB'), dtype=np.float32) - # Calculate L2 norm of difference (using RGB channels, ignoring Alpha to match 'difference' behavior) - diff = a1[:, :, :3] - a2[:, :, :3] - dist = np.sqrt(np.sum(diff**2, axis=2)) + # Calculate squared difference + diff_sq = np.sum((a1 - a2) ** 2, axis=2) - # Clip to 255 and convert to uint8 - dist = np.clip(dist, 0, 255).astype(np.uint8) + if self.region_size > 1: + # Mean squared error over region + mse = scipy.ndimage.uniform_filter(diff_sq, size=self.region_size, mode='reflect') + else: + mse = diff_sq - return Image.fromarray(dist, mode='L') + return np.sqrt(mse) - def _heatmap_ssim(self, img1: Image.Image, img2: Image.Image) -> Image.Image: + def _heatmap_ssim(self, img1: Image.Image, img2: Image.Image) -> np.ndarray: # Convert to numpy and RGB (ignore alpha) a1 = np.array(img1.convert('RGB')) a2 = np.array(img2.convert('RGB')) # Calculate SSIM - # win_size needs to be odd. Default is 7. - min_side = min(a1.shape[0], a1.shape[1]) - win_size = min(7, min_side) - if win_size % 2 == 0: - win_size -= 1 + # win_size needs to be odd. + win_size = self.region_size if self.region_size % 2 == 1 else self.region_size + 1 + if win_size < 3: + win_size = 3 - kwargs = {'channel_axis': 2, 'data_range': 255} - if win_size >= 3: - kwargs['win_size'] = win_size + kwargs = {'channel_axis': 2, 'data_range': 255, 'win_size': win_size} _, diff = structural_similarity(a1, a2, full=True, **kwargs) # diff is the SSIM map (1 is identical, lower is different) # Invert to make differences bright diff_map = (1 - diff) * 255 - diff_map = np.clip(diff_map, 0, 255).astype(np.uint8) # Average across channels if diff_map.ndim == 3: - diff_map = np.mean(diff_map, axis=2).astype(np.uint8) + diff_map = np.mean(diff_map, axis=2) - return Image.fromarray(diff_map, mode='L') + return diff_map @dataclass diff --git a/nodes.py b/nodes.py index b023d76..c5cb572 100644 --- a/nodes.py +++ b/nodes.py @@ -14,6 +14,15 @@ class ImgPatchNode: "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": ( @@ -42,6 +51,7 @@ class ImgPatchNode: edited_image: torch.Tensor, detection_alg: str, region_size: int, + stride: int, min_area_size: int, threshold: int, blur_radius: float, @@ -71,6 +81,7 @@ class ImgPatchNode: original_img=pil_1, edited_img=pil_2, region_size=region_size, + stride=stride, allow_rescale=allow_rescale, detection_alg=detection_alg, )