108 lines
4.6 KiB
Python
108 lines
4.6 KiB
Python
"""DeepMosaics restorer — real generative mosaic removal.
|
||
|
||
Rather than vendoring DeepMosaics' GPL network code (which must match the exact
|
||
checkpoint), we drive a **user-installed** DeepMosaics (https://github.com/HypoX64/DeepMosaics)
|
||
as a subprocess: write the frame to a temp file, run ``deepmosaic.py --mode clean``,
|
||
read the cleaned image back. This reuses their tested pipeline (incl. their own
|
||
mosaic locator ``mosaic_position.pth``) and respects the GPL boundary.
|
||
|
||
Setup the user must do once (see README → Восстановление):
|
||
1. ``git clone https://github.com/HypoX64/DeepMosaics`` and install its deps.
|
||
2. Download clean weights (e.g. ``clean_youknow_video.pth``) AND ``mosaic_position.pth``
|
||
into one folder.
|
||
3. In the app: Восстановление… → engine "deepmosaics", set the DeepMosaics folder
|
||
and the clean-model path (a CUDA GPU is strongly recommended).
|
||
|
||
NOTE: DeepMosaics finds the mosaic itself; our detections are used for navigation,
|
||
not passed to it.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
|
||
from ..detection.types import Detection
|
||
from ..imageio import imread_unicode, imwrite_unicode
|
||
from .base import Restorer
|
||
|
||
_IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp"}
|
||
|
||
|
||
class DeepMosaicsRestorer(Restorer):
|
||
def __init__(
|
||
self,
|
||
deepmosaics_dir: str | None,
|
||
model_path: str | None,
|
||
python_exe: str | None = None,
|
||
gpu_id: str = "0",
|
||
) -> None:
|
||
if not deepmosaics_dir or not (Path(deepmosaics_dir) / "deepmosaic.py").is_file():
|
||
raise ValueError(
|
||
"Не указана папка DeepMosaics (с deepmosaic.py).\n"
|
||
"Установите DeepMosaics и укажите её в «Восстановление…». См. README."
|
||
)
|
||
if not model_path or not Path(model_path).is_file():
|
||
raise ValueError(
|
||
"Не найдены веса DeepMosaics (clean_*.pth).\n"
|
||
"Скачайте clean_youknow_video.pth + mosaic_position.pth в одну папку. См. README."
|
||
)
|
||
self._dir = Path(deepmosaics_dir)
|
||
self._model = model_path
|
||
self._python = python_exe or sys.executable
|
||
self._gpu = gpu_id
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return f"DeepMosaics(gpu={self._gpu})"
|
||
|
||
def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray:
|
||
with tempfile.TemporaryDirectory(prefix="hvt_dm_") as tmp:
|
||
tmpd = Path(tmp)
|
||
src = tmpd / "frame.jpg"
|
||
result_dir = tmpd / "result"
|
||
result_dir.mkdir()
|
||
imwrite_unicode(str(src), image)
|
||
|
||
cmd = [
|
||
self._python, "deepmosaic.py",
|
||
"--media_path", str(src),
|
||
"--model_path", str(self._model),
|
||
"--mode", "clean",
|
||
"--result_dir", str(result_dir),
|
||
"--temp_dir", str(tmpd / "dmtmp"),
|
||
"--gpu_id", str(self._gpu),
|
||
"--no_preview",
|
||
]
|
||
proc = subprocess.run(
|
||
cmd, cwd=str(self._dir),
|
||
stdin=subprocess.DEVNULL, # so DeepMosaics' error input() can't hang
|
||
capture_output=True, text=True,
|
||
)
|
||
outputs = [p for p in result_dir.iterdir() if p.suffix.lower() in _IMG_EXTS]
|
||
if outputs:
|
||
newest = max(outputs, key=lambda p: p.stat().st_mtime)
|
||
restored = imread_unicode(str(newest))
|
||
if restored is None:
|
||
raise RuntimeError("Не удалось прочитать результат DeepMosaics.")
|
||
return restored
|
||
|
||
# No output file — figure out why.
|
||
log = (proc.stderr or "") + (proc.stdout or "")
|
||
if "BVDNet.forward()" in log or "argument: 'previous'" in log:
|
||
raise RuntimeError(
|
||
"Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен "
|
||
"соседний кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
|
||
)
|
||
if proc.returncode == 0:
|
||
# DeepMosaics ran fine but found no mosaic to clean — keep the frame as is.
|
||
return image.copy()
|
||
tail = log.strip().splitlines()[-6:]
|
||
raise RuntimeError(
|
||
f"DeepMosaics не вернул результат (код {proc.returncode}).\n" + "\n".join(tail)
|
||
)
|