217 lines
9.3 KiB
Python
217 lines
9.3 KiB
Python
"""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
|