Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.
This commit is contained in:
Leonid Pershin
@@ -1,9 +1,10 @@
|
||||
"""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/filled. This mirrors the ``Detector``
|
||||
abstraction so different engines (classic inpaint now; a generative model like
|
||||
DeepMosaics / LADA later) plug in behind the same interface.
|
||||
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
|
||||
@@ -24,7 +25,18 @@ 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
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return type(self).__name__
|
||||
@@ -42,3 +54,24 @@ class Restorer(ABC):
|
||||
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))
|
||||
|
||||
@@ -26,18 +26,29 @@ from types import SimpleNamespace
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Cancelled, Restorer
|
||||
from .base import (
|
||||
CancelCheck,
|
||||
Cancelled,
|
||||
DetGetter,
|
||||
FrameGetter,
|
||||
ResultSink,
|
||||
Restorer,
|
||||
)
|
||||
|
||||
_VENDOR = Path(__file__).parent / "_deepmosaics"
|
||||
# Default place to drop DeepMosaics clean weights (gitignored — see models/).
|
||||
DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"
|
||||
|
||||
|
||||
def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
|
||||
"""Find usable per-frame clean models: (display_name, full_path).
|
||||
def discover_models(
|
||||
extra_dir: str | None = None, include_video: bool = False
|
||||
) -> list[tuple[str, str]]:
|
||||
"""Find usable clean models: (display_name, full_path).
|
||||
|
||||
Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given)
|
||||
for ``clean_*.pth``. The video model is skipped — it can't run per-frame.
|
||||
for ``clean_*.pth``. By default the video model is skipped — it can't run
|
||||
per-frame; pass ``include_video=True`` for the temporal (BVDNet) engine, which
|
||||
*needs* ``clean_*_video.pth``.
|
||||
"""
|
||||
dirs = [DEFAULT_WEIGHTS_DIR]
|
||||
if extra_dir:
|
||||
@@ -48,13 +59,23 @@ def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
|
||||
if not d.is_dir():
|
||||
continue
|
||||
for p in sorted(d.glob("clean_*.pth")):
|
||||
if "video" in p.name.lower() or p.name in seen:
|
||||
if p.name in seen:
|
||||
continue
|
||||
if "video" in p.name.lower() and not include_video:
|
||||
continue
|
||||
seen.add(p.name)
|
||||
out.append((p.stem, str(p)))
|
||||
return out
|
||||
|
||||
|
||||
def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
|
||||
"""Locate ``mosaic_position.pth`` (the BiSeNet mosaic locator) for ``model``."""
|
||||
candidates = [model.parent / "mosaic_position.pth"]
|
||||
if dm_dir:
|
||||
candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
|
||||
return next((p for p in candidates if p.is_file()), None)
|
||||
|
||||
|
||||
def _netg_kind(model_name: str) -> str:
|
||||
"""Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
|
||||
n = model_name.lower()
|
||||
@@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer):
|
||||
)
|
||||
model = Path(model_path)
|
||||
self._netg = _netg_kind(model.name) # raises on a video model
|
||||
pos = self._find_mosaic_position(model, deepmosaics_dir)
|
||||
pos = _find_mosaic_position(model, deepmosaics_dir)
|
||||
if pos is None:
|
||||
raise ValueError(
|
||||
"Рядом с clean-моделью не найден mosaic_position.pth.\n"
|
||||
@@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer):
|
||||
self._gpu = gpu_id
|
||||
self._loaded = False # models loaded lazily on first restore
|
||||
|
||||
@staticmethod
|
||||
def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
|
||||
candidates = [model.parent / "mosaic_position.pth"]
|
||||
if dm_dir:
|
||||
candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
|
||||
return next((p for p in candidates if p.is_file()), None)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"DeepMosaics(gpu={self._gpu})"
|
||||
@@ -167,3 +181,147 @@ class DeepMosaicsRestorer(Restorer):
|
||||
img_mosaic = work[y - size:y + size, x - size:x + size]
|
||||
img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
|
||||
return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather)
|
||||
|
||||
|
||||
class DeepMosaicsVideoRestorer(Restorer):
|
||||
"""Temporal DeepMosaics (BVDNet) — un-censors using *neighbouring* frames.
|
||||
|
||||
Reproduces DeepMosaics' ``cleanmosaic_video_fusion`` per target frame: for frame
|
||||
``i`` it feeds the network a temporal window of ``T`` frames (sampled at step ``S``
|
||||
around ``i``) plus its own previous output (recurrent), so the reconstruction is
|
||||
temporally coherent. Because of that recurrence the frames MUST be processed in
|
||||
order over a contiguous range — see :meth:`restore_sequence` (the batch run).
|
||||
|
||||
Needs the **video** weights ``clean_youknow_video.pth`` + ``mosaic_position.pth``
|
||||
(beside it). Single-frame :meth:`restore` degrades to a window of the same frame.
|
||||
"""
|
||||
|
||||
temporal = True
|
||||
|
||||
# DeepMosaics fusion window: N before/after at step S → T = 2N+1 frames, INPUT_SIZE px.
|
||||
_N, _T, _S = 2, 5, 3
|
||||
_INPUT_SIZE = 256
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
deepmosaics_dir: str | None,
|
||||
model_path: str | None,
|
||||
python_exe: str | None = None, # unused (in-process); kept for factory parity
|
||||
gpu_id: str = "0",
|
||||
) -> None:
|
||||
chosen: Path | None = None
|
||||
if model_path and Path(model_path).is_file() and "video" in Path(model_path).name.lower():
|
||||
chosen = Path(model_path)
|
||||
else: # configured model missing or not a video model → auto-pick a video model
|
||||
vids = [p for _n, p in discover_models(include_video=True) if "video" in Path(p).name.lower()]
|
||||
if vids:
|
||||
chosen = Path(vids[0])
|
||||
if chosen is None:
|
||||
raise ValueError(
|
||||
"Не найдены веса видеомодели DeepMosaics (clean_*_video.pth).\n"
|
||||
"Положите clean_youknow_video.pth + mosaic_position.pth в models/deepmosaics "
|
||||
"(или выберите в «Восстановление…»). См. README."
|
||||
)
|
||||
pos = _find_mosaic_position(chosen, deepmosaics_dir)
|
||||
if pos is None:
|
||||
raise ValueError(
|
||||
"Рядом с видеомоделью не найден mosaic_position.pth.\n"
|
||||
"Положите mosaic_position.pth в ту же папку, что и clean_*_video.pth. См. README."
|
||||
)
|
||||
self._model = str(chosen)
|
||||
self._pos = str(pos)
|
||||
self._gpu = gpu_id
|
||||
self._loaded = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"DeepMosaicsVideo(gpu={self._gpu})"
|
||||
|
||||
# ------------------------------------------------------------------ engine
|
||||
def _ensure_loaded(self) -> None:
|
||||
if self._loaded:
|
||||
return
|
||||
if str(_VENDOR) not in sys.path:
|
||||
sys.path.insert(0, str(_VENDOR))
|
||||
|
||||
import torch # noqa: E402
|
||||
if self._gpu != "-1" and not torch.cuda.is_available():
|
||||
self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why)
|
||||
|
||||
from models import loadmodel, runmodel # type: ignore # noqa: E402
|
||||
import util.data as data # type: ignore # noqa: E402
|
||||
import util.image_processing as impro # type: ignore # noqa: E402
|
||||
|
||||
self._torch = torch
|
||||
self._runmodel = runmodel
|
||||
self._data = data
|
||||
self._impro = impro
|
||||
self._opt = SimpleNamespace(
|
||||
gpu_id=self._gpu,
|
||||
model_path=self._model,
|
||||
mosaic_position_model_path=self._pos,
|
||||
mask_threshold=64,
|
||||
all_mosaic_area=False,
|
||||
ex_mult=1.5,
|
||||
no_feather=False,
|
||||
)
|
||||
self._netM = loadmodel.bisenet(self._opt, "mosaic")
|
||||
self._netG = loadmodel.video(self._opt) # BVDNet
|
||||
self._loaded = True
|
||||
|
||||
def restore(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
detections: list[Detection],
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Single-frame restore — no neighbours, so the window is the same frame."""
|
||||
out: dict[int, np.ndarray] = {}
|
||||
self.restore_sequence(
|
||||
1,
|
||||
lambda _i: image,
|
||||
lambda _i: detections,
|
||||
lambda i, r: out.__setitem__(i, r),
|
||||
should_cancel,
|
||||
)
|
||||
return out.get(0, image.copy())
|
||||
|
||||
def restore_sequence(
|
||||
self,
|
||||
count: int,
|
||||
get_frame: FrameGetter,
|
||||
get_dets: DetGetter,
|
||||
emit: ResultSink,
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> None:
|
||||
self._ensure_loaded()
|
||||
torch, data, impro, opt = self._torch, self._data, self._impro, self._opt
|
||||
N, T, S, SZ = self._N, self._T, self._S, self._INPUT_SIZE
|
||||
|
||||
previous = None # recurrent state: the network's previous output (a tensor)
|
||||
for i in range(count):
|
||||
if should_cancel is not None and should_cancel():
|
||||
raise Cancelled("Восстановление отменено")
|
||||
img_origin = get_frame(i)
|
||||
x, y, size, mask = self._runmodel.get_mosaic_position(img_origin, self._netM, opt)
|
||||
if size <= 50:
|
||||
emit(i, img_origin.copy()) # no mosaic here; recurrence carries over
|
||||
continue
|
||||
|
||||
stream = []
|
||||
for k in range(T):
|
||||
j = min(max(i + (k - N) * S, 0), count - 1) # clamp window to range edges
|
||||
frame = img_origin if j == i else get_frame(j)
|
||||
crop = frame[y - size:y + size, x - size:x + size]
|
||||
stream.append(impro.resize(crop, SZ)[:, :, ::-1]) # BGR→RGB, SZ×SZ
|
||||
|
||||
if previous is None: # seed recurrence with the (centre) input crop
|
||||
previous = data.im2tensor(stream[N], bgr2rgb=False, gpu_id=opt.gpu_id)
|
||||
|
||||
arr = np.array(stream).reshape(1, T, SZ, SZ, 3).transpose((0, 4, 1, 2, 3))
|
||||
tensor = data.to_tensor(data.normalize(arr), gpu_id=opt.gpu_id)
|
||||
with torch.no_grad():
|
||||
pred = self._netG(tensor, previous)
|
||||
previous = pred
|
||||
img_fake = data.tensor2im(pred, rgb2bgr=True)
|
||||
emit(i, impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather))
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
"""Restorer factory: build a Restorer from the app config.
|
||||
|
||||
- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct).
|
||||
- ``deepmosaics``: real generative mosaic removal. The DeepMosaics network code is
|
||||
vendored (``_deepmosaics/``, GPL-3.0) and run in-process; the user supplies only the
|
||||
clean weights (+ ``mosaic_position.pth`` alongside). A CUDA GPU is recommended.
|
||||
Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but
|
||||
did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``,
|
||||
GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth``
|
||||
alongside). A CUDA GPU is recommended.
|
||||
|
||||
- ``deepmosaics``: per-frame generative mosaic removal (image model).
|
||||
- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for
|
||||
coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,25 +15,28 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .base import Restorer
|
||||
from .inpaint import InpaintRestorer
|
||||
|
||||
if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed)
|
||||
from ...config import AppConfig
|
||||
|
||||
|
||||
def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer:
|
||||
if name == "inpaint":
|
||||
return InpaintRestorer()
|
||||
def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer:
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
if name == "deepmosaics":
|
||||
from .deepmosaics import DeepMosaicsRestorer
|
||||
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
return DeepMosaicsRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "deepmosaics_video":
|
||||
from .deepmosaics import DeepMosaicsVideoRestorer
|
||||
|
||||
return DeepMosaicsVideoRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "lada":
|
||||
raise ValueError(
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README."
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics. См. README."
|
||||
)
|
||||
raise ValueError(f"Неизвестный режим восстановления: {name!r}")
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Classic inpainting restorer (cv2) — the always-available baseline.
|
||||
|
||||
HONEST LIMITATION: cv2 inpainting fills the masked region by propagating
|
||||
surrounding pixels. It removes the mosaic/bar but does NOT reconstruct the hidden
|
||||
detail — it smooths/guesses. For real reconstruction a generative model
|
||||
(DeepMosaics / LADA) is needed; this is the no-weights, no-GPU fallback so the
|
||||
"Расцензурить кадр" flow works end-to-end today.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Restorer
|
||||
from .mask import detections_to_mask
|
||||
|
||||
|
||||
class InpaintRestorer(Restorer):
|
||||
def __init__(self, radius: int = 3, dilate: int = 2, method: str = "telea") -> None:
|
||||
self.radius = radius
|
||||
self.dilate = dilate
|
||||
self.method = method
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"InpaintRestorer({self.method})"
|
||||
|
||||
def restore(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
detections: list[Detection],
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> np.ndarray:
|
||||
# Single cv2.inpaint call — effectively instant, so cancellation is moot.
|
||||
if not detections:
|
||||
return image.copy()
|
||||
mask = detections_to_mask(image.shape, detections, dilate=self.dilate)
|
||||
flags = cv2.INPAINT_TELEA if self.method == "telea" else cv2.INPAINT_NS
|
||||
return cv2.inpaint(image, mask, self.radius, flags)
|
||||
@@ -1,26 +0,0 @@
|
||||
"""Build a binary mask of the censored regions from detections."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
|
||||
|
||||
def detections_to_mask(
|
||||
shape: tuple[int, int], detections: list[Detection], dilate: int = 0
|
||||
) -> np.ndarray:
|
||||
"""White (255) over every detected region (polygon if present, else bbox)."""
|
||||
h, w = shape[:2]
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
for d in detections:
|
||||
if len(d.polygon) >= 3:
|
||||
cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255)
|
||||
else:
|
||||
x, y, bw, bh = d.bbox
|
||||
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1)
|
||||
if dilate > 0:
|
||||
k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8)
|
||||
mask = cv2.dilate(mask, k)
|
||||
return mask
|
||||
Reference in New Issue
Block a user