Refactor HVideoTool to exclusively use YOLO for detection and DeepMosaics for restoration: removed classic CV and composite detectors, updated configuration and UI accordingly. Enhanced documentation in README and CLAUDE.md to reflect these changes, including new batch processing capabilities and device diagnostics.
This commit is contained in:
Leonid Pershin
@@ -12,7 +12,7 @@ import sys
|
||||
|
||||
from . import settings_store
|
||||
from .app import run
|
||||
from .config import AppConfig
|
||||
from .config import AppConfig, normalize_config
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -21,15 +21,13 @@ def main() -> int:
|
||||
description="Инспектор детекции уже наложенной цензуры на картинках.",
|
||||
)
|
||||
parser.add_argument("target", nargs="?", help="путь к проекту для открытия (папка или project.json)")
|
||||
parser.add_argument("--detector", choices=["classic", "yolo", "combined"], default=None)
|
||||
parser.add_argument("--model", dest="model_path", default=None, help="путь к весам (YOLO)")
|
||||
args = parser.parse_args()
|
||||
|
||||
config = AppConfig()
|
||||
settings_store.apply(config) # persisted defaults first
|
||||
normalize_config(config) # drop any legacy classic/inpaint values
|
||||
|
||||
if args.detector:
|
||||
config.detector = args.detector
|
||||
if args.model_path:
|
||||
config.model_path = args.model_path
|
||||
|
||||
|
||||
+21
-27
@@ -1,40 +1,21 @@
|
||||
"""Application configuration and tunable defaults.
|
||||
|
||||
Plain dataclasses. The detection thresholds matter most here — this tool is now
|
||||
an image-folder inspector for tuning the detectors, so keep them easy to tweak.
|
||||
Plain dataclasses. Detection is YOLO-only and restoration is DeepMosaics-only, so the
|
||||
knobs here are the YOLO inference params, the overlay style, and the DeepMosaics weights.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
DETECTORS = ("yolo",)
|
||||
RESTORERS = ("deepmosaics", "deepmosaics_video")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DetectionConfig:
|
||||
"""Parameters for the detector. Thresholds tuned for the classic-CV detector."""
|
||||
"""Parameters for the YOLO detector."""
|
||||
|
||||
proc_max_dim: int = 720 # downscale longer side to this before detection (speed)
|
||||
min_area_frac: float = 0.0008 # ignore regions smaller than this fraction of the image
|
||||
|
||||
# --- solid bars (black or white, achromatic, rectangular) ---
|
||||
black_intensity: int = 40 # V below this = dark-bar candidate
|
||||
white_intensity: int = 225 # V above this = light-bar candidate
|
||||
bar_saturation_max: int = 45 # S below this = achromatic (excludes colored fills)
|
||||
bar_min_extent: float = 0.80 # contour area / bbox area — how rectangular a bar must be
|
||||
|
||||
# --- mosaic / pixelation ---
|
||||
mosaic_block_sizes: tuple[int, ...] = (8, 12, 16, 24) # candidate tile sizes (px, proc space)
|
||||
mosaic_residual_max: float = 6.0 # max reconstruction error to count as "blocky"
|
||||
mosaic_contrast_min: float = 14.0 # min local contrast (excludes flat gradients)
|
||||
mosaic_grad_min: float = 8.0 # min edge energy in BOTH x and y (excludes straight edges)
|
||||
mosaic_min_side: int = 24 # reject thin regions (px) — kills edge false-positives
|
||||
|
||||
# --- blur ---
|
||||
blur_window: int = 31 # sliding window for local sharpness (odd)
|
||||
blur_sharpness_ratio: float = 0.35 # below this fraction of median sharpness => blurry
|
||||
blur_contrast_min: float = 8.0 # min local contrast (excludes flat regions)
|
||||
|
||||
# --- YOLO detector (used only when detector == "yolo"/"combined") ---
|
||||
yolo_conf: float = 0.2 # confidence threshold (LADA recommends ~0.2)
|
||||
yolo_imgsz: int = 640 # inference image size
|
||||
yolo_device: str | None = None # None => auto ("cuda" if available, else "cpu")
|
||||
@@ -62,13 +43,26 @@ class OverlayConfig:
|
||||
class AppConfig:
|
||||
detection: DetectionConfig = field(default_factory=DetectionConfig)
|
||||
overlay: OverlayConfig = field(default_factory=OverlayConfig)
|
||||
detector: str = "classic" # "classic" | "yolo" | "combined"
|
||||
detector: str = "yolo" # only "yolo"
|
||||
model_path: str | None = None # weights path, used by the YOLO detector
|
||||
default_threshold: float = 0.20 # initial overlay confidence threshold
|
||||
|
||||
# --- restoration ("расцензурить") ---
|
||||
restorer: str = "inpaint" # "inpaint" | "deepmosaics"
|
||||
restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
|
||||
dm_dir: str | None = None # DeepMosaics repo dir (contains deepmosaic.py)
|
||||
dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
|
||||
dm_python: str | None = None # python exe for DeepMosaics (None = current)
|
||||
dm_gpu: str = "0" # CUDA device id, "-1" for CPU
|
||||
|
||||
|
||||
def normalize_config(cfg: AppConfig) -> None:
|
||||
"""Coerce legacy/removed settings to supported values (mutates ``cfg``).
|
||||
|
||||
Old projects / settings.json may carry the removed ``classic``/``combined``
|
||||
detectors or the ``inpaint`` restorer — map those onto the survivors so loading
|
||||
them doesn't blow up at build time.
|
||||
"""
|
||||
if cfg.detector not in DETECTORS:
|
||||
cfg.detector = "yolo"
|
||||
if cfg.restorer not in RESTORERS:
|
||||
cfg.restorer = "deepmosaics"
|
||||
|
||||
@@ -1,216 +0,0 @@
|
||||
"""Weights-free, heuristic censorship detector (classic computer vision).
|
||||
|
||||
APPROXIMATE BY DESIGN. This detector uses hand-tuned CV heuristics, not a
|
||||
trained model. Its purpose is to make the whole pipeline runnable end-to-end
|
||||
and to exercise the :class:`Detector` interface. For real-world accuracy,
|
||||
replace it with a trained model (see ``yolo.py``, to be implemented) — the rest
|
||||
of the app does not need to change.
|
||||
|
||||
Heuristics:
|
||||
- black_bar: large, near-uniform very dark regions (classic censor bars).
|
||||
- mosaic: regions that reconstruct well from a coarse block grid (low
|
||||
residual) yet have high coarse-scale contrast (i.e. blocky, not flat).
|
||||
- blur: regions with local high-frequency energy far below the frame median,
|
||||
while still being textured (excludes genuinely flat areas).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ...config import DetectionConfig
|
||||
from ..video.frame import Frame
|
||||
from .base import Detector
|
||||
from .types import CensorType, Detection
|
||||
|
||||
|
||||
class ClassicCVDetector(Detector):
|
||||
def __init__(
|
||||
self,
|
||||
config: DetectionConfig | None = None,
|
||||
types: "set[CensorType] | None" = None,
|
||||
) -> None:
|
||||
self.cfg = config or DetectionConfig()
|
||||
# Which censorship kinds to look for. Default: all. The composite detector
|
||||
# restricts this to black_bar/blur (mosaic comes from the YOLO model).
|
||||
self.types = (
|
||||
types if types is not None
|
||||
else {CensorType.MOSAIC, CensorType.BLUR, CensorType.BLACK_BAR}
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ public
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
bgr = frame.image
|
||||
h0, w0 = bgr.shape[:2]
|
||||
scale = self._proc_scale(w0, h0)
|
||||
proc = (
|
||||
cv2.resize(bgr, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
|
||||
if scale != 1.0
|
||||
else bgr
|
||||
)
|
||||
gray = cv2.cvtColor(proc, cv2.COLOR_BGR2GRAY)
|
||||
ph, pw = gray.shape
|
||||
min_area = self.cfg.min_area_frac * pw * ph
|
||||
|
||||
dets: list[Detection] = []
|
||||
for ctype, fn, factor in (
|
||||
(CensorType.BLACK_BAR, self._detect_bars, 1.0),
|
||||
(CensorType.MOSAIC, self._detect_mosaic, 4.0),
|
||||
(CensorType.BLUR, self._detect_blur, 6.0),
|
||||
):
|
||||
if ctype not in self.types:
|
||||
continue
|
||||
try:
|
||||
dets += fn(proc, gray, min_area * factor)
|
||||
except Exception:
|
||||
# A failing heuristic must not break playback; skip it for this frame.
|
||||
continue
|
||||
|
||||
# Map proc-space coordinates back to source-frame pixels.
|
||||
inv = 1.0 / scale
|
||||
for d in dets:
|
||||
x, y, w, h = d.bbox
|
||||
d.bbox = (round(x * inv), round(y * inv), round(w * inv), round(h * inv))
|
||||
d.polygon = [(round(px * inv), round(py * inv)) for px, py in d.polygon]
|
||||
return self._dedup(dets)
|
||||
|
||||
# ----------------------------------------------------------------- helpers
|
||||
def _proc_scale(self, w: int, h: int) -> float:
|
||||
longest = max(w, h)
|
||||
if longest <= self.cfg.proc_max_dim:
|
||||
return 1.0
|
||||
return self.cfg.proc_max_dim / longest
|
||||
|
||||
@staticmethod
|
||||
def _local_std(g: np.ndarray, win: int) -> np.ndarray:
|
||||
"""Per-pixel standard deviation over a (win x win) box window."""
|
||||
mean = cv2.boxFilter(g, -1, (win, win))
|
||||
sqmean = cv2.boxFilter(g * g, -1, (win, win))
|
||||
var = np.maximum(sqmean - mean * mean, 0.0)
|
||||
return np.sqrt(var)
|
||||
|
||||
def _mask_to_detections(
|
||||
self,
|
||||
mask: np.ndarray,
|
||||
ctype: CensorType,
|
||||
min_area: float,
|
||||
base_score: float,
|
||||
min_extent: float = 0.0,
|
||||
min_side: int = 0,
|
||||
) -> list[Detection]:
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((3, 3), np.uint8))
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((9, 9), np.uint8))
|
||||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
out: list[Detection] = []
|
||||
for c in contours:
|
||||
area = cv2.contourArea(c)
|
||||
if area < min_area:
|
||||
continue
|
||||
x, y, w, h = cv2.boundingRect(c)
|
||||
if min(w, h) < min_side:
|
||||
continue # reject thin strips (e.g. edge false-positives)
|
||||
extent = area / float(w * h + 1e-6) # how rectangular the blob is
|
||||
if extent < min_extent:
|
||||
continue
|
||||
approx = cv2.approxPolyDP(c, 0.01 * cv2.arcLength(c, True), True)
|
||||
poly = [(int(p[0][0]), int(p[0][1])) for p in approx]
|
||||
score = float(np.clip(base_score + 0.25 * extent, 0.0, 1.0))
|
||||
out.append(Detection(type=ctype, score=score, bbox=(x, y, w, h), polygon=poly))
|
||||
return out
|
||||
|
||||
# --------------------------------------------------------------- detectors
|
||||
def _detect_bars(self, bgr, gray, min_area) -> list[Detection]:
|
||||
# Solid censor bars are achromatic (black OR white) rectangles. Requiring
|
||||
# low saturation + high rectangularity excludes large flat *colored* fills
|
||||
# that are common in drawn/anime backgrounds.
|
||||
hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
|
||||
sat, val = hsv[:, :, 1], hsv[:, :, 2]
|
||||
achromatic = sat < self.cfg.bar_saturation_max
|
||||
dark = (val < self.cfg.black_intensity) & achromatic
|
||||
light = (val > self.cfg.white_intensity) & achromatic
|
||||
mask = (dark | light).astype(np.uint8) * 255
|
||||
return self._mask_to_detections(
|
||||
mask, CensorType.BLACK_BAR, min_area, base_score=0.55,
|
||||
min_extent=self.cfg.bar_min_extent,
|
||||
)
|
||||
|
||||
def _detect_mosaic(self, bgr, gray, min_area) -> list[Detection]:
|
||||
g = gray.astype(np.float32)
|
||||
h, w = gray.shape
|
||||
win = 17
|
||||
# Lowest reconstruction residual across candidate tile sizes AND grid phases.
|
||||
# Real mosaics aren't aligned to the origin, so we try a few offsets per size
|
||||
# (phase-invariant) and keep the best fit.
|
||||
best_residual = np.full((h, w), np.inf, np.float32)
|
||||
for b in self.cfg.mosaic_block_sizes:
|
||||
half = b // 2
|
||||
for oy, ox in ((0, 0), (half, 0), (0, half), (half, half)):
|
||||
sub = g[oy:, ox:]
|
||||
sh, sw = sub.shape
|
||||
if sh < b or sw < b:
|
||||
continue
|
||||
small = cv2.resize(sub, (max(1, sw // b), max(1, sh // b)), interpolation=cv2.INTER_AREA)
|
||||
restored = cv2.resize(small, (sw, sh), interpolation=cv2.INTER_NEAREST)
|
||||
region = best_residual[oy:oy + sh, ox:ox + sw]
|
||||
np.minimum(region, np.abs(sub - restored), out=region)
|
||||
best_residual = cv2.boxFilter(best_residual, -1, (win, win))
|
||||
|
||||
contrast = self._local_std(g, win)
|
||||
# Mosaic has edges in BOTH directions; a lone straight boundary (flat-region
|
||||
# border, bar edge) has edge energy in only one — exclude those.
|
||||
gx = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)), -1, (win, win))
|
||||
gy = cv2.boxFilter(np.abs(cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)), -1, (win, win))
|
||||
both_dirs = (gx > self.cfg.mosaic_grad_min) & (gy > self.cfg.mosaic_grad_min)
|
||||
|
||||
blocky = best_residual < self.cfg.mosaic_residual_max
|
||||
textured = contrast > self.cfg.mosaic_contrast_min
|
||||
mask = (blocky & textured & both_dirs).astype(np.uint8) * 255
|
||||
return self._mask_to_detections(
|
||||
mask, CensorType.MOSAIC, min_area, base_score=0.50, min_side=self.cfg.mosaic_min_side
|
||||
)
|
||||
|
||||
def _detect_blur(self, bgr, gray, min_area) -> list[Detection]:
|
||||
g = gray.astype(np.float32)
|
||||
win = self.cfg.blur_window | 1 # force odd
|
||||
lap = cv2.Laplacian(g, cv2.CV_32F, ksize=3)
|
||||
sharpness = cv2.boxFilter(lap * lap, -1, (win, win)) # local high-freq energy
|
||||
median = float(np.median(sharpness)) + 1e-6
|
||||
contrast = self._local_std(g, win)
|
||||
|
||||
blurry = sharpness < median * self.cfg.blur_sharpness_ratio
|
||||
textured = contrast > self.cfg.blur_contrast_min
|
||||
mask = (blurry & textured).astype(np.uint8) * 255
|
||||
return self._mask_to_detections(
|
||||
mask, CensorType.BLUR, min_area, base_score=0.40, min_side=self.cfg.mosaic_min_side
|
||||
)
|
||||
|
||||
# ----------------------------------------------------------------- dedup
|
||||
def _dedup(self, dets: list[Detection]) -> list[Detection]:
|
||||
"""Greedy IoU suppression; prefer black_bar > mosaic > blur, then score."""
|
||||
priority = {
|
||||
CensorType.BLACK_BAR: 3,
|
||||
CensorType.MOSAIC: 2,
|
||||
CensorType.BLUR: 1,
|
||||
CensorType.UNKNOWN: 0,
|
||||
}
|
||||
dets = sorted(dets, key=lambda d: (priority[d.type], d.score), reverse=True)
|
||||
kept: list[Detection] = []
|
||||
for d in dets:
|
||||
if all(self._iou(d.bbox, k.bbox) < 0.5 for k in kept):
|
||||
kept.append(d)
|
||||
return kept
|
||||
|
||||
@staticmethod
|
||||
def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
|
||||
ax, ay, aw, ah = a
|
||||
bx, by, bw, bh = b
|
||||
ix = max(ax, bx)
|
||||
iy = max(ay, by)
|
||||
ix2 = min(ax + aw, bx + bw)
|
||||
iy2 = min(ay + ah, by + bh)
|
||||
iw, ih = max(0, ix2 - ix), max(0, iy2 - iy)
|
||||
inter = iw * ih
|
||||
union = aw * ah + bw * bh - inter
|
||||
return inter / union if union > 0 else 0.0
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Composite detector: runs several detectors and merges their results.
|
||||
|
||||
Used for the "combined" mode = YOLO (mosaic) + classic-CV (black bars / blur).
|
||||
Detections from all sub-detectors are concatenated, then de-duplicated by IoU
|
||||
(higher score wins) so overlapping hits from different detectors don't stack.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ..video.frame import Frame
|
||||
from .base import Detector
|
||||
from .types import Detection
|
||||
|
||||
|
||||
class CompositeDetector(Detector):
|
||||
def __init__(self, detectors: list[Detector], iou_threshold: float = 0.6) -> None:
|
||||
if not detectors:
|
||||
raise ValueError("CompositeDetector requires at least one detector")
|
||||
self._detectors = detectors
|
||||
self._iou = iou_threshold
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "Composite(" + " + ".join(d.name for d in self._detectors) + ")"
|
||||
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
merged: list[Detection] = []
|
||||
for detector in self._detectors:
|
||||
try:
|
||||
merged += detector.detect(frame)
|
||||
except Exception: # noqa: BLE001 - one detector failing must not kill the frame
|
||||
continue
|
||||
return self._dedup(merged)
|
||||
|
||||
def _dedup(self, dets: list[Detection]) -> list[Detection]:
|
||||
dets = sorted(dets, key=lambda d: d.score, reverse=True)
|
||||
kept: list[Detection] = []
|
||||
for d in dets:
|
||||
if all(self._iou_of(d.bbox, k.bbox) < self._iou for k in kept):
|
||||
kept.append(d)
|
||||
return kept
|
||||
|
||||
@staticmethod
|
||||
def _iou_of(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
|
||||
ax, ay, aw, ah = a
|
||||
bx, by, bw, bh = b
|
||||
ix, iy = max(ax, bx), max(ay, by)
|
||||
ix2, iy2 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
|
||||
inter = max(0, ix2 - ix) * max(0, iy2 - iy)
|
||||
union = aw * ah + bw * bh - inter
|
||||
return inter / union if union > 0 else 0.0
|
||||
@@ -3,14 +3,15 @@
|
||||
Kept separate from ``app.py`` so both the app bootstrap and the UI can build
|
||||
detectors without an import cycle. Raises ``ValueError`` (not ``SystemExit``) on
|
||||
bad config so the GUI can show the message instead of exiting.
|
||||
|
||||
Only the YOLO detector is supported — the classic-CV heuristic (and the composite
|
||||
mode that combined them) were removed: they were noisy/approximate on real footage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ...config import AppConfig
|
||||
from .base import Detector
|
||||
from .classic_cv import ClassicCVDetector
|
||||
from .types import CensorType
|
||||
|
||||
|
||||
def _require_model(config: AppConfig) -> str:
|
||||
@@ -23,19 +24,6 @@ def _require_model(config: AppConfig) -> str:
|
||||
|
||||
|
||||
def build_detector(config: AppConfig) -> Detector:
|
||||
if config.detector == "classic":
|
||||
return ClassicCVDetector(config.detection)
|
||||
if config.detector == "yolo":
|
||||
from .yolo import YoloDetector # lazy: pulls torch/ultralytics
|
||||
from .yolo import YoloDetector # lazy: pulls torch/ultralytics
|
||||
|
||||
return YoloDetector(_require_model(config), config.detection)
|
||||
if config.detector == "combined":
|
||||
# YOLO handles mosaic; classic-CV handles black bars / blur.
|
||||
from .composite import CompositeDetector
|
||||
from .yolo import YoloDetector
|
||||
|
||||
return CompositeDetector([
|
||||
YoloDetector(_require_model(config), config.detection),
|
||||
ClassicCVDetector(config.detection, types={CensorType.BLACK_BAR, CensorType.BLUR}),
|
||||
])
|
||||
raise ValueError(f"Неизвестный детектор: {config.detector!r}")
|
||||
return YoloDetector(_require_model(config), config.detection)
|
||||
|
||||
@@ -7,7 +7,7 @@ a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose cl
|
||||
names map onto :class:`CensorType`.
|
||||
|
||||
Heavy imports (``ultralytics``/``torch``) happen lazily in ``__init__`` so the
|
||||
rest of the app — and the classic-CV detector — never pull them in.
|
||||
rest of the app never pulls them in until detection actually runs.
|
||||
|
||||
Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
|
||||
"""
|
||||
|
||||
@@ -30,6 +30,7 @@ FRAMES_DIR = "frames"
|
||||
CACHE_FILE = "detections.json"
|
||||
COLLECTIONS_DIR = "collections"
|
||||
FAVORITES_DIR = "Избранное" # the single default collection ("в избранное")
|
||||
RESTORED_DIR = "restored" # batch "расцензурить все" output (kept out of frames/)
|
||||
_VERSION = 1
|
||||
|
||||
# The subset of AppConfig fields a project remembers (mirrored to/from project.json).
|
||||
@@ -78,6 +79,12 @@ class Project:
|
||||
"""The single default collection — frames moved "to favorites" land here."""
|
||||
return self.collections_dir / FAVORITES_DIR
|
||||
|
||||
@property
|
||||
def restored_dir(self) -> Path:
|
||||
"""Batch restoration output ("расцензурить все") — mirrors frame basenames.
|
||||
Kept out of ``frames/`` so results aren't listed/re-detected/re-restored."""
|
||||
return self.root / RESTORED_DIR
|
||||
|
||||
# ------------------------------------------------------------- lifecycle
|
||||
@classmethod
|
||||
def create(
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
"""Restorer interface — "un-censor" detected regions of an image.
|
||||
|
||||
A Restorer takes an image plus the detected censored regions and returns a new
|
||||
image with those regions reconstructed/filled. This mirrors the ``Detector``
|
||||
abstraction so different engines (classic inpaint now; a generative model like
|
||||
DeepMosaics / LADA later) plug in behind the same interface.
|
||||
image with those regions reconstructed. This mirrors the ``Detector`` abstraction
|
||||
so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug
|
||||
in behind the same interface. Note DeepMosaics locates the mosaic itself, so the
|
||||
``detections`` argument is currently advisory (unused by the DeepMosaics engines).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -24,7 +25,18 @@ class Cancelled(Exception):
|
||||
"""Raised by a Restorer when ``should_cancel`` asked it to stop."""
|
||||
|
||||
|
||||
# Sequence (batch) restore callbacks — see ``Restorer.restore_sequence``.
|
||||
FrameGetter = Callable[[int], np.ndarray] # index -> BGR image
|
||||
DetGetter = Callable[[int], list[Detection]] # index -> that frame's detections
|
||||
ResultSink = Callable[[int, np.ndarray], None] # (index, restored image) -> None
|
||||
|
||||
|
||||
class Restorer(ABC):
|
||||
#: Whether this engine uses *neighbouring* frames (so a batch run must feed it a
|
||||
#: contiguous, ordered sequence — see :meth:`restore_sequence`). Per-frame engines
|
||||
#: leave this False; the temporal DeepMosaics (BVDNet) sets it True.
|
||||
temporal: bool = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return type(self).__name__
|
||||
@@ -42,3 +54,24 @@ class Restorer(ABC):
|
||||
True the engine should abort and raise :class:`Cancelled`.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def restore_sequence(
|
||||
self,
|
||||
count: int,
|
||||
get_frame: FrameGetter,
|
||||
get_dets: DetGetter,
|
||||
emit: ResultSink,
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> None:
|
||||
"""Restore ``count`` frames *in order*, calling ``emit(i, restored)`` for each.
|
||||
|
||||
The default treats every frame independently (just loops :meth:`restore`).
|
||||
Temporal engines override this to pull neighbouring frames via ``get_frame``
|
||||
and carry recurrent state across the sequence. ``get_frame``/``get_dets`` are
|
||||
lazy so the engine only reads the frames it needs; ``emit`` lets the caller
|
||||
stream results to disk instead of holding them all in memory.
|
||||
"""
|
||||
for i in range(count):
|
||||
if should_cancel is not None and should_cancel():
|
||||
raise Cancelled("Восстановление отменено")
|
||||
emit(i, self.restore(get_frame(i), get_dets(i), should_cancel))
|
||||
|
||||
@@ -26,18 +26,29 @@ from types import SimpleNamespace
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Cancelled, Restorer
|
||||
from .base import (
|
||||
CancelCheck,
|
||||
Cancelled,
|
||||
DetGetter,
|
||||
FrameGetter,
|
||||
ResultSink,
|
||||
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).
|
||||
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``. The video model is skipped — it can't run per-frame.
|
||||
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:
|
||||
@@ -48,13 +59,23 @@ def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
|
||||
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:
|
||||
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()
|
||||
@@ -90,7 +111,7 @@ class DeepMosaicsRestorer(Restorer):
|
||||
)
|
||||
model = Path(model_path)
|
||||
self._netg = _netg_kind(model.name) # raises on a video model
|
||||
pos = self._find_mosaic_position(model, deepmosaics_dir)
|
||||
pos = _find_mosaic_position(model, deepmosaics_dir)
|
||||
if pos is None:
|
||||
raise ValueError(
|
||||
"Рядом с clean-моделью не найден mosaic_position.pth.\n"
|
||||
@@ -101,13 +122,6 @@ class DeepMosaicsRestorer(Restorer):
|
||||
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})"
|
||||
@@ -167,3 +181,147 @@ class DeepMosaicsRestorer(Restorer):
|
||||
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,
|
||||
python_exe: str | None = None, # unused (in-process); kept for factory parity
|
||||
gpu_id: str = "0",
|
||||
) -> None:
|
||||
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 # noqa: E402
|
||||
if self._gpu != "-1" and not torch.cuda.is_available():
|
||||
self._gpu = "-1" # CPU fallback (see DeepMosaicsRestorer for why)
|
||||
|
||||
from models import loadmodel, runmodel # type: ignore # noqa: E402
|
||||
import util.data as data # type: ignore # noqa: E402
|
||||
import util.image_processing as impro # type: ignore # noqa: E402
|
||||
|
||||
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
|
||||
|
||||
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:
|
||||
emit(i, img_origin.copy()) # no mosaic here; recurrence carries over
|
||||
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)
|
||||
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)
|
||||
emit(i, impro.replace_mosaic(img_origin.copy(), img_fake, mask, x, y, size, opt.no_feather))
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
"""Restorer factory: build a Restorer from the app config.
|
||||
|
||||
- ``inpaint``: cv2 baseline (no weights, no GPU; fills, doesn't reconstruct).
|
||||
- ``deepmosaics``: real generative mosaic removal. The DeepMosaics network code is
|
||||
vendored (``_deepmosaics/``, GPL-3.0) and run in-process; the user supplies only the
|
||||
clean weights (+ ``mosaic_position.pth`` alongside). A CUDA GPU is recommended.
|
||||
Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but
|
||||
did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``,
|
||||
GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth``
|
||||
alongside). A CUDA GPU is recommended.
|
||||
|
||||
- ``deepmosaics``: per-frame generative mosaic removal (image model).
|
||||
- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for
|
||||
coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,25 +15,28 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .base import Restorer
|
||||
from .inpaint import InpaintRestorer
|
||||
|
||||
if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed)
|
||||
from ...config import AppConfig
|
||||
|
||||
|
||||
def build_restorer(name: str = "inpaint", config: "AppConfig | None" = None) -> Restorer:
|
||||
if name == "inpaint":
|
||||
return InpaintRestorer()
|
||||
def build_restorer(name: str = "deepmosaics", config: "AppConfig | None" = None) -> Restorer:
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
if name == "deepmosaics":
|
||||
from .deepmosaics import DeepMosaicsRestorer
|
||||
|
||||
if config is None:
|
||||
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
||||
return DeepMosaicsRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "deepmosaics_video":
|
||||
from .deepmosaics import DeepMosaicsVideoRestorer
|
||||
|
||||
return DeepMosaicsVideoRestorer(
|
||||
config.dm_dir, config.dm_model, config.dm_python, config.dm_gpu
|
||||
)
|
||||
if name == "lada":
|
||||
raise ValueError(
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics или inpaint. См. README."
|
||||
"Движок LADA пока не подключён. Используйте DeepMosaics. См. README."
|
||||
)
|
||||
raise ValueError(f"Неизвестный режим восстановления: {name!r}")
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
"""Classic inpainting restorer (cv2) — the always-available baseline.
|
||||
|
||||
HONEST LIMITATION: cv2 inpainting fills the masked region by propagating
|
||||
surrounding pixels. It removes the mosaic/bar but does NOT reconstruct the hidden
|
||||
detail — it smooths/guesses. For real reconstruction a generative model
|
||||
(DeepMosaics / LADA) is needed; this is the no-weights, no-GPU fallback so the
|
||||
"Расцензурить кадр" flow works end-to-end today.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
from .base import CancelCheck, Restorer
|
||||
from .mask import detections_to_mask
|
||||
|
||||
|
||||
class InpaintRestorer(Restorer):
|
||||
def __init__(self, radius: int = 3, dilate: int = 2, method: str = "telea") -> None:
|
||||
self.radius = radius
|
||||
self.dilate = dilate
|
||||
self.method = method
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"InpaintRestorer({self.method})"
|
||||
|
||||
def restore(
|
||||
self,
|
||||
image: np.ndarray,
|
||||
detections: list[Detection],
|
||||
should_cancel: CancelCheck | None = None,
|
||||
) -> np.ndarray:
|
||||
# Single cv2.inpaint call — effectively instant, so cancellation is moot.
|
||||
if not detections:
|
||||
return image.copy()
|
||||
mask = detections_to_mask(image.shape, detections, dilate=self.dilate)
|
||||
flags = cv2.INPAINT_TELEA if self.method == "telea" else cv2.INPAINT_NS
|
||||
return cv2.inpaint(image, mask, self.radius, flags)
|
||||
@@ -1,26 +0,0 @@
|
||||
"""Build a binary mask of the censored regions from detections."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..detection.types import Detection
|
||||
|
||||
|
||||
def detections_to_mask(
|
||||
shape: tuple[int, int], detections: list[Detection], dilate: int = 0
|
||||
) -> np.ndarray:
|
||||
"""White (255) over every detected region (polygon if present, else bbox)."""
|
||||
h, w = shape[:2]
|
||||
mask = np.zeros((h, w), np.uint8)
|
||||
for d in detections:
|
||||
if len(d.polygon) >= 3:
|
||||
cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255)
|
||||
else:
|
||||
x, y, bw, bh = d.bbox
|
||||
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1)
|
||||
if dilate > 0:
|
||||
k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8)
|
||||
mask = cv2.dilate(mask, k)
|
||||
return mask
|
||||
+191
-48
@@ -1,18 +1,66 @@
|
||||
"""Probe the PyTorch / CUDA situation, so the UI can show a device badge.
|
||||
"""Diagnose the PyTorch / CUDA situation so the UI can explain *why* it's on CPU.
|
||||
|
||||
Pure (no Qt). ``gather()`` imports torch (slow / heavy) — call it off the GUI
|
||||
thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and
|
||||
how to enable the GPU.
|
||||
Pure (no Qt). ``gather()`` is the heavy part — it imports torch and shells out to
|
||||
``nvidia-smi`` — so call it off the GUI thread. ``analyze()`` is fast formatting on
|
||||
the gathered dict and figures out the most likely cause + concrete fix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# pip index for the CUDA build (matches the README).
|
||||
CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121"
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# pip indexes for the CUDA builds (see README). cu121 needs a driver with CUDA >= 12.1.
|
||||
CUDA_WHEELS = {
|
||||
"cu121": "https://download.pytorch.org/whl/cu121",
|
||||
"cu118": "https://download.pytorch.org/whl/cu118",
|
||||
}
|
||||
|
||||
_NO_WINDOW = 0x08000000 if sys.platform == "win32" else 0 # CREATE_NO_WINDOW
|
||||
|
||||
|
||||
def _run_nvidia_smi() -> dict:
|
||||
"""Probe the NVIDIA driver/GPU via nvidia-smi. Never raises."""
|
||||
out: dict = {"found": False, "gpus": [], "driver_version": None, "cuda_driver": None}
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe and sys.platform == "win32":
|
||||
candidate = r"C:\Windows\System32\nvidia-smi.exe"
|
||||
exe = candidate if shutil.os.path.isfile(candidate) else None
|
||||
if not exe:
|
||||
return out
|
||||
# GPU names + driver version (robust CSV form).
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[exe, "--query-gpu=name,driver_version", "--format=csv,noheader,nounits"],
|
||||
capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
||||
)
|
||||
if r.returncode == 0:
|
||||
out["found"] = True
|
||||
for line in r.stdout.strip().splitlines():
|
||||
parts = [p.strip() for p in line.split(",")]
|
||||
if parts and parts[0]:
|
||||
out["gpus"].append(parts[0])
|
||||
if len(parts) > 1 and parts[1]:
|
||||
out["driver_version"] = parts[1]
|
||||
except Exception: # noqa: BLE001 - any failure => "not found"
|
||||
return out
|
||||
# Max CUDA version the driver supports (only in the plain header).
|
||||
try:
|
||||
r2 = subprocess.run(
|
||||
[exe], capture_output=True, text=True, timeout=10, creationflags=_NO_WINDOW,
|
||||
)
|
||||
m = re.search(r"CUDA Version:\s*([\d.]+)", r2.stdout)
|
||||
if m:
|
||||
out["cuda_driver"] = m.group(1)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
return out
|
||||
|
||||
|
||||
def gather() -> dict:
|
||||
"""Collect torch/CUDA facts. Never raises — missing torch is a valid result."""
|
||||
"""Collect torch + NVIDIA facts. Never raises — missing torch/GPU are valid."""
|
||||
info: dict = {
|
||||
"installed": False,
|
||||
"version": None, # torch.__version__ (e.g. "2.12.0+cpu")
|
||||
@@ -20,27 +68,38 @@ def gather() -> dict:
|
||||
"cuda_available": False,
|
||||
"device_name": None, # the active GPU's name, if any
|
||||
"import_error": None,
|
||||
# filled by nvidia-smi:
|
||||
"nvidia_smi": False,
|
||||
"gpus": [],
|
||||
"driver_version": None,
|
||||
"cuda_driver": None,
|
||||
}
|
||||
try:
|
||||
import torch
|
||||
except Exception as exc: # noqa: BLE001 - report any import failure, not just ImportError
|
||||
except Exception as exc: # noqa: BLE001 - report any import failure
|
||||
info["import_error"] = str(exc)
|
||||
return info
|
||||
info["installed"] = True
|
||||
info["version"] = getattr(torch, "__version__", None)
|
||||
try:
|
||||
info["built_cuda"] = torch.version.cuda
|
||||
except Exception: # noqa: BLE001
|
||||
info["built_cuda"] = None
|
||||
try:
|
||||
info["cuda_available"] = bool(torch.cuda.is_available())
|
||||
except Exception: # noqa: BLE001
|
||||
info["cuda_available"] = False
|
||||
if info["cuda_available"]:
|
||||
else:
|
||||
info["installed"] = True
|
||||
info["version"] = getattr(torch, "__version__", None)
|
||||
try:
|
||||
info["device_name"] = torch.cuda.get_device_name(0)
|
||||
info["built_cuda"] = torch.version.cuda
|
||||
except Exception: # noqa: BLE001
|
||||
info["device_name"] = None
|
||||
info["built_cuda"] = None
|
||||
try:
|
||||
info["cuda_available"] = bool(torch.cuda.is_available())
|
||||
except Exception: # noqa: BLE001
|
||||
info["cuda_available"] = False
|
||||
if info["cuda_available"]:
|
||||
try:
|
||||
info["device_name"] = torch.cuda.get_device_name(0)
|
||||
except Exception: # noqa: BLE001
|
||||
info["device_name"] = None
|
||||
|
||||
smi = _run_nvidia_smi()
|
||||
info["nvidia_smi"] = smi["found"]
|
||||
info["gpus"] = smi["gpus"]
|
||||
info["driver_version"] = smi["driver_version"]
|
||||
info["cuda_driver"] = smi["cuda_driver"]
|
||||
return info
|
||||
|
||||
|
||||
@@ -48,33 +107,117 @@ def device_label(info: dict) -> str:
|
||||
return "CUDA" if info.get("cuda_available") else "CPU"
|
||||
|
||||
|
||||
def reason(info: dict) -> str:
|
||||
"""One-sentence human explanation of the current device choice."""
|
||||
if not info.get("installed"):
|
||||
return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics "
|
||||
"работают на CPU (классический детектор torch не требует).")
|
||||
if info.get("cuda_available"):
|
||||
name = info.get("device_name") or "GPU"
|
||||
return f"PyTorch использует CUDA: {name}. Вычисления идут на видеокарте."
|
||||
version = info.get("version") or "?"
|
||||
built = info.get("built_cuda")
|
||||
if not built:
|
||||
return (f"Установлена CPU-сборка PyTorch ({version}) — без поддержки CUDA, "
|
||||
"поэтому вычисления идут на процессоре (медленно).")
|
||||
return (f"PyTorch собран с CUDA {built} ({version}), но GPU недоступен: нет "
|
||||
"NVIDIA-видеокарты, не установлен/устарел драйвер, либо версия CUDA "
|
||||
"несовместима с драйвером.")
|
||||
def _ver_tuple(v: str | None) -> tuple[int, ...]:
|
||||
try:
|
||||
return tuple(int(x) for x in str(v).split(".")[:2])
|
||||
except (ValueError, AttributeError):
|
||||
return ()
|
||||
|
||||
|
||||
def install_hint() -> str:
|
||||
"""Steps to enable the GPU (shown when running on CPU)."""
|
||||
def recommend_channel(info: dict) -> str:
|
||||
"""Pick the pip CUDA wheel index that matches the driver (cu118 for older)."""
|
||||
cd = _ver_tuple(info.get("cuda_driver"))
|
||||
if cd and cd < (12, 1):
|
||||
return "cu118"
|
||||
return "cu121"
|
||||
|
||||
|
||||
def install_command(channel: str = "cu121") -> str:
|
||||
"""The pip commands to (re)install the chosen CUDA build.
|
||||
|
||||
Targets the **running interpreter** (``sys.executable -m pip``) so the command
|
||||
hits the same venv that runs the app — not whatever ``pip`` is on PATH. (A common
|
||||
trap: running bare ``pip`` in a global shell while torch lives in the project venv.)
|
||||
"""
|
||||
index = CUDA_WHEELS.get(channel, CUDA_WHEELS["cu121"])
|
||||
py = sys.executable or "python"
|
||||
q = f'"{py}"' if " " in py else py
|
||||
return (
|
||||
"Как включить GPU (NVIDIA):\n"
|
||||
"1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n"
|
||||
"2. Переустановите PyTorch со сборкой CUDA:\n\n"
|
||||
" pip uninstall -y torch torchvision\n"
|
||||
f" pip install torch torchvision --index-url {CUDA_WHEEL_INDEX}\n\n"
|
||||
"3. Перезапустите приложение.\n\n"
|
||||
"Классический детектор работает и без CUDA. На CPU детекция и расцензуривание "
|
||||
"просто медленнее."
|
||||
f"{q} -m pip uninstall -y torch torchvision\n"
|
||||
f"{q} -m pip install torch torchvision --index-url {index}"
|
||||
)
|
||||
|
||||
|
||||
def analyze(info: dict) -> dict:
|
||||
"""Turn the raw facts into {summary, details[list], steps, command}."""
|
||||
installed = info.get("installed")
|
||||
cuda = info.get("cuda_available")
|
||||
built = info.get("built_cuda")
|
||||
gpus = info.get("gpus") or []
|
||||
smi = info.get("nvidia_smi")
|
||||
driver = info.get("driver_version")
|
||||
cuda_driver = info.get("cuda_driver")
|
||||
channel = recommend_channel(info)
|
||||
command = install_command(channel)
|
||||
|
||||
if not installed:
|
||||
build_str = "не установлен"
|
||||
elif built:
|
||||
build_str = f"CUDA {built}"
|
||||
else:
|
||||
build_str = "CPU-only (+cpu)"
|
||||
gpu_str = (
|
||||
", ".join(gpus) if gpus
|
||||
else ("не обнаружена" if smi or info.get("nvidia_smi") is False else "nvidia-smi не найден")
|
||||
)
|
||||
if not gpus and not smi:
|
||||
gpu_str = "nvidia-smi не найден (нет драйвера NVIDIA?)"
|
||||
details = [
|
||||
f"PyTorch: {info.get('version') or 'не установлен'}",
|
||||
f"Сборка PyTorch: {build_str}",
|
||||
f"CUDA доступна в PyTorch: {'да' if cuda else 'нет'}",
|
||||
f"Видеокарта (nvidia-smi): {gpu_str}",
|
||||
f"Драйвер NVIDIA: {driver or '—'}",
|
||||
f"Макс. CUDA драйвера: {cuda_driver or '—'}",
|
||||
f"Интерпретатор (venv): {sys.executable}",
|
||||
]
|
||||
if info.get("import_error"):
|
||||
details.append(f"Ошибка импорта torch: {info['import_error']}")
|
||||
|
||||
if not installed:
|
||||
summary = "PyTorch не установлен — детекция YOLO и DeepMosaics идут на CPU."
|
||||
steps = ("Установите PyTorch (CUDA-сборку, если есть NVIDIA-видеокарта):\n\n"
|
||||
+ command + "\n\nДля YOLO также: pip install -e \".[yolo]\"")
|
||||
elif cuda:
|
||||
gpu = info.get("device_name") or (gpus[0] if gpus else "GPU")
|
||||
summary = f"Всё в порядке: PyTorch использует CUDA. Активный GPU: {gpu}."
|
||||
steps = "GPU уже задействован — ничего делать не нужно."
|
||||
elif not built: # CPU-only torch build — the usual case
|
||||
if gpus:
|
||||
summary = (
|
||||
"Главная причина: установлена CPU-сборка PyTorch (+cpu) — она физически "
|
||||
f"не умеет в CUDA. Видеокарта ({gpus[0]}) и драйвер {driver or '?'} на месте, "
|
||||
"поэтому достаточно переустановить PyTorch со сборкой CUDA."
|
||||
)
|
||||
steps = ("Переустановите PyTorch под CUDA, затем перезапустите приложение:\n\n"
|
||||
+ command)
|
||||
else:
|
||||
summary = (
|
||||
"Установлена CPU-сборка PyTorch (+cpu), и видеокарта NVIDIA не обнаружена "
|
||||
"(nvidia-smi не отвечает). Либо нет NVIDIA GPU, либо не установлен драйвер."
|
||||
)
|
||||
steps = ("1. Проверьте видеокарту и драйвер: в консоли выполните nvidia-smi\n"
|
||||
"2. Если NVIDIA GPU есть — переустановите PyTorch под CUDA:\n\n"
|
||||
+ command +
|
||||
"\n\nБез NVIDIA GPU всё работает на CPU — просто медленнее.")
|
||||
else: # CUDA-enabled torch build, but CUDA still not available
|
||||
if not gpus:
|
||||
summary = (
|
||||
f"PyTorch собран под CUDA {built}, но видеокарта/драйвер NVIDIA не найдены. "
|
||||
"Скорее всего не установлен драйвер NVIDIA или нет GPU."
|
||||
)
|
||||
steps = "Установите свежий драйвер NVIDIA и перезапустите. Проверка: nvidia-smi."
|
||||
elif cuda_driver and _ver_tuple(cuda_driver) < _ver_tuple(built):
|
||||
summary = (
|
||||
f"Драйвер поддерживает CUDA {cuda_driver}, а PyTorch собран под CUDA {built} — "
|
||||
"версия драйвера слишком старая."
|
||||
)
|
||||
steps = ("Вариант A — обновите драйвер NVIDIA (рекомендуется).\n"
|
||||
"Вариант B — поставьте PyTorch под CUDA вашего драйвера:\n\n" + command)
|
||||
else:
|
||||
summary = (
|
||||
f"PyTorch собран под CUDA {built}, GPU ({gpus[0]}) есть, но CUDA недоступна — "
|
||||
"вероятен конфликт версий или повреждённая установка."
|
||||
)
|
||||
steps = "Переустановите PyTorch под CUDA:\n\n" + command
|
||||
return {"summary": summary, "details": details, "steps": steps, "command": command}
|
||||
|
||||
+156
-68
@@ -26,11 +26,10 @@ import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from PySide6.QtCore import Qt, QThreadPool
|
||||
from PySide6.QtGui import QAction, QBrush, QColor, QKeySequence, QShortcut
|
||||
from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcut
|
||||
from PySide6.QtWidgets import (
|
||||
QAbstractItemView,
|
||||
QApplication,
|
||||
QComboBox,
|
||||
QDialog,
|
||||
QDoubleSpinBox,
|
||||
QFileDialog,
|
||||
@@ -41,6 +40,7 @@ from PySide6.QtWidgets import (
|
||||
QListWidgetItem,
|
||||
QMainWindow,
|
||||
QMessageBox,
|
||||
QPlainTextEdit,
|
||||
QProgressBar,
|
||||
QPushButton,
|
||||
QSplitter,
|
||||
@@ -51,7 +51,7 @@ from PySide6.QtWidgets import (
|
||||
)
|
||||
|
||||
from .. import settings_store
|
||||
from ..config import AppConfig
|
||||
from ..config import AppConfig, normalize_config
|
||||
from ..core.detection import cache as detection_cache
|
||||
from ..core.detection.factory import build_detector
|
||||
from ..core.detection.types import Detection
|
||||
@@ -68,7 +68,6 @@ from .workers import Job
|
||||
|
||||
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||||
_VIDEO_FILTER = "Видео (*.mp4 *.mkv *.avi *.mov *.webm *.m4v);;Все файлы (*.*)"
|
||||
_DETECTORS = ["classic", "yolo", "combined"]
|
||||
|
||||
|
||||
class MainWindow(QMainWindow):
|
||||
@@ -120,6 +119,8 @@ class MainWindow(QMainWindow):
|
||||
file_menu.addAction("Детектировать все заново", lambda: self._detect_all(True))
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("Движок восстановления…", self._open_restore_settings)
|
||||
file_menu.addAction("Расцензурить все (дозапуск)", lambda: self._restore_all(False))
|
||||
file_menu.addAction("Расцензурить все заново", lambda: self._restore_all(True))
|
||||
file_menu.addSeparator()
|
||||
file_menu.addAction("В избранное", self._move_to_favorites).setShortcut("Ctrl+M")
|
||||
file_menu.addSeparator()
|
||||
@@ -136,14 +137,9 @@ class MainWindow(QMainWindow):
|
||||
tb.addAction(from_video)
|
||||
tb.addSeparator()
|
||||
|
||||
tb.addWidget(QLabel(" Детектор: "))
|
||||
self.detector_combo = QComboBox()
|
||||
self.detector_combo.addItems(_DETECTORS)
|
||||
self.detector_combo.setCurrentText(self._cfg.detector)
|
||||
self.detector_combo.currentTextChanged.connect(self._on_detector_changed)
|
||||
tb.addWidget(self.detector_combo)
|
||||
|
||||
tb.addWidget(QLabel(" Детектор: YOLO "))
|
||||
self.model_action = QAction("Модель…", self, triggered=self._choose_model)
|
||||
self.model_action.setToolTip("Выбрать веса YOLO (.pt) — модель LADA для мозаики")
|
||||
tb.addAction(self.model_action)
|
||||
|
||||
tb.addSeparator()
|
||||
@@ -166,6 +162,12 @@ class MainWindow(QMainWindow):
|
||||
restore = QAction("Расцензурить кадр", self, triggered=self._restore_current)
|
||||
restore.setToolTip("Восстановить найденные области на текущем кадре")
|
||||
tb.addAction(restore)
|
||||
restore_all = QAction("Расцензурить все", self, triggered=lambda: self._restore_all(False))
|
||||
restore_all.setToolTip("Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон)")
|
||||
tb.addAction(restore_all)
|
||||
restore_regen = QAction("Все заново (расцензур)", self, triggered=lambda: self._restore_all(True))
|
||||
restore_regen.setToolTip("Перерасцензурить все кадры заново (перезапись restored/)")
|
||||
tb.addAction(restore_regen)
|
||||
self.toggle_restored_action = QAction("Показать оригинал", self, triggered=self._toggle_restored)
|
||||
self.toggle_restored_action.setEnabled(False)
|
||||
tb.addAction(self.toggle_restored_action)
|
||||
@@ -359,27 +361,58 @@ class MainWindow(QMainWindow):
|
||||
|
||||
info = self._device_info if self._device_info else torch_info.gather()
|
||||
cuda = bool(info.get("cuda_available"))
|
||||
a = torch_info.analyze(info)
|
||||
self._install_command = a["command"] # what the Copy button will copy
|
||||
|
||||
lines = [
|
||||
torch_info.reason(info),
|
||||
"ВЕРДИКТ:",
|
||||
a["summary"],
|
||||
"",
|
||||
"Диагностика:",
|
||||
f" • PyTorch: {info.get('version') or 'не установлен'}",
|
||||
f" • Сборка CUDA: {info.get('built_cuda') or '— (CPU-сборка)'}",
|
||||
f" • CUDA доступна: {'да' if cuda else 'нет'}",
|
||||
*(f" • {d}" for d in a["details"]),
|
||||
"",
|
||||
"Что делать:",
|
||||
a["steps"],
|
||||
]
|
||||
if info.get("device_name"):
|
||||
lines.append(f" • GPU: {info['device_name']}")
|
||||
if info.get("import_error"):
|
||||
lines.append(f" • Ошибка импорта torch: {info['import_error']}")
|
||||
|
||||
box = QMessageBox(self)
|
||||
box.setIcon(QMessageBox.Information if cuda else QMessageBox.Warning)
|
||||
box.setWindowTitle("Устройство: " + ("CUDA (GPU)" if cuda else "CPU"))
|
||||
box.setText("\n".join(lines))
|
||||
# A real dialog (not QMessageBox) so the text — incl. the install command — is
|
||||
# selectable, and a Copy button drops the pip command straight onto the clipboard.
|
||||
dlg = QDialog(self)
|
||||
dlg.setWindowTitle("Почему " + ("GPU" if cuda else "CPU") + " — диагностика PyTorch/CUDA")
|
||||
dlg.resize(620, 480)
|
||||
layout = QVBoxLayout(dlg)
|
||||
|
||||
text = QPlainTextEdit()
|
||||
text.setReadOnly(True)
|
||||
text.setPlainText("\n".join(lines))
|
||||
mono = QFont("Consolas")
|
||||
mono.setStyleHint(QFont.Monospace)
|
||||
text.setFont(mono)
|
||||
layout.addWidget(text, 1)
|
||||
|
||||
buttons = QHBoxLayout()
|
||||
if not cuda:
|
||||
box.setInformativeText(torch_info.install_hint())
|
||||
box.setTextInteractionFlags(Qt.TextSelectableByMouse) # let the user copy commands
|
||||
box.exec()
|
||||
copy_btn = QPushButton("Скопировать команду установки")
|
||||
copy_btn.clicked.connect(self._copy_install_command)
|
||||
buttons.addWidget(copy_btn)
|
||||
recheck = QPushButton("Проверить заново")
|
||||
recheck.setToolTip("Перепроверить torch/CUDA (например, после переустановки)")
|
||||
recheck.clicked.connect(lambda: (self._probe_device(), dlg.accept()))
|
||||
buttons.addWidget(recheck)
|
||||
buttons.addStretch(1)
|
||||
close_btn = QPushButton("Закрыть")
|
||||
close_btn.clicked.connect(dlg.accept)
|
||||
buttons.addWidget(close_btn)
|
||||
layout.addLayout(buttons)
|
||||
dlg.exec()
|
||||
|
||||
def _copy_install_command(self) -> None:
|
||||
command = getattr(self, "_install_command", None)
|
||||
if not command:
|
||||
from ..core import torch_info
|
||||
command = torch_info.install_command()
|
||||
QApplication.clipboard().setText(command)
|
||||
self.statusBar().showMessage("Команда установки скопирована в буфер обмена")
|
||||
|
||||
# ------------------------------------------------------------- cancellation
|
||||
def _begin_busy(self, total: int | None = None) -> None:
|
||||
@@ -388,7 +421,6 @@ class MainWindow(QMainWindow):
|
||||
self._cancel = False
|
||||
self.stop_action.setEnabled(True)
|
||||
# Disable inputs that would race a running job (they clear cache / rebuild engines).
|
||||
self.detector_combo.setEnabled(False)
|
||||
self.model_action.setEnabled(False)
|
||||
if total is None:
|
||||
self.progress.setRange(0, 0) # indeterminate
|
||||
@@ -400,7 +432,6 @@ class MainWindow(QMainWindow):
|
||||
def _end_busy(self) -> None:
|
||||
self._busy = False
|
||||
self.stop_action.setEnabled(False)
|
||||
self.detector_combo.setEnabled(True)
|
||||
self.model_action.setEnabled(True)
|
||||
self.progress.setVisible(False)
|
||||
self.progress.setRange(0, 100) # leave it determinate for the next user
|
||||
@@ -459,19 +490,19 @@ class MainWindow(QMainWindow):
|
||||
self._detector_key = key
|
||||
return self._detector
|
||||
|
||||
def _on_detector_changed(self, name: str) -> None:
|
||||
self._cfg.detector = name
|
||||
# YOLO/combined need a model. Auto-pick a known one from models/ if we have it;
|
||||
# only prompt when nothing suitable is found (don't nag when the path is obvious).
|
||||
if name in ("yolo", "combined") and not self._cfg.model_path:
|
||||
found = self._auto_find_model()
|
||||
if found:
|
||||
self._cfg.model_path = found
|
||||
self.statusBar().showMessage(f"Модель найдена автоматически: {found}")
|
||||
else:
|
||||
self._choose_model()
|
||||
self._persist_settings()
|
||||
self._invalidate_results()
|
||||
def _ensure_model(self) -> None:
|
||||
"""Make sure the YOLO detector has weights — auto-pick from ./models silently.
|
||||
|
||||
Called on project open. Doesn't prompt (the user can pick via "Модель…"); the
|
||||
detector factory raises a clear message if a detect is attempted without one.
|
||||
"""
|
||||
if self._cfg.model_path and Path(self._cfg.model_path).is_file():
|
||||
return
|
||||
found = self._auto_find_model()
|
||||
if found:
|
||||
self._cfg.model_path = found
|
||||
self.statusBar().showMessage(f"Модель YOLO найдена автоматически: {found}")
|
||||
self._persist_settings()
|
||||
|
||||
@staticmethod
|
||||
def _auto_find_model() -> str | None:
|
||||
@@ -667,6 +698,8 @@ class MainWindow(QMainWindow):
|
||||
self._project = project
|
||||
project.frames_dir.mkdir(parents=True, exist_ok=True)
|
||||
project.apply_to_config(self._cfg) # per-project settings -> live config
|
||||
normalize_config(self._cfg) # coerce any legacy classic/inpaint values
|
||||
self._ensure_model() # YOLO needs weights — auto-pick if missing
|
||||
self._sync_settings_ui()
|
||||
self._detector_key = None
|
||||
self._restorer_key = None
|
||||
@@ -679,9 +712,6 @@ class MainWindow(QMainWindow):
|
||||
|
||||
def _sync_settings_ui(self) -> None:
|
||||
"""Reflect the (project's) config onto the toolbar widgets without signal loops."""
|
||||
self.detector_combo.blockSignals(True)
|
||||
self.detector_combo.setCurrentText(self._cfg.detector)
|
||||
self.detector_combo.blockSignals(False)
|
||||
self.threshold_spin.blockSignals(True)
|
||||
self.threshold_spin.setValue(self._cfg.default_threshold)
|
||||
self.threshold_spin.blockSignals(False)
|
||||
@@ -912,41 +942,32 @@ class MainWindow(QMainWindow):
|
||||
key, dets = payload
|
||||
self._results[key] = dets
|
||||
self._tag_file(Path(key), len(dets))
|
||||
# If the frame being viewed was just computed, show its overlay live.
|
||||
if not self._showing_restored and self._current is not None and str(self._current) == key:
|
||||
self._show(self._current)
|
||||
self._tick_count += 1
|
||||
if self._tick_count % 25 == 0:
|
||||
self._refresh_marks() # let marks appear progressively (throttled)
|
||||
|
||||
# ------------------------------------------------------------- restoration
|
||||
def _restore_current(self) -> None:
|
||||
"""Restore the current frame's regions on a background thread, then show it.
|
||||
"""Restore the current frame on a background thread, then show it.
|
||||
|
||||
Detections are computed first (in the same job) if not cached. The DeepMosaics
|
||||
engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
|
||||
DeepMosaics locates the mosaic itself, so no detection step is needed — we just
|
||||
run the engine on the frame (if there's no mosaic the frame comes back unchanged).
|
||||
The engine polls ``job.cancelled`` so "■ Стоп" stops it promptly."""
|
||||
if self._current is None or self._busy:
|
||||
return
|
||||
path = self._current
|
||||
key = str(path)
|
||||
|
||||
def fn(job):
|
||||
dets = self._results.get(key)
|
||||
if dets is None:
|
||||
dets = self._compute(self._make_detector(), path)
|
||||
job.tick(("dets", key, dets)) # cache them on the GUI thread
|
||||
if not dets:
|
||||
return ("empty", key)
|
||||
img = imread_unicode(key)
|
||||
if img is None:
|
||||
raise RuntimeError(f"Не удалось прочитать: {path.name}")
|
||||
restorer = self._make_restorer()
|
||||
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
|
||||
return ("restored", key, restored, len(dets), restorer.name)
|
||||
|
||||
def tick(payload):
|
||||
if payload[0] == "dets":
|
||||
_, k, dets = payload
|
||||
self._results[k] = dets
|
||||
self._tag_file(Path(k), len(dets))
|
||||
self._refresh_marks()
|
||||
restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled)
|
||||
return ("restored", key, restored, restorer.name)
|
||||
|
||||
def done(result, cancelled):
|
||||
if cancelled:
|
||||
@@ -954,19 +975,86 @@ class MainWindow(QMainWindow):
|
||||
return
|
||||
if result is None:
|
||||
return
|
||||
if result[0] == "empty":
|
||||
self.statusBar().showMessage("Нет найденных областей — нечего расцензуривать")
|
||||
return
|
||||
_, k, restored, n, engine = result
|
||||
_, k, restored, engine = result
|
||||
self._restored[k] = restored
|
||||
if self._current is not None and str(self._current) == k:
|
||||
self._showing_restored = True
|
||||
self.view.set_image(restored, [])
|
||||
self._update_restore_actions()
|
||||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name} — {n} обл.")
|
||||
self.statusBar().showMessage(f"Расцензурено ({engine}): {Path(k).name}")
|
||||
|
||||
self.statusBar().showMessage(f"Восстановление: {path.name}…")
|
||||
self._start_job(fn, None, on_tick=tick, on_done=done)
|
||||
self._start_job(fn, None, on_done=done)
|
||||
|
||||
def _restore_all(self, force: bool = False) -> None:
|
||||
"""Restore every frame on a background thread, writing results to ``restored/``.
|
||||
|
||||
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
|
||||
engine skips frames already restored (resume) unless ``force``. The temporal
|
||||
engine (DeepMosaics-video) runs the whole contiguous sequence in order via
|
||||
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
|
||||
"""
|
||||
if not self._files or self._project is None or self._busy:
|
||||
return
|
||||
files = list(self._files) # snapshot — favorites/move mutate self._files
|
||||
total = len(files)
|
||||
out_dir = self._project.restored_dir
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def out_path(p: Path) -> Path:
|
||||
return out_dir / f"{p.stem}.jpg"
|
||||
|
||||
def fn(job):
|
||||
restorer = self._make_restorer() # built on the worker (may raise)
|
||||
frame_cache: dict[int, object] = {} # small cache so the temporal window reuses reads
|
||||
|
||||
def get_frame(i):
|
||||
img = frame_cache.get(i)
|
||||
if img is None:
|
||||
img = imread_unicode(str(files[i]))
|
||||
if img is None:
|
||||
raise RuntimeError(f"Не удалось прочитать: {files[i].name}")
|
||||
if len(frame_cache) > 24:
|
||||
frame_cache.clear()
|
||||
frame_cache[i] = img
|
||||
return img
|
||||
|
||||
def emit(i, restored):
|
||||
imwrite_unicode(str(out_path(files[i])), restored)
|
||||
job.progress(i + 1, total, f"Расцензуривание {i + 1}/{total}: {files[i].name}")
|
||||
|
||||
if restorer.temporal:
|
||||
restorer.restore_sequence(
|
||||
total, get_frame, lambda _i: [], emit, should_cancel=lambda: job.cancelled
|
||||
)
|
||||
else:
|
||||
for i, p in enumerate(files):
|
||||
if job.cancelled:
|
||||
break
|
||||
if not force and out_path(p).is_file():
|
||||
job.progress(i + 1, total, f"Пропуск {i + 1}/{total}: {p.name}")
|
||||
continue
|
||||
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled))
|
||||
frame_cache.pop(i, None) # per-frame: don't accumulate
|
||||
return None
|
||||
|
||||
def done(_result, cancelled):
|
||||
if self._current is not None: # live-preview the current frame's result, if any
|
||||
rp = out_path(self._current)
|
||||
if rp.is_file():
|
||||
img = imread_unicode(str(rp))
|
||||
if img is not None:
|
||||
self._restored[str(self._current)] = img
|
||||
self._showing_restored = True
|
||||
self.view.set_image(img, [])
|
||||
self._update_restore_actions()
|
||||
self.statusBar().showMessage(
|
||||
"Расцензуривание отменено" if cancelled
|
||||
else f"Готово: результаты в {out_dir.name}/ ({total} кадров)"
|
||||
)
|
||||
|
||||
self.statusBar().showMessage("Пакетное расцензуривание…")
|
||||
self._start_job(fn, total, on_done=done)
|
||||
|
||||
def _make_restorer(self):
|
||||
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model,
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
"""Configure the restoration ("расцензурить") engine.
|
||||
|
||||
inpaint — no setup. deepmosaics — the network code is vendored (built-in); the user
|
||||
picks a clean model from a dropdown of the bundled ``models/deepmosaics`` weights
|
||||
(or browses to another ``clean_*.pth``). ``mosaic_position.pth`` must sit beside the
|
||||
chosen model. A CUDA GPU is strongly recommended (GPU id, -1 = CPU/slow).
|
||||
Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео")
|
||||
engine. The network code is vendored (built-in); the user picks a clean model from a
|
||||
dropdown of the bundled ``models/deepmosaics`` weights (or browses to another
|
||||
``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered.
|
||||
``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly
|
||||
recommended (GPU id, -1 = CPU/slow).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -35,9 +37,9 @@ class RestoreDialog(QDialog):
|
||||
self.setMinimumWidth(560)
|
||||
|
||||
self.engine = QComboBox()
|
||||
self.engine.addItem("Инпейнт (быстро, замазывает — без модели)", "inpaint")
|
||||
self.engine.addItem("DeepMosaics (реальное расцензуривание, нужна модель+GPU)", "deepmosaics")
|
||||
self.engine.setCurrentIndex(1 if config.restorer == "deepmosaics" else 0)
|
||||
self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics")
|
||||
self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video")
|
||||
self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer)))
|
||||
self.engine.currentIndexChanged.connect(self._sync)
|
||||
|
||||
# Model dropdown — bundled clean models, plus the configured one if external.
|
||||
@@ -52,9 +54,11 @@ class RestoreDialog(QDialog):
|
||||
form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model))
|
||||
form.addRow("GPU id:", self.dm_gpu)
|
||||
hint = QLabel(
|
||||
"Модели берутся из models/deepmosaics. Нужны clean_youknow_resnet_9blocks.pth "
|
||||
"и mosaic_position.pth (рядом). Видеомодель clean_*_video.pth покадрово не "
|
||||
"работает и в списке не показывается. На CPU медленно — лучше GPU. См. README."
|
||||
"Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n"
|
||||
"• Картинка: clean_youknow_resnet_9blocks.pth — покадрово.\n"
|
||||
"• Видео: clean_youknow_video.pth — использует соседние кадры (когерентнее на "
|
||||
"роликах), требует прогона по диапазону («Расцензурить все»).\n"
|
||||
"На CPU медленно — лучше GPU. См. README."
|
||||
)
|
||||
hint.setWordWrap(True)
|
||||
form.addRow(hint)
|
||||
@@ -67,7 +71,12 @@ class RestoreDialog(QDialog):
|
||||
|
||||
def _populate_models(self, current: str | None) -> None:
|
||||
self.model_combo.clear()
|
||||
for name, path in discover_models():
|
||||
is_video = self.engine.currentData() == "deepmosaics_video"
|
||||
if is_video: # temporal engine: only the video weights (clean_*_video.pth)
|
||||
models = [(n, p) for n, p in discover_models(include_video=True) if "video" in n.lower()]
|
||||
else:
|
||||
models = discover_models() # per-frame engine: image clean models only
|
||||
for name, path in models:
|
||||
self.model_combo.addItem(name, path)
|
||||
# Keep an externally-configured model selectable even if it's outside the folder.
|
||||
if current and self.model_combo.findData(current) < 0:
|
||||
@@ -88,7 +97,9 @@ class RestoreDialog(QDialog):
|
||||
return w
|
||||
|
||||
def _sync(self) -> None:
|
||||
is_dm = self.engine.currentData() == "deepmosaics"
|
||||
is_dm = self.engine.currentData() in ("deepmosaics", "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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user