84 lines
3.3 KiB
Python
84 lines
3.3 KiB
Python
"""Restorer interface — "un-censor" detected regions of an image.
|
|
|
|
A Restorer takes an image plus the detected censored regions and returns a new
|
|
image with those regions reconstructed. This mirrors the ``Detector`` abstraction
|
|
so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug
|
|
in behind the same interface. Note DeepMosaics locates the mosaic itself, so the
|
|
``detections`` argument is currently advisory (unused by the DeepMosaics engines).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable
|
|
|
|
import numpy as np
|
|
|
|
from ..detection.types import Detection
|
|
|
|
# Optional cooperative-cancel hook: returns True to abort. Engines that loop or
|
|
# drive a subprocess should poll it; instant engines may ignore it.
|
|
CancelCheck = Callable[[], bool]
|
|
|
|
|
|
class Cancelled(Exception):
|
|
"""Raised by a Restorer when ``should_cancel`` asked it to stop."""
|
|
|
|
|
|
# Sequence (batch) restore callbacks — see ``Restorer.restore_sequence``.
|
|
FrameGetter = Callable[[int], np.ndarray] # index -> BGR image
|
|
DetGetter = Callable[[int], list[Detection]] # index -> that frame's detections
|
|
ResultSink = Callable[[int, np.ndarray], None] # (index, restored image) -> None
|
|
|
|
|
|
class Restorer(ABC):
|
|
#: Whether this engine uses *neighbouring* frames (so a batch run must feed it a
|
|
#: contiguous, ordered sequence — see :meth:`restore_sequence`). Per-frame engines
|
|
#: leave this False; the temporal DeepMosaics (BVDNet) sets it True.
|
|
temporal: bool = False
|
|
|
|
#: Whether this engine needs the frame's detections (it builds an inpaint mask from
|
|
#: them). DeepMosaics locates the mosaic itself, so it leaves this False and the
|
|
#: caller passes ``[]``; the diffusion engine sets it True and must be fed the real
|
|
#: detections (a frame with none comes back unchanged).
|
|
needs_detections: bool = False
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return type(self).__name__
|
|
|
|
@abstractmethod
|
|
def restore(
|
|
self,
|
|
image: np.ndarray,
|
|
detections: list[Detection],
|
|
should_cancel: CancelCheck | None = None,
|
|
) -> np.ndarray:
|
|
"""Return a copy of ``image`` with the detected regions reconstructed.
|
|
|
|
``should_cancel`` (if given) is polled periodically; when it returns
|
|
True the engine should abort and raise :class:`Cancelled`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def restore_sequence(
|
|
self,
|
|
count: int,
|
|
get_frame: FrameGetter,
|
|
get_dets: DetGetter,
|
|
emit: ResultSink,
|
|
should_cancel: CancelCheck | None = None,
|
|
) -> None:
|
|
"""Restore ``count`` frames *in order*, calling ``emit(i, restored)`` for each.
|
|
|
|
The default treats every frame independently (just loops :meth:`restore`).
|
|
Temporal engines override this to pull neighbouring frames via ``get_frame``
|
|
and carry recurrent state across the sequence. ``get_frame``/``get_dets`` are
|
|
lazy so the engine only reads the frames it needs; ``emit`` lets the caller
|
|
stream results to disk instead of holding them all in memory.
|
|
"""
|
|
for i in range(count):
|
|
if should_cancel is not None and should_cancel():
|
|
raise Cancelled("Восстановление отменено")
|
|
emit(i, self.restore(get_frame(i), get_dets(i), should_cancel))
|