117 lines
3.0 KiB
Python
117 lines
3.0 KiB
Python
from PIL import Image
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from .patch_detector import PatchDetector
|
|
from .patch_merger import PatchMerger
|
|
|
|
|
|
@dataclass
|
|
class EditorCtx:
|
|
original_img: Image.Image
|
|
edited_img: Image.Image
|
|
detection_alg: str
|
|
region_size: int
|
|
stride: int
|
|
allow_rescale: bool
|
|
min_area_size: int
|
|
threshold: int
|
|
blur_radius: float
|
|
mask_expansion: int
|
|
edit_max_alpha: float
|
|
|
|
|
|
@dataclass
|
|
class PatchEditor:
|
|
ctx: EditorCtx
|
|
|
|
@cached_property
|
|
def patch_detector(self) -> PatchDetector:
|
|
return PatchDetector(
|
|
original_img=self.ctx.original_img,
|
|
edited_img=self.ctx.edited_img,
|
|
region_size=self.ctx.region_size,
|
|
stride=self.ctx.stride,
|
|
allow_rescale=self.ctx.allow_rescale,
|
|
detection_alg=self.ctx.detection_alg,
|
|
)
|
|
|
|
@cached_property
|
|
def patch_merger(self) -> PatchMerger:
|
|
return PatchMerger(
|
|
detector=self.patch_detector,
|
|
min_area_size=self.ctx.min_area_size,
|
|
threshold=self.ctx.threshold,
|
|
blur_radius=self.ctx.blur_radius,
|
|
mask_expansion=self.ctx.mask_expansion,
|
|
edit_max_alpha=self.ctx.edit_max_alpha,
|
|
)
|
|
|
|
@staticmethod
|
|
def from_node(
|
|
original_image: Image.Image,
|
|
edited_image: Image.Image,
|
|
detection_alg: str,
|
|
region_size: int,
|
|
stride: int,
|
|
min_area_size: int,
|
|
threshold: int,
|
|
blur_radius: float,
|
|
mask_expansion: int,
|
|
edit_max_alpha: float,
|
|
allow_rescale: bool,
|
|
) -> "PatchEditor":
|
|
ctx = EditorCtx(
|
|
original_img=original_image,
|
|
edited_img=edited_image,
|
|
detection_alg=detection_alg,
|
|
region_size=region_size,
|
|
stride=stride,
|
|
allow_rescale=allow_rescale,
|
|
min_area_size=min_area_size,
|
|
threshold=threshold,
|
|
blur_radius=blur_radius,
|
|
mask_expansion=mask_expansion,
|
|
edit_max_alpha=edit_max_alpha,
|
|
)
|
|
return PatchEditor(ctx)
|
|
|
|
|
|
def main() -> None:
|
|
work_dir = Path('/tmp/work')
|
|
original_path = work_dir / 'a.png'
|
|
edited_path = work_dir / 'a-patch.png'
|
|
output_path = work_dir / 'a-patched.png'
|
|
|
|
original_img = Image.open(original_path)
|
|
edited_img = Image.open(edited_path)
|
|
|
|
editor = PatchEditor.from_node(
|
|
original_image=original_img,
|
|
edited_image=edited_img,
|
|
detection_alg="difference",
|
|
region_size=8,
|
|
stride=1,
|
|
min_area_size=4,
|
|
threshold=10,
|
|
blur_radius=3.0,
|
|
mask_expansion=0,
|
|
edit_max_alpha=1.0,
|
|
allow_rescale=False,
|
|
)
|
|
|
|
merger = editor.patch_merger
|
|
|
|
print(f"Processing {original_path} and {edited_path}...")
|
|
merger.save_merged(output_path)
|
|
print(f"Saved merged image to {output_path}")
|
|
|
|
merger.save_debug(work_dir)
|
|
print(f"Saved debug images to {work_dir}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|