80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
from PIL import Image
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
from typing import Tuple, Optional
|
|
import numpy as np
|
|
|
|
from .image_warper import ImageWarper
|
|
from .diff_algs import HeatmapDifference, HeatmapL2Norm, HeatmapSSIM
|
|
|
|
|
|
@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
|
|
|
|
alg_map = {
|
|
"difference": HeatmapDifference,
|
|
"l2_norm": HeatmapL2Norm,
|
|
"ssim": HeatmapSSIM,
|
|
}
|
|
alg_class = alg_map.get(self.detection_alg)
|
|
|
|
if not alg_class:
|
|
raise ValueError(f"Unknown detection algorithm: {self.detection_alg}")
|
|
|
|
alg = alg_class(ctx=self)
|
|
heatmap_arr = alg.calculate(img1, img2)
|
|
|
|
# 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
|