119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
from PIL import Image, ImageFilter, ImageDraw
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
import numpy as np
|
|
|
|
from .patch_detector import PatchDetector
|
|
|
|
|
|
from .editor_ctx import EditorCtx
|
|
|
|
|
|
@dataclass
|
|
class PatchMerger:
|
|
"""
|
|
Merges two images based on a detected patch, with various processing options.
|
|
|
|
This class takes the heatmap from the PatchDetector and processes it to create
|
|
a final mask. This mask is then used to composite the edited image onto the
|
|
original image. It also provides methods for saving the merged image and
|
|
debug artifacts.
|
|
"""
|
|
|
|
detector: PatchDetector
|
|
|
|
@cached_property
|
|
def ctx(self) -> EditorCtx:
|
|
return self.detector.ctx
|
|
|
|
@cached_property
|
|
def raw_mask(self) -> Image.Image:
|
|
heatmap = self.detector.heatmap
|
|
threshold = self.detector.detection_alg.get_threshold(self.ctx.tolerance)
|
|
# Apply threshold
|
|
mask = heatmap.point(lambda p: 255 if p > threshold else 0)
|
|
|
|
# Drop smaller areas
|
|
if self.ctx.min_area_size > 1:
|
|
# Filter size must be odd
|
|
filter_size = self.ctx.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 sized_mask(self) -> Image.Image:
|
|
mask = self.raw_mask
|
|
if self.ctx.mask_expansion == 0:
|
|
return mask
|
|
|
|
# The kernel size for a radius of N is 2*N+1
|
|
size = abs(self.ctx.mask_expansion) * 2 + 1
|
|
|
|
if self.ctx.mask_expansion > 0:
|
|
# Expand (Dilation)
|
|
return mask.filter(ImageFilter.MaxFilter(size))
|
|
else:
|
|
# Shrink (Erosion)
|
|
return mask.filter(ImageFilter.MinFilter(size))
|
|
|
|
@cached_property
|
|
def final_mask(self) -> Image.Image:
|
|
mask = self.sized_mask
|
|
if self.ctx.blur_radius > 0:
|
|
mask = mask.filter(ImageFilter.GaussianBlur(self.ctx.blur_radius))
|
|
return mask
|
|
|
|
@cached_property
|
|
def merged_image(self) -> Image.Image:
|
|
img_orig, img_edit = self.detector.images
|
|
mask = self.final_mask
|
|
|
|
img_orig_rgba = img_orig.convert('RGBA')
|
|
img_edit_rgba = img_edit.convert('RGBA')
|
|
|
|
if self.ctx.edit_max_alpha < 1.0:
|
|
img_edit_rgba = Image.blend(img_orig_rgba, img_edit_rgba, self.ctx.edit_max_alpha)
|
|
|
|
# Create a new image for the patch to be applied, with the mask as alpha
|
|
patch_img = img_edit_rgba.copy()
|
|
patch_img.putalpha(mask)
|
|
|
|
# Composite the patch onto the original image
|
|
return Image.alpha_composite(img_orig_rgba, patch_img)
|
|
|
|
@cached_property
|
|
def debug_overlay(self) -> Image.Image:
|
|
# 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))
|
|
|
|
# Blend highlight over overlay
|
|
blended = Image.alpha_composite(overlay, highlight)
|
|
|
|
# Composite using mask
|
|
return Image.composite(blended, overlay, self.sized_mask)
|
|
|
|
@property
|
|
def base_name(self) -> str:
|
|
if self.ctx.original_img_path:
|
|
return self.ctx.original_img_path.stem
|
|
return "patch_result"
|
|
|
|
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")
|
|
self.debug_overlay.save(output_dir / f"{self.base_name}-overlay.png")
|