249 lines
8.1 KiB
Python
249 lines
8.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 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()
|