170 lines
6.9 KiB
Python
170 lines
6.9 KiB
Python
"""DeepMosaics restorer — 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 the per-frame clean path
|
|
in-process — far faster than spawning a subprocess per frame (which reloaded the
|
|
models every time). Only the model *weights* are user-supplied.
|
|
|
|
Per-frame clean = DeepMosaics' ``cleanmosaic_img_server`` logic, reimplemented
|
|
here (so we don't pull in their video/ffmpeg modules):
|
|
locate mosaic (BiSeNet ``mosaic_position.pth``) → run the clean generator on the
|
|
crop → feather it back. DeepMosaics finds the mosaic itself; our detections are
|
|
used for navigation, not passed to it.
|
|
|
|
Setup (see README → Восстановление): download the **image** clean weights
|
|
``clean_youknow_resnet_9blocks.pth`` + ``mosaic_position.pth`` into one folder and
|
|
point the app at the clean-model file. The video model ``clean_youknow_video.pth``
|
|
(BVDNet) needs neighbour frames and does NOT work per-frame.
|
|
"""
|
|
|
|
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, 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).
|
|
|
|
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.
|
|
"""
|
|
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 "video" in p.name.lower() or p.name in seen:
|
|
continue
|
|
seen.add(p.name)
|
|
out.append((p.stem, str(p)))
|
|
return out
|
|
|
|
|
|
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, # kept for factory/config compatibility (weights hint)
|
|
model_path: str | None,
|
|
python_exe: str | None = None, # unused now (in-process)
|
|
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 = self._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
|
|
|
|
@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})"
|
|
|
|
# ------------------------------------------------------------------ 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 # noqa: E402
|
|
if self._gpu != "-1" and not torch.cuda.is_available():
|
|
self._gpu = "-1"
|
|
|
|
from models import loadmodel, runmodel # type: ignore # noqa: E402
|
|
import util.image_processing as impro # type: ignore # noqa: E402
|
|
|
|
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)
|