356 lines
14 KiB
Python
356 lines
14 KiB
Python
"""DeepMosaics restorers — real generative mosaic removal, in-process.
|
||
|
||
The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see its
|
||
NOTICE/LICENSE). We load the models **once** and run in-process — far faster than
|
||
spawning a subprocess per frame (which reloaded the models every time). Only the model
|
||
*weights* are user-supplied. Both engines locate the mosaic themselves (BiSeNet
|
||
``mosaic_position.pth``); detections are not passed to them.
|
||
|
||
Two engines:
|
||
- :class:`DeepMosaicsRestorer` (per-frame): reproduces ``cleanmosaic_img_server`` —
|
||
locate mosaic → run the image generator on the crop → feather it back. Image weights
|
||
``clean_youknow_resnet_9blocks.pth``.
|
||
- :class:`DeepMosaicsVideoRestorer` (temporal/BVDNet): reproduces
|
||
``cleanmosaic_video_fusion`` — a window of neighbouring frames + recurrence, for
|
||
temporal coherence. Video weights ``clean_youknow_video.pth``; needs a contiguous
|
||
sequence (see :meth:`Restorer.restore_sequence`).
|
||
|
||
Setup (see README → Восстановление): drop the chosen ``clean_*.pth`` + ``mosaic_position.pth``
|
||
into one folder (``models/deepmosaics``) and pick it in the restore dialog.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
import numpy as np
|
||
|
||
from ..detection.types import Detection
|
||
from .base import (
|
||
CancelCheck,
|
||
Cancelled,
|
||
DetGetter,
|
||
FrameGetter,
|
||
Restorer,
|
||
ResultSink,
|
||
)
|
||
|
||
_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, 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``. 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:
|
||
dirs.insert(0, Path(extra_dir))
|
||
out: list[tuple[str, str]] = []
|
||
seen: set[str] = set()
|
||
for d in dirs:
|
||
if not d.is_dir():
|
||
continue
|
||
for p in sorted(d.glob("clean_*.pth")):
|
||
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()
|
||
if "video" in n:
|
||
raise ValueError(
|
||
"Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен соседний "
|
||
"кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
|
||
)
|
||
if "unet_128" in n:
|
||
return "unet_128"
|
||
if "hd" in n:
|
||
return "HD"
|
||
return "resnet_9blocks"
|
||
|
||
|
||
class DeepMosaicsRestorer(Restorer):
|
||
def __init__(
|
||
self,
|
||
deepmosaics_dir: str | None, # optional extra dir to find mosaic_position.pth
|
||
model_path: str | None,
|
||
gpu_id: str = "0",
|
||
) -> None:
|
||
if not model_path or not Path(model_path).is_file():
|
||
discovered = discover_models() # fall back to a bundled model
|
||
if discovered:
|
||
model_path = discovered[0][1]
|
||
else:
|
||
raise ValueError(
|
||
"Не найдены веса DeepMosaics (clean_*.pth).\n"
|
||
"Положите clean_youknow_resnet_9blocks.pth + mosaic_position.pth в "
|
||
"models/deepmosaics (или выберите в «Восстановление…»). См. README."
|
||
)
|
||
model = Path(model_path)
|
||
self._netg = _netg_kind(model.name) # raises on a video model
|
||
pos = _find_mosaic_position(model, deepmosaics_dir)
|
||
if pos is None:
|
||
raise ValueError(
|
||
"Рядом с clean-моделью не найден mosaic_position.pth.\n"
|
||
"Положите mosaic_position.pth в ту же папку, что и clean_*.pth. См. README."
|
||
)
|
||
self._model = str(model)
|
||
self._pos = str(pos)
|
||
self._gpu = gpu_id
|
||
self._loaded = False # models loaded lazily on first restore
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return f"DeepMosaics(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)) # so vendored `from models/util import …` resolve
|
||
|
||
# Fall back to CPU when CUDA isn't available: DeepMosaics calls `.cuda()`
|
||
# whenever gpu_id != "-1", which raises "Torch not compiled with CUDA enabled"
|
||
# on a CPU-only torch build.
|
||
import torch
|
||
if self._gpu != "-1" and not torch.cuda.is_available():
|
||
self._gpu = "-1"
|
||
|
||
import util.image_processing as impro # type: ignore
|
||
|
||
from models import loadmodel, runmodel # type: ignore
|
||
|
||
self._runmodel = runmodel
|
||
self._impro = impro
|
||
self._opt = SimpleNamespace(
|
||
gpu_id=self._gpu,
|
||
netG=self._netg,
|
||
model_path=self._model,
|
||
mosaic_position_model_path=self._pos,
|
||
mask_threshold=64,
|
||
all_mosaic_area=False,
|
||
ex_mult=1.5,
|
||
no_feather=False,
|
||
traditional=False,
|
||
)
|
||
self._netM = loadmodel.bisenet(self._opt, "mosaic")
|
||
self._netG = loadmodel.pix2pix(self._opt)
|
||
self._loaded = True
|
||
|
||
def restore(
|
||
self,
|
||
image: np.ndarray,
|
||
detections: list[Detection],
|
||
should_cancel: CancelCheck | None = None,
|
||
) -> np.ndarray:
|
||
if should_cancel is not None and should_cancel():
|
||
raise Cancelled("Восстановление отменено")
|
||
self._ensure_loaded()
|
||
rm, impro, opt = self._runmodel, self._impro, self._opt
|
||
|
||
# DeepMosaics' cleanmosaic_img_server, faithfully reproduced.
|
||
x, y, size, mask = rm.get_mosaic_position(image, self._netM, opt)
|
||
if size <= 100:
|
||
return image.copy() # no mosaic located — leave the frame untouched
|
||
if should_cancel is not None and should_cancel():
|
||
raise Cancelled("Восстановление отменено")
|
||
work = image.copy()
|
||
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,
|
||
gpu_id: str = "0",
|
||
feed_restored: bool = True,
|
||
) -> None:
|
||
# When True, already-restored PAST frames are fed into the temporal window
|
||
# (instead of the original mosaic frames). Future neighbours and the centre
|
||
# frame stay original — they aren't restored yet / are what we're cleaning.
|
||
self._feed_restored = feed_restored
|
||
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
|
||
if self._gpu != "-1" and not torch.cuda.is_available():
|
||
self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why)
|
||
|
||
import util.data as data # type: ignore
|
||
import util.image_processing as impro # type: ignore
|
||
|
||
from models import loadmodel, runmodel # type: ignore
|
||
|
||
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
|
||
reach = N * S # how far back/forward the window samples (frames)
|
||
|
||
# Rolling cache of already-restored frames, used as window neighbours when
|
||
# ``feed_restored`` is on. Only the last ``reach`` frames are ever needed.
|
||
restored: dict[int, np.ndarray] = {}
|
||
|
||
def remember(idx: int, frame: np.ndarray) -> None:
|
||
if not self._feed_restored:
|
||
return
|
||
restored[idx] = frame
|
||
for old in [k for k in restored if k < idx - reach]:
|
||
restored.pop(old, None)
|
||
|
||
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:
|
||
clean = img_origin.copy() # no mosaic here; recurrence carries over
|
||
emit(i, clean)
|
||
remember(i, clean)
|
||
continue
|
||
|
||
stream = []
|
||
for k in range(T):
|
||
j = min(max(i + (k - N) * S, 0), count - 1) # clamp window to range edges
|
||
if j == i:
|
||
frame = img_origin
|
||
elif self._feed_restored and j < i and j in restored:
|
||
frame = restored[j] # already-restored past neighbour
|
||
else:
|
||
frame = get_frame(j) # original (future neighbour / not yet cached)
|
||
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)
|
||
result = impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather)
|
||
emit(i, result)
|
||
remember(i, result)
|