from PIL import Image, ImageChops, ImageFilter, ImageDraw from dataclasses import dataclass, field 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 from .image_warper import ImageWarper # from image_warper import ImageWarper # import pillow_avif @dataclass class PatchDetector: original_img_path: Optional[Path] = None edited_img_path: Optional[Path] = None 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" def __post_init__(self) -> None: if self.original_img is None and self.original_img_path is not None: self.original_img = Image.open(self.original_img_path).convert('RGBA') if self.edited_img is None and self.edited_img_path is not None: self.edited_img = Image.open(self.edited_img_path).convert('RGBA') if self.original_img is None or self.edited_img is None: raise ValueError("Images must be provided either via paths or direct Image objects") @cached_property def images(self) -> Tuple[Image.Image, Image.Image]: assert self.original_img is not None assert self.edited_img is not None img1 = self.original_img.convert('RGBA') img2 = self.edited_img.convert('RGBA') if self.allow_rescale: warper = ImageWarper() img2 = warper.warp(img1, img2) if img1.size != img2.size: raise ValueError(f"Images must be same size: {img1.size} != {img2.size}. Enable allow_rescale to fix automatically.") return img1, img2 @cached_property def heatmap(self) -> Image.Image: img1, img2 = self.images if self.detection_alg == "difference": heatmap_arr = self._heatmap_difference(img1, img2) elif self.detection_alg == "l2_norm": heatmap_arr = self._heatmap_l2_norm(img1, img2) elif self.detection_alg == "ssim": heatmap_arr = self._heatmap_ssim(img1, img2) else: raise ValueError(f"Unknown detection algorithm: {self.detection_alg}") # 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) -> np.ndarray: # Calculate difference diff = ImageChops.difference(img1, img2) # Convert to grayscale to get intensity of difference diff_l = diff.convert('L') arr = np.array(diff_l, dtype=np.float32) 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.convert('RGB'), dtype=np.float32) a2 = np.array(img2.convert('RGB'), dtype=np.float32) # Calculate squared difference diff_sq = np.sum((a1 - a2) ** 2, axis=2) 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 np.sqrt(mse) 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. 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, '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 # Average across channels if diff_map.ndim == 3: diff_map = np.mean(diff_map, axis=2) return diff_map @dataclass class PatchMerger: detector: PatchDetector min_area_size: int = 4 threshold: int = 10 blur_radius: float = 3.0 mask_expansion: int = 0 edit_max_alpha: float = 1.0 @cached_property def raw_mask(self) -> Image.Image: heatmap = self.detector.heatmap # Apply threshold mask = heatmap.point(lambda p: 255 if p > self.threshold else 0) # Drop smaller areas if self.min_area_size > 1: # Filter size must be odd filter_size = self.min_area_size | 1 # Erode then Dilate (Opening) to remove small noise mask = mask.filter(ImageFilter.MinFilter(filter_size)) mask = mask.filter(ImageFilter.MaxFilter(filter_size)) return mask @cached_property def sized_mask(self) -> Image.Image: mask = self.raw_mask if self.mask_expansion == 0: return mask # The kernel size for a radius of N is 2*N+1 size = abs(self.mask_expansion) * 2 + 1 if self.mask_expansion > 0: # Expand (Dilation) return mask.filter(ImageFilter.MaxFilter(size)) else: # Shrink (Erosion) return mask.filter(ImageFilter.MinFilter(size)) @cached_property def final_mask(self) -> Image.Image: mask = self.sized_mask if self.blur_radius > 0: mask = mask.filter(ImageFilter.GaussianBlur(self.blur_radius)) return mask @cached_property def merged_image(self) -> Image.Image: img1, img2 = self.detector.images mask = self.final_mask if self.edit_max_alpha >= 1.0: # Composite img2 onto img1 using mask return Image.composite(img2, img1, mask) # Blend original and edited images blended_img = Image.blend(img1, img2, self.edit_max_alpha) return Image.composite(blended_img, img1, mask) @cached_property def debug_overlay(self) -> Image.Image: # Overlay img1 = self.detector.images[0] overlay = img1.copy().convert('RGBA') # Create a red highlight layer highlight = Image.new('RGBA', overlay.size, (255, 0, 0, 128)) # Blend highlight over overlay blended = Image.alpha_composite(overlay, highlight) # Composite using mask return Image.composite(blended, overlay, self.sized_mask) @property def base_name(self) -> str: if self.detector.original_img_path: return self.detector.original_img_path.stem return "patch_result" def save_merged(self, path: Path) -> None: self.merged_image.save(path, qmin=0, qmax=0) def save_debug(self, output_dir: Path) -> None: output_dir.mkdir(parents=True, exist_ok=True) self.detector.heatmap.save(output_dir / f"{self.base_name}-heatmap.png") self.raw_mask.save(output_dir / f"{self.base_name}-raw_mask.png") self.final_mask.save(output_dir / f"{self.base_name}-final_mask.png") self.debug_overlay.save(output_dir / f"{self.base_name}-overlay.png") def main() -> None: work_dir = Path('/tmp/work') original_path = work_dir / 'a.png' edited_path = work_dir / 'a-patch.png' output_path = work_dir / 'a-patched.png' detector = PatchDetector(original_path, edited_path) merger = PatchMerger(detector) print(f"Processing {original_path} and {edited_path}...") merger.save_merged(output_path) print(f"Saved merged image to {output_path}") merger.save_debug(work_dir) print(f"Saved debug images to {work_dir}") if __name__ == '__main__': main()