86 lines
2.3 KiB
Python
86 lines
2.3 KiB
Python
from PIL import Image, ImageChops
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
import numpy as np
|
|
import scipy.ndimage
|
|
from skimage.metrics import structural_similarity
|
|
|
|
if TYPE_CHECKING:
|
|
from .patch_editor import EditorCtx
|
|
|
|
|
|
@dataclass
|
|
class DiffAlgorithm:
|
|
ctx: "EditorCtx"
|
|
half: int = 0
|
|
full: int = 255
|
|
|
|
def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
|
|
raise NotImplementedError
|
|
|
|
def get_threshold(self, tolerance: float) -> int:
|
|
if tolerance <= 1.0:
|
|
return int(self.half * tolerance)
|
|
else:
|
|
return int(self.half + (self.full - self.half) * (tolerance - 1.0))
|
|
|
|
|
|
@dataclass
|
|
class HeatmapDifference(DiffAlgorithm):
|
|
half: int = 10
|
|
full: int = 40
|
|
|
|
def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
|
|
diff = ImageChops.difference(img1, img2)
|
|
diff_l = diff.convert('L')
|
|
arr = np.array(diff_l, dtype=np.float32)
|
|
|
|
if self.ctx.region_size > 1:
|
|
arr = scipy.ndimage.uniform_filter(arr, size=self.ctx.region_size, mode='reflect')
|
|
|
|
return arr
|
|
|
|
|
|
@dataclass
|
|
class HeatmapL2Norm(DiffAlgorithm):
|
|
half: int = 10
|
|
full: int = 100
|
|
|
|
def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
|
|
a1 = np.array(img1.convert('RGB'), dtype=np.float32)
|
|
a2 = np.array(img2.convert('RGB'), dtype=np.float32)
|
|
|
|
diff_sq = np.sum((a1 - a2) ** 2, axis=2)
|
|
|
|
if self.ctx.region_size > 1:
|
|
mse = scipy.ndimage.uniform_filter(diff_sq, size=self.ctx.region_size, mode='reflect')
|
|
else:
|
|
mse = diff_sq
|
|
|
|
return np.sqrt(mse)
|
|
|
|
|
|
@dataclass
|
|
class HeatmapSSIM(DiffAlgorithm):
|
|
half: int = 40
|
|
full: int = 100
|
|
|
|
def calculate(self, img1: Image.Image, img2: Image.Image) -> np.ndarray:
|
|
a1 = np.array(img1.convert('RGB'))
|
|
a2 = np.array(img2.convert('RGB'))
|
|
|
|
win_size = self.ctx.region_size if self.ctx.region_size % 2 == 1 else self.ctx.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_map = (1 - diff) * 255
|
|
|
|
if diff_map.ndim == 3:
|
|
diff_map = np.mean(diff_map, axis=2)
|
|
|
|
return diff_map
|