Stride
This commit is contained in:
Premik
+43
-34
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user