patch_editor.py

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:
    """
    Orchestrates the patch detection and merging process.

    This class acts as a high-level container for the patch editing workflow.
    It initializes the PatchDetector and PatchMerger, and provides a convenient
    factory method to be used from the ComfyUI node.
    """

    @cached_property
    def ctx(self) -> EditorCtx:
        return EditorCtx()

    @cached_property
    def patch_detector(self) -> PatchDetector:
        return PatchDetector(self.ctx)

    @cached_property
    def patch_merger(self) -> PatchMerger:
        return PatchMerger(self.patch_detector)


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()
    editor.ctx.original_img = original_img
    editor.ctx.edited_img = edited_img

    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()