Introduce diffusion-inpaint restoration engine in HVideoTool: added support for a new restoration method that regenerates masked regions via an external SwarmUI server, requiring YOLO detections for mask creation. Updated configuration management to include diffusion parameters, enhanced the UI for engine selection, and improved documentation in README and CLAUDE.md to guide users on the new functionality.
This commit is contained in:
@@ -67,7 +67,7 @@ from ..core.detection.factory import build_detector
|
||||
from ..core.detection.types import Detection
|
||||
from ..core.imageio import imread_unicode, imwrite_unicode
|
||||
from ..core.project import PROJECT_FILE, Project
|
||||
from ..core.restore.factory import build_restorer
|
||||
from ..core.restore.factory import build_restorer, restorer_needs_detections
|
||||
from ..core.video.extract import extract_frames
|
||||
from ..core.video.frame import Frame
|
||||
from .extract_dialog import ExtractDialog
|
||||
@@ -1220,13 +1220,24 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
path = self._current
|
||||
key = str(path)
|
||||
# The diffusion engine builds its mask from detections — capture them on the GUI
|
||||
# thread (DeepMosaics ignores them). Warn if it needs them but none are computed.
|
||||
needs_dets = restorer_needs_detections(self._cfg.restorer)
|
||||
dets_for_restore = list(self._results.get(key, []))
|
||||
if needs_dets and not dets_for_restore:
|
||||
self.statusBar().showMessage(
|
||||
"Diffusion перерисовывает по детекции — на этом кадре цензура не найдена "
|
||||
"(сначала «Рассчитать кадр»)"
|
||||
)
|
||||
return
|
||||
|
||||
def fn(job):
|
||||
img = imread_unicode(key)
|
||||
if img is None:
|
||||
raise RuntimeError(f"Не удалось прочитать: {path.name}")
|
||||
restorer = self._make_restorer()
|
||||
restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled)
|
||||
dets = dets_for_restore if restorer.needs_detections else []
|
||||
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
|
||||
return ("restored", key, restored, restorer.name)
|
||||
|
||||
def done(result, cancelled):
|
||||
@@ -1275,9 +1286,17 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
files = list(self._files) # snapshot — favorites/move mutate self._files
|
||||
is_temporal = self._cfg.restorer == "deepmosaics_video"
|
||||
# The diffusion engine masks the detections, so it inherently runs only on hits and
|
||||
# requires detection to be computed — same gating as the explicit "найденное" mode.
|
||||
needs_dets = restorer_needs_detections(self._cfg.restorer)
|
||||
hits = [i for i, p in enumerate(files) if self._results.get(str(p))]
|
||||
# Per-index detection snapshot (GUI-thread read) — fed to engines that need a mask.
|
||||
dets_by_index = (
|
||||
{i: list(self._results.get(str(files[i]), [])) for i in hits} if needs_dets else {}
|
||||
)
|
||||
need_hits = only_detected or needs_dets
|
||||
|
||||
if only_detected:
|
||||
if need_hits:
|
||||
if not self._results:
|
||||
self.statusBar().showMessage(
|
||||
"Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)"
|
||||
@@ -1331,7 +1350,8 @@ class MainWindow(QMainWindow):
|
||||
should_cancel=lambda: job.cancelled,
|
||||
)
|
||||
else:
|
||||
indices = hits if only_detected else range(len(files))
|
||||
indices = hits if need_hits else range(len(files))
|
||||
use_dets = restorer.needs_detections
|
||||
for i in indices:
|
||||
if job.cancelled:
|
||||
break
|
||||
@@ -1339,7 +1359,8 @@ class MainWindow(QMainWindow):
|
||||
if not force and out_path(p).is_file():
|
||||
emit(i, None, verb="Пропуск") # already restored — count, don't rewrite
|
||||
continue
|
||||
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
||||
dets = dets_by_index.get(i, []) if use_dets else []
|
||||
emit(i, restorer.restore(get_frame(i), dets, should_cancel=lambda: job.cancelled))
|
||||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||||
return None
|
||||
|
||||
@@ -1359,7 +1380,14 @@ class MainWindow(QMainWindow):
|
||||
self._start_job(fn, total, on_done=done)
|
||||
|
||||
def _make_restorer(self):
|
||||
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, self._cfg.dm_gpu)
|
||||
c = self._cfg
|
||||
key = (
|
||||
c.restorer, c.dm_dir, c.dm_model, c.dm_gpu,
|
||||
# diffusion identity — changing any of these must rebuild the engine
|
||||
c.diff_backend, c.diff_url, c.diff_model, c.diff_prompt, c.diff_negative,
|
||||
c.diff_steps, c.diff_cfg, c.diff_denoise, c.diff_seed,
|
||||
c.diff_mask_dilate, c.diff_mask_blur,
|
||||
)
|
||||
if key != self._restorer_key:
|
||||
self._restorer = build_restorer(self._cfg.restorer, self._cfg) # may raise
|
||||
self._restorer_key = key
|
||||
|
||||
Reference in New Issue
Block a user