70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
from PIL import Image
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from typing import Tuple
|
|
import numpy as np
|
|
|
|
from .editor_ctx import EditorCtx
|
|
from .image_warper import ImageWarper
|
|
from .diff_algs import DiffAlgorithm, HeatmapDifference, HeatmapL2Norm, HeatmapSSIM
|
|
|
|
|
|
@dataclass
|
|
class PatchDetector:
|
|
"""
|
|
Detects differences between two images and generates a heatmap of the changes.
|
|
|
|
This class is responsible for comparing the original and edited images.
|
|
It first aligns the images if `allow_rescale` is true, then uses a selected
|
|
difference algorithm to calculate a heatmap of the changes.
|
|
"""
|
|
|
|
ctx: EditorCtx
|
|
|
|
@cached_property
|
|
def images(self) -> Tuple[Image.Image, Image.Image]:
|
|
img1 = self.ctx.original_img.convert('RGBA')
|
|
img2 = self.ctx.edited_img.convert('RGBA')
|
|
|
|
if self.ctx.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 detection_alg(self) -> DiffAlgorithm:
|
|
alg_map = {
|
|
"difference": HeatmapDifference,
|
|
"l2_norm": HeatmapL2Norm,
|
|
"ssim": HeatmapSSIM,
|
|
}
|
|
alg_class = alg_map.get(self.ctx.detection_alg)
|
|
|
|
if not alg_class:
|
|
raise ValueError(f"Unknown detection algorithm: {self.ctx.detection_alg}")
|
|
|
|
return alg_class(ctx=self.ctx)
|
|
|
|
@cached_property
|
|
def heatmap(self) -> Image.Image:
|
|
img_orig, img_edit = self.images
|
|
heatmap_arr = self.detection_alg.calculate(img_orig, img_edit)
|
|
|
|
# Handle Stride
|
|
if self.ctx.stride > 1:
|
|
# Subsample
|
|
small = heatmap_arr[:: self.ctx.stride, :: self.ctx.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
|