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
@@ -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.
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user