94 lines
2.5 KiB
Python
94 lines
2.5 KiB
Python
from PIL import Image
|
|
from dataclasses import dataclass
|
|
from functools import cached_property
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from .editor_ctx import EditorCtx
|
|
from .patch_detector import PatchDetector
|
|
from .patch_merger import PatchMerger
|
|
|
|
|
|
@dataclass
|
|
class PatchEditor:
|
|
ctx: EditorCtx
|
|
|
|
@cached_property
|
|
def patch_detector(self) -> PatchDetector:
|
|
return PatchDetector(self.ctx)
|
|
|
|
@cached_property
|
|
def patch_merger(self) -> PatchMerger:
|
|
return PatchMerger(self.patch_detector)
|
|
|
|
@staticmethod
|
|
def from_node(
|
|
original_image: Image.Image,
|
|
edited_image: Image.Image,
|
|
detection_alg: str,
|
|
region_size: int,
|
|
stride: int,
|
|
min_area_size: int,
|
|
tolerance: float,
|
|
blur_radius: float,
|
|
mask_expansion: int,
|
|
edit_max_alpha: float,
|
|
allow_rescale: bool,
|
|
original_image_path: Optional[Path] = None,
|
|
edited_image_path: Optional[Path] = None,
|
|
) -> "PatchEditor":
|
|
ctx = EditorCtx(
|
|
original_img=original_image,
|
|
edited_img=edited_image,
|
|
original_img_path=original_image_path,
|
|
edited_img_path=edited_image_path,
|
|
detection_alg=detection_alg,
|
|
region_size=region_size,
|
|
stride=stride,
|
|
allow_rescale=allow_rescale,
|
|
min_area_size=min_area_size,
|
|
tolerance=tolerance,
|
|
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,
|
|
original_image_path=original_path,
|
|
detection_alg="difference",
|
|
region_size=8,
|
|
stride=1,
|
|
min_area_size=4,
|
|
tolerance=0.5,
|
|
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()
|