128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
import pillow_avif
|
|
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
|
|
|
|
|
|
@dataclass
|
|
class PatchDetector:
|
|
img_path_1: Path
|
|
img_path_2: Path
|
|
region_size: int = 8
|
|
|
|
@cached_property
|
|
def images(self) -> Tuple[Image.Image, Image.Image]:
|
|
img1 = Image.open(self.img_path_1).convert('RGBA')
|
|
img2 = Image.open(self.img_path_2).convert('RGBA')
|
|
|
|
if img1.size != img2.size:
|
|
raise ValueError(f"Images must be same size: {img1.size} != {img2.size}")
|
|
|
|
return img1, img2
|
|
|
|
@cached_property
|
|
def heatmap(self) -> Image.Image:
|
|
img1, img2 = self.images
|
|
# Calculate difference
|
|
diff = ImageChops.difference(img1, img2)
|
|
# Convert to grayscale to get intensity of difference
|
|
heatmap = diff.convert('L')
|
|
|
|
# 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
|
|
|
|
|
|
@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)
|
|
|
|
@property
|
|
def base_name(self) -> str:
|
|
return self.detector.img_path_1.stem
|
|
|
|
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")
|
|
|
|
# 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))
|
|
|
|
# Paste highlight using the raw mask
|
|
overlay.paste(highlight, (0, 0), self.raw_mask)
|
|
|
|
overlay.save(output_dir / f"{self.base_name}-overlay.png")
|
|
|
|
|
|
def main() -> None:
|
|
work_dir = Path('/tmp/work')
|
|
img1_path = work_dir / 'a.png'
|
|
img2_path = work_dir / 'a-patch.png'
|
|
output_path = work_dir / 'a-patched.png'
|
|
|
|
detector = PatchDetector(img1_path, img2_path)
|
|
merger = PatchMerger(detector)
|
|
|
|
print(f"Processing {img1_path} and {img2_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()
|