180 lines
6.1 KiB
Python
180 lines
6.1 KiB
Python
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 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
|
|
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 img1.size != img2.size:
|
|
if self.allow_rescale:
|
|
img2 = img2.resize(img1.size, resample=Image.Resampling.LANCZOS)
|
|
else:
|
|
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 = self._heatmap_difference(img1, img2)
|
|
elif self.detection_alg == "l2_norm":
|
|
heatmap = self._heatmap_l2_norm(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)
|
|
|
|
return heatmap
|
|
|
|
def _heatmap_difference(self, img1: Image.Image, img2: Image.Image) -> Image.Image:
|
|
# Calculate difference
|
|
diff = ImageChops.difference(img1, img2)
|
|
# Convert to grayscale to get intensity of difference
|
|
return diff.convert('L')
|
|
|
|
def _heatmap_l2_norm(self, img1: Image.Image, img2: Image.Image) -> Image.Image:
|
|
# Convert to numpy
|
|
a1 = np.array(img1, dtype=np.float32)
|
|
a2 = np.array(img2, 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))
|
|
|
|
# Clip to 255 and convert to uint8
|
|
dist = np.clip(dist, 0, 255).astype(np.uint8)
|
|
|
|
return Image.fromarray(dist, mode='L')
|
|
|
|
|
|
@dataclass
|
|
class PatchMerger:
|
|
detector: PatchDetector
|
|
min_area_size: int = 4
|
|
threshold: int = 10
|
|
blur_radius: float = 3.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 final_mask(self) -> Image.Image:
|
|
mask = self.raw_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
|
|
# Composite img2 onto img1 using mask
|
|
return Image.composite(img2, 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.raw_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()
|