Enhance HVideoTool's detection and restoration features: added support for tracking the model used in detections, improved temporal coherence by allowing the use of already-restored frames in the restoration process, and updated the UI to reflect these changes with new indicators and configuration options. Documentation in CLAUDE.md has been updated accordingly.

CLAUDE.md

@@ -149,7 +149,9 @@ hvideotool/
   paths and, if nothing is selected, default-ticks every discovered model. `build_detector`
   builds one `YoloDetector` per selected model (tagged `label=category`) wrapped in a
   `MultiYoloDetector` that concatenates their detections (no cross-model dedup). Each
-  `Detection` carries `label` (category); overlay colour + table group by `Detection.display`
+  `Detection` carries `label` (category) **and `model`** (the producing `.pt` stem, tagged in
+  `YoloDetector` — shown as its own "Модель" column in the detail table since a category folder
+  may hold several models); overlay colour + table group by `Detection.display`
   (label, else the CensorType) via `OverlayConfig.colors` + a stable `palette` fallback.
 - **Background jobs (`ui/workers.py`).** Detection and restoration are CPU-heavy and
   would freeze the GUI, so they run on a `QThreadPool` thread via `Job` (a `QRunnable`
@@ -248,6 +250,12 @@ hvideotool/
     `restore_sequence(count, get_frame, get_dets, emit, should_cancel)` (the batch run uses
     it; single-frame `restore` degrades to a window of the same frame). Needs the **video**
     weights `clean_youknow_video.pth` (+ `mosaic_position.pth` beside). `INPUT_SIZE=256`.
+    **`feed_restored`** (config `dm_feed_restored`, default on, checkbox in `RestoreDialog`):
+    when set, the **past** neighbours in the window (`j<i`) are taken from the engine's own
+    already-restored outputs (a small rolling cache, `reach=N*S` deep) instead of the original
+    mosaic frames — stronger temporal coherence. The **centre** (the frame being cleaned) and
+    **future** neighbours (`j>i`, not yet restored) stay original. Slightly out-of-distribution
+    for the BVDNet (trained on mosaic windows), so it's a toggle; off = faithful DeepMosaics.
 
   `restore_sequence` is on the `Restorer` ABC (default = independent per-frame loop);
   `_restore_all` dispatches on `restorer.temporal` (temporal → `restore_sequence` over the
@@ -266,10 +274,15 @@ hvideotool/
   `_step_hit` scans `_results` for the next non-empty frame). The slider and file list
   are kept in sync via `_update_nav` guarded by `_nav_sync` (avoids signal loops); all
   navigation ultimately drives `file_list.setCurrentRow`. The scrubber is a custom
-  `MarkerSlider` (`ui/marker_slider.py`) that paints cyan ticks at frames with
-  detections (`_refresh_marks` projects `_results` onto row indices; per-pixel deduped
-  so big folders stay cheap). File-list rows are tinted too (`_tag_file`): red =
-  censorship found, green = checked & clean. Both reset on `_invalidate_results`.
+  `MarkerSlider` (`ui/marker_slider.py`) that paints **two mark layers**: cyan ticks
+  (upper half) at frames with detections (`_refresh_marks` projects `_results`) and
+  **green ticks (lower half) at restored frames** (`_refresh_restored_marks` scans
+  `restored/` + in-memory `_restored`; called on load and after each restore op, not
+  per-frame); per-pixel deduped so big folders stay cheap. Under the scrubber a
+  **progress summary** `stats_label` reads "Кадров: N · детектировано: D/N (с цензурой:
+  H) · расцензурено: R/N" (`_update_counts_label`, cheap counts; `_restored_count`
+  cached by `_refresh_restored_marks`). File-list rows are tinted too (`_tag_file`):
+  red = censorship found, green = checked & clean. Both reset on `_invalidate_results`.
 - **Cancellation (cooperative).** A single "■ Стоп" toolbar action (Esc) cancels the
   running op. `_begin_busy(total)` / `_end_busy()` toggle `self._busy` + the Stop button +
   the progress bar (`total=None` → indeterminate). For **background jobs** (detection,

hvideotool/config.py

@@ -62,3 +62,7 @@ class AppConfig:
     dm_dir: str | None = None          # optional extra dir to search for mosaic_position.pth
     dm_model: str | None = None        # DeepMosaics clean weights (clean_*.pth)
     dm_gpu: str = "0"                  # CUDA device id, "-1" for CPU
+    # Temporal engine only: feed already-restored PAST frames into the BVDNet window
+    # (instead of the original mosaic frames) for stronger temporal coherence. Slightly
+    # out-of-distribution for the net (trained on mosaic windows) — toggle in the dialog.
+    dm_feed_restored: bool = True

hvideotool/core/detection/types.py

@@ -24,6 +24,7 @@ class Detection:
     bbox: tuple[int, int, int, int]           # x, y, w, h
     polygon: list[tuple[int, int]] = field(default_factory=list)  # contour points
     label: str = ""                           # model category (models/yolo/<label>); drives colour/grouping
+    model: str = ""                           # weights file that produced it (the .pt stem)
 
     @property
     def display(self) -> str:
@@ -37,6 +38,7 @@ class Detection:
             "bbox": list(self.bbox),
             "polygon": [list(p) for p in self.polygon],
             "label": self.label,
+            "model": self.model,
         }
 
     @classmethod
@@ -47,4 +49,5 @@ class Detection:
             bbox=tuple(data["bbox"]),  # type: ignore[arg-type]
             polygon=[tuple(p) for p in data.get("polygon", [])],
             label=data.get("label", ""),
+            model=data.get("model", ""),
         )

hvideotool/core/detection/yolo.py

@@ -35,10 +35,14 @@ def _name_to_type(name: str) -> CensorType:
 
 class YoloDetector(Detector):
     def __init__(
-        self, model_path: str, config: DetectionConfig | None = None, label: str = ""
+        self, model_path: str, config: DetectionConfig | None = None,
+        label: str = "", model_name: str = "",
     ) -> None:
         self.cfg = config or DetectionConfig()
         self._label = label  # category (models/yolo/<label>) tagged onto every detection
+        # weights filename (stem) tagged onto every detection, so the UI can show
+        # which specific model predicted it (a category folder may hold several).
+        self._model_name = model_name or os.path.splitext(os.path.basename(model_path))[0]
         if not os.path.isfile(model_path):
             raise FileNotFoundError(
                 f"Файл весов не найден: {model_path}\n"
@@ -108,6 +112,7 @@ class YoloDetector(Detector):
                 poly = [(int(px), int(py)) for px, py in polygons[i]]
             ctype = _name_to_type(names.get(int(classes[i]), ""))
             out.append(Detection(
-                type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly, label=self._label
+                type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly,
+                label=self._label, model=self._model_name,
             ))
         return out

hvideotool/core/project.py

@@ -42,6 +42,7 @@ _SETTING_KEYS = (
     "dm_dir",
     "dm_model",
     "dm_gpu",
+    "dm_feed_restored",
 )
 
 

hvideotool/core/restore/deepmosaics.py

@@ -209,7 +209,12 @@ class DeepMosaicsVideoRestorer(Restorer):
         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)
@@ -299,6 +304,18 @@ class DeepMosaicsVideoRestorer(Restorer):
         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):
@@ -307,13 +324,20 @@ class DeepMosaicsVideoRestorer(Restorer):
             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
+                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
-                frame = img_origin if j == i else get_frame(j)
+                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
 
@@ -326,4 +350,6 @@ class DeepMosaicsVideoRestorer(Restorer):
                 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))
+            result = impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather)
+            emit(i, result)
+            remember(i, result)

hvideotool/core/restore/factory.py

@@ -33,7 +33,8 @@ def build_restorer(name: str = "deepmosaics", config: AppConfig | None = None) -
         from .deepmosaics import DeepMosaicsVideoRestorer
 
         return DeepMosaicsVideoRestorer(
-            config.dm_dir, config.dm_model, config.dm_gpu
+            config.dm_dir, config.dm_model, config.dm_gpu,
+            feed_restored=getattr(config, "dm_feed_restored", True),
         )
     if name == "lada":
         raise ValueError(

hvideotool/settings_store.py

@@ -42,6 +42,8 @@ def apply(config: AppConfig) -> None:
     for key in ("dm_dir", "dm_model", "dm_gpu"):
         if key in data:
             setattr(config, key, data[key])
+    if "dm_feed_restored" in data:
+        config.dm_feed_restored = bool(data["dm_feed_restored"])
 
 
 def save(config: AppConfig) -> None:
@@ -55,6 +57,7 @@ def save(config: AppConfig) -> None:
         dm_dir=config.dm_dir,
         dm_model=config.dm_model,
         dm_gpu=config.dm_gpu,
+        dm_feed_restored=config.dm_feed_restored,
     )
     _write(data)
 

hvideotool/ui/main_window.py

@@ -88,6 +88,7 @@ class MainWindow(QMainWindow):
         self._restorer = None           # un-censor engine, built lazily from config
         self._restorer_key = None
         self._restored: dict[str, object] = {}    # path -> restored image (BGR ndarray)
+        self._restored_count = 0  # frames with output in restored/ (for the progress summary)
         self._showing_restored = False
         self._nav_sync = False  # guard against slider<->list signal loops
         self._busy = False      # a long operation is running
@@ -218,6 +219,11 @@ class MainWindow(QMainWindow):
         clayout.setSpacing(2)
         clayout.addWidget(self.view, 1)
         clayout.addWidget(self._build_nav_bar())
+        # Processing summary under the scrubber: how much of the sequence is done.
+        self.stats_label = QLabel("")
+        self.stats_label.setAlignment(Qt.AlignCenter)
+        self.stats_label.setStyleSheet("QLabel{color:#888; padding:1px;}")
+        clayout.addWidget(self.stats_label)
 
         right = QWidget()
         rlayout = QVBoxLayout(right)
@@ -225,8 +231,10 @@ class MainWindow(QMainWindow):
         self.detail_header = QLabel("Детекции")
         self.detail_header.setWordWrap(True)
         rlayout.addWidget(self.detail_header)
-        self.detail_table = QTableWidget(0, 4)
-        self.detail_table.setHorizontalHeaderLabels(["Категория", "Увер.", "BBox (x,y,w,h)", "Полигон"])
+        self.detail_table = QTableWidget(0, 5)
+        self.detail_table.setHorizontalHeaderLabels(
+            ["Категория", "Модель", "Увер.", "BBox (x,y,w,h)", "Полигон"]
+        )
         self.detail_table.verticalHeader().setVisible(False)
         self.detail_table.setSelectionBehavior(QTableWidget.SelectRows)
         self.detail_table.setEditTriggers(QTableWidget.NoEditTriggers)
@@ -864,6 +872,7 @@ class MainWindow(QMainWindow):
         self._end_busy()
         loaded = self._load_cached_results()  # reuse a matching on-disk cache
         self._refresh_marks()
+        self._refresh_restored_marks()
 
         if not files:
             self.view.set_image(None, [])
@@ -1031,6 +1040,7 @@ class MainWindow(QMainWindow):
                 self._showing_restored = True
                 self.view.set_image(restored, [])
                 self._update_restore_actions()
+            self._refresh_restored_marks()
             self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}")
 
         self.statusBar().showMessage(f"Восстановление: {path.name}…")
@@ -1098,6 +1108,7 @@ class MainWindow(QMainWindow):
                         self._showing_restored = True
                         self.view.set_image(img, [])
                         self._update_restore_actions()
+            self._refresh_restored_marks()
             self.statusBar().showMessage(
                 "Расцензуривание отменено" if cancelled
                 else f"Готово: результаты в {out_dir.name}/ ({total} кадров)"
@@ -1196,6 +1207,7 @@ class MainWindow(QMainWindow):
         elif self.file_list.count() == 0:
             self.view.set_image(None, [])
         self._refresh_marks()  # rows shifted — remap marks to new indices
+        self._refresh_restored_marks()
         self._update_nav()
 
     @staticmethod
@@ -1276,6 +1288,39 @@ class MainWindow(QMainWindow):
             if self._results.get(self.file_list.item(i).data(Qt.UserRole))
         }
         self.frame_slider.set_marks(marks)
+        self._update_counts_label()
+
+    def _update_counts_label(self) -> None:
+        """Refresh the processing summary under the scrubber (cheap; counts only)."""
+        n = len(self._files)
+        detected = sum(1 for p in self._files if str(p) in self._results)
+        hits = sum(1 for p in self._files if self._results.get(str(p)))
+        if n == 0:
+            self.stats_label.setText("")
+            return
+        self.stats_label.setText(
+            f"Кадров: {n}  ·  детектировано: {detected}/{n} (с цензурой: {hits})"
+            f"  ·  расцензурено: {self._restored_count}/{n}"
+        )
+
+    def _refresh_restored_marks(self) -> None:
+        """Scan ``restored/`` and project restored frames onto the scrubber (green).
+
+        Done only on load / after a restore op (not per-frame) — it touches the disk.
+        """
+        stems: set[str] = set()
+        if self._project is not None and self._project.restored_dir.is_dir():
+            stems = {p.stem for p in self._project.restored_dir.glob("*.jpg")}
+        mem = set(self._restored)  # single-frame restores held in memory (not on disk yet)
+        rows: set[int] = set()
+        if stems or mem:
+            for i in range(self.file_list.count()):
+                fp = Path(self.file_list.item(i).data(Qt.UserRole))
+                if fp.stem in stems or str(fp) in mem:
+                    rows.add(i)
+        self._restored_count = len(rows)
+        self.frame_slider.set_restored_marks(rows)
+        self._update_counts_label()
 
     def _fill_detail_table(self, path: Path, img, dets: list[Detection] | None) -> None:
         h, w = (img.shape[0], img.shape[1]) if img is not None else (0, 0)
@@ -1298,7 +1343,7 @@ class MainWindow(QMainWindow):
         self.detail_table.setRowCount(len(dets))
         for row, d in enumerate(dets):
             x, y, bw, bh = d.bbox
-            cells = [d.display, f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
+            cells = [d.display, d.model or "—", f"{d.score:.2f}", f"{x},{y},{bw},{bh}", str(len(d.polygon))]
             for col, text in enumerate(cells):
                 self.detail_table.setItem(row, col, QTableWidgetItem(text))
         self.detail_table.blockSignals(False)

hvideotool/ui/marker_slider.py

@@ -1,8 +1,11 @@
-"""A horizontal QSlider that paints marks at frames where censorship was found.
+"""A horizontal QSlider that paints marks at frames along the sequence.
 
-Used as the navigation scrubber under the image: the file list / `_results`
-cache is projected onto the groove as small vertical ticks, so you can see at a
-glance where the detected regions are along the whole sequence.
+Used as the navigation scrubber under the image. Two independent layers are
+projected onto the groove as small vertical ticks, so you can see at a glance how
+the whole sequence is processed:
+
+* **cyan** (upper half) — frames where censorship was *detected* (`_results`);
+* **green** (lower half) — frames that have been *restored* (``restored/``).
 """
 
 from __future__ import annotations
@@ -17,19 +20,29 @@ from PySide6.QtWidgets import QSlider, QStyle, QStyleOptionSlider
 class MarkerSlider(QSlider):
     def __init__(self, orientation=Qt.Horizontal, parent=None) -> None:
         super().__init__(orientation, parent)
-        self._marks: set[int] = set()
+        self._marks: set[int] = set()           # detected (censorship found)
+        self._restored_marks: set[int] = set()  # restored (un-censored)
         # Bright cyan stands out against both the dark track and the orange fill.
         self._mark_color = QColor(0, 220, 255)
+        self._restored_color = QColor(80, 230, 120)  # green = restored
 
     def set_marks(self, marks: Iterable[int]) -> None:
+        """Frames with detections (painted cyan, upper half)."""
         marks = set(marks)
         if marks != self._marks:
             self._marks = marks
             self.update()
 
+    def set_restored_marks(self, marks: Iterable[int]) -> None:
+        """Frames that have been restored (painted green, lower half)."""
+        marks = set(marks)
+        if marks != self._restored_marks:
+            self._restored_marks = marks
+            self.update()
+
     def paintEvent(self, event) -> None:
         super().paintEvent(event)
-        if not self._marks or self.maximum() <= self.minimum():
+        if (not self._marks and not self._restored_marks) or self.maximum() <= self.minimum():
             return
 
         opt = QStyleOptionSlider()
@@ -45,22 +58,32 @@ class MarkerSlider(QSlider):
             return
         lo, hi = self.minimum(), self.maximum()
         half = handle.width() // 2
-        # Span (almost) the full widget height so marks read over the fill/handle.
         top = self.rect().top() + 1
         bottom = self.rect().bottom() - 1
+        mid = (top + bottom) // 2
 
         painter = QPainter(self)
+        # Detections in the upper half, restored in the lower half, so a frame that's
+        # both detected and restored shows both ticks instead of one hiding the other.
+        self._paint_layer(painter, self._marks, self._mark_color, top, mid,
+                          groove.x(), half, span, lo, hi, opt.upsideDown)
+        self._paint_layer(painter, self._restored_marks, self._restored_color, mid, bottom,
+                          groove.x(), half, span, lo, hi, opt.upsideDown)
+        painter.end()
+
+    def _paint_layer(self, painter, marks, color, y0, y1, gx, half, span, lo, hi, upside) -> None:
+        if not marks:
+            return
         pen = painter.pen()
-        pen.setColor(self._mark_color)
+        pen.setColor(color)
         pen.setWidth(2)
         painter.setPen(pen)
         # Many frames can collapse onto the same pixel column — dedupe to keep
         # repaint cheap on big folders (tens of thousands of frames).
         seen_x: set[int] = set()
-        for m in self._marks:
-            pos = QStyle.sliderPositionFromValue(lo, hi, m, span, opt.upsideDown)
-            x = groove.x() + half + pos
+        for m in marks:
+            pos = QStyle.sliderPositionFromValue(lo, hi, m, span, upside)
+            x = gx + half + pos
             if x not in seen_x:
                 seen_x.add(x)
-                painter.drawLine(x, top, x, bottom)
-        painter.end()
+                painter.drawLine(x, y0, x, y1)

hvideotool/ui/restore_dialog.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 from pathlib import Path
 
 from PySide6.QtWidgets import (
+    QCheckBox,
     QComboBox,
     QDialog,
     QDialogButtonBox,
@@ -49,10 +50,23 @@ class RestoreDialog(QDialog):
         self.dm_gpu = QLineEdit(config.dm_gpu or "0")
         self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)")
 
+        # Video engine only: feed already-restored past frames into the temporal window.
+        self.feed_restored = QCheckBox(
+            "Подавать уже расцензуренные прошлые кадры в окно (эксперим.)"
+        )
+        self.feed_restored.setChecked(bool(getattr(config, "dm_feed_restored", True)))
+        self.feed_restored.setToolTip(
+            "Только для видеодвижка: прошлые соседние кадры в окне берутся из уже\n"
+            "восстановленных результатов, а не из оригинала с мозаикой — больше\n"
+            "временной связности. Сеть обучалась на мозаичных окнах, так что эффект\n"
+            "не гарантирован; выключите для точной реализации DeepMosaics."
+        )
+
         form = QFormLayout(self)
         form.addRow("Движок:", self.engine)
         form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model))
         form.addRow("GPU id:", self.dm_gpu)
+        form.addRow("", self.feed_restored)
         hint = QLabel(
             "Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n"
             "• Картинка: clean_youknow_resnet_9blocks.pth — покадрово.\n"
@@ -98,10 +112,12 @@ class RestoreDialog(QDialog):
 
     def _sync(self) -> None:
         is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video")
+        is_video = self.engine.currentData() == "deepmosaics_video"
         # The model list differs per engine (image vs video weights) — repopulate.
         self._populate_models(self._cfg.dm_model)
         self.model_combo.setEnabled(is_dm)
         self.dm_gpu.setEnabled(is_dm)
+        self.feed_restored.setEnabled(is_video)  # only the temporal engine has a window
 
     def _browse_model(self) -> None:
         p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)")
@@ -114,3 +130,4 @@ class RestoreDialog(QDialog):
         self._cfg.restorer = self.engine.currentData()
         self._cfg.dm_model = self.model_combo.currentData()
         self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0"
+        self._cfg.dm_feed_restored = self.feed_restored.isChecked()