patch_detector.py

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