Добавлено описание и документация для HVideoTool, включая функционал, требования, установку и запуск приложения для обнаружения цензуры на изображениях.
This commit is contained in:
Leonid Pershin
@@ -0,0 +1,26 @@
|
||||
"""Detector interface. Implement this to plug in a new model (e.g. YOLO)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ..video.frame import Frame
|
||||
from .types import Detection
|
||||
|
||||
|
||||
class Detector(ABC):
|
||||
"""Abstract censorship detector.
|
||||
|
||||
Implementations must be safe to call repeatedly on consecutive frames. They
|
||||
receive a :class:`Frame` and return detections in *source-frame* pixel
|
||||
coordinates (the same resolution as ``frame.image``).
|
||||
"""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return type(self).__name__
|
||||
|
||||
@abstractmethod
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
"""Return detected censored regions for ``frame`` (possibly empty)."""
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,216 @@
|
||||
"""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
|
||||
@@ -0,0 +1,51 @@
|
||||
"""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
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Detector factory: build a Detector from AppConfig.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
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:
|
||||
if not config.model_path:
|
||||
raise ValueError(
|
||||
"Для детектора YOLO укажите путь к весам (.pt) в Параметрах "
|
||||
"или скачайте модель LADA — см. README."
|
||||
)
|
||||
return config.model_path
|
||||
|
||||
|
||||
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
|
||||
|
||||
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}")
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Detection result types shared across detectors and the UI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CensorType(str, Enum):
|
||||
"""Kind of already-applied censorship a detection represents."""
|
||||
|
||||
MOSAIC = "mosaic"
|
||||
BLUR = "blur"
|
||||
BLACK_BAR = "black_bar"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Detection:
|
||||
"""A single detected censored region, in source-frame pixel coordinates."""
|
||||
|
||||
type: CensorType
|
||||
score: float # confidence, 0..1
|
||||
bbox: tuple[int, int, int, int] # x, y, w, h
|
||||
polygon: list[tuple[int, int]] = field(default_factory=list) # contour points
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"type": self.type.value,
|
||||
"score": self.score,
|
||||
"bbox": list(self.bbox),
|
||||
"polygon": [list(p) for p in self.polygon],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "Detection":
|
||||
return cls(
|
||||
type=CensorType(data["type"]),
|
||||
score=float(data["score"]),
|
||||
bbox=tuple(data["bbox"]), # type: ignore[arg-type]
|
||||
polygon=[tuple(p) for p in data.get("polygon", [])],
|
||||
)
|
||||
@@ -0,0 +1,104 @@
|
||||
"""Ultralytics YOLO detector.
|
||||
|
||||
Wraps an Ultralytics YOLO model (detection or segmentation) behind the
|
||||
:class:`Detector` interface. Designed for the LADA mosaic-detection weights
|
||||
(https://huggingface.co/ladaapp/lada), which are YOLO *segmentation* models with
|
||||
a single ``mosaic`` class — but it works with any Ultralytics ``.pt`` whose class
|
||||
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.
|
||||
|
||||
Licensing: Ultralytics YOLO and the LADA weights are AGPL-3.0. See README.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ...config import DetectionConfig
|
||||
from ..video.frame import Frame
|
||||
from .base import Detector
|
||||
from .types import CensorType, Detection
|
||||
|
||||
|
||||
def _name_to_type(name: str) -> CensorType:
|
||||
n = name.lower()
|
||||
if "mosaic" in n or "pixel" in n:
|
||||
return CensorType.MOSAIC
|
||||
if "blur" in n:
|
||||
return CensorType.BLUR
|
||||
if "bar" in n or "black" in n:
|
||||
return CensorType.BLACK_BAR
|
||||
return CensorType.UNKNOWN
|
||||
|
||||
|
||||
class YoloDetector(Detector):
|
||||
def __init__(self, model_path: str, config: DetectionConfig | None = None) -> None:
|
||||
self.cfg = config or DetectionConfig()
|
||||
if not os.path.isfile(model_path):
|
||||
raise FileNotFoundError(
|
||||
f"Файл весов не найден: {model_path}\n"
|
||||
"Скачайте модель детекции мозаики LADA, например:\n"
|
||||
" curl.exe -L -o models\\lada_mosaic_detection_model_v4_accurate.pt "
|
||||
'"https://huggingface.co/ladaapp/lada/resolve/main/'
|
||||
'lada_mosaic_detection_model_v4_accurate.pt?download=true"'
|
||||
)
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
except ImportError as exc: # pragma: no cover - environment dependent
|
||||
raise ImportError(
|
||||
"Не установлен ultralytics. Установите: pip install -e \".[yolo]\" "
|
||||
"(и PyTorch с CUDA отдельно — см. README)."
|
||||
) from exc
|
||||
|
||||
# Resolve the device: explicit override, else CUDA when available.
|
||||
device = self.cfg.yolo_device
|
||||
if device is None:
|
||||
try:
|
||||
import torch
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
except Exception: # noqa: BLE001
|
||||
device = "cpu"
|
||||
self._device = device
|
||||
self._model = YOLO(model_path)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return f"YoloDetector(device={self._device})"
|
||||
|
||||
def detect(self, frame: Frame) -> list[Detection]:
|
||||
results = self._model.predict(
|
||||
source=frame.image, # BGR ndarray; ultralytics handles it
|
||||
conf=self.cfg.yolo_conf,
|
||||
imgsz=self.cfg.yolo_imgsz,
|
||||
device=self._device,
|
||||
verbose=False,
|
||||
)
|
||||
if not results:
|
||||
return []
|
||||
res = results[0]
|
||||
boxes = getattr(res, "boxes", None)
|
||||
if boxes is None or len(boxes) == 0:
|
||||
return []
|
||||
|
||||
names = res.names # {class_index: class_name}
|
||||
xyxy = boxes.xyxy.cpu().numpy()
|
||||
confs = boxes.conf.cpu().numpy()
|
||||
classes = boxes.cls.cpu().numpy().astype(int)
|
||||
# Segmentation polygons in source-pixel coords, one per detection (if any).
|
||||
polygons = res.masks.xy if getattr(res, "masks", None) is not None else None
|
||||
|
||||
out: list[Detection] = []
|
||||
for i in range(len(xyxy)):
|
||||
x1, y1, x2, y2 = xyxy[i]
|
||||
bbox = (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
|
||||
poly: list[tuple[int, int]] = []
|
||||
if polygons is not None and i < len(polygons):
|
||||
poly = [(int(px), int(py)) for px, py in polygons[i]]
|
||||
ctype = _name_to_type(names.get(int(classes[i]), ""))
|
||||
out.append(Detection(type=ctype, score=float(confs[i]), bbox=bbox, polygon=poly))
|
||||
return out
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Unicode-safe image read/write.
|
||||
|
||||
``cv2.imread``/``cv2.imwrite`` mishandle non-ASCII paths on Windows. These
|
||||
helpers go through ``np.fromfile``/``ndarray.tofile`` + ``imdecode``/``imencode``
|
||||
so paths with Cyrillic (etc.) work regardless of the system locale.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def imread_unicode(path: str) -> np.ndarray | None:
|
||||
try:
|
||||
data = np.fromfile(path, dtype=np.uint8)
|
||||
except OSError:
|
||||
return None
|
||||
if data.size == 0:
|
||||
return None
|
||||
return cv2.imdecode(data, cv2.IMREAD_COLOR)
|
||||
|
||||
|
||||
def imwrite_unicode(path: str, image: np.ndarray, params: list[int] | None = None) -> bool:
|
||||
ext = os.path.splitext(path)[1] or ".jpg"
|
||||
ok, buf = cv2.imencode(ext, image, params or [])
|
||||
if not ok:
|
||||
return False
|
||||
buf.tofile(path)
|
||||
return True
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Extract frames from a video into a folder of JPGs.
|
||||
|
||||
Two engines:
|
||||
|
||||
* **ffmpeg** (preferred, used when the ``ffmpeg`` binary is on PATH) — one
|
||||
subprocess does decode + sampling + optional downscale + JPEG encode, which is
|
||||
faster than pulling frames into Python one by one, and unlocks the big win:
|
||||
*keyframe-only* extraction (``-skip_frame nokey`` decodes only I-frames, ~10×
|
||||
faster than decoding every frame).
|
||||
* **OpenCV** fallback (``cv2.VideoCapture``) when ffmpeg is absent.
|
||||
|
||||
Decoding H.264/HEVC frame-by-frame is inherently the cost; hardware accel doesn't
|
||||
help for this (GPU transfer overhead). The only way to be dramatically faster is
|
||||
to decode fewer frames — hence the keyframe mode.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import cv2
|
||||
|
||||
from ..imageio import imwrite_unicode
|
||||
|
||||
# progress(done_seconds, total_seconds) -> return False to cancel.
|
||||
Progress = Callable[[float, float], bool]
|
||||
|
||||
|
||||
def _find_ffmpeg() -> str | None:
|
||||
"""ffmpeg on PATH, else the binary bundled with imageio-ffmpeg, else None."""
|
||||
exe = shutil.which("ffmpeg")
|
||||
if exe:
|
||||
return exe
|
||||
try:
|
||||
import imageio_ffmpeg
|
||||
|
||||
return imageio_ffmpeg.get_ffmpeg_exe()
|
||||
except Exception: # noqa: BLE001 - package missing or no bundled binary
|
||||
return None
|
||||
|
||||
|
||||
def _video_duration_seconds(video_path: str) -> float:
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
try:
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
|
||||
frames = cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0
|
||||
return frames / fps if fps > 0 else 0.0
|
||||
finally:
|
||||
cap.release()
|
||||
|
||||
|
||||
def _quality_to_qscale(jpg_quality: int) -> int:
|
||||
"""Map JPEG quality 0..100 to ffmpeg -q:v (2 best .. 31 worst)."""
|
||||
q = round(2 + (100 - max(0, min(100, jpg_quality))) / 100 * 29)
|
||||
return max(2, min(31, q))
|
||||
|
||||
|
||||
def extract_frames(
|
||||
video_path: str,
|
||||
out_dir: str,
|
||||
step: int = 15,
|
||||
keyframes_only: bool = False,
|
||||
max_dim: int = 0,
|
||||
jpg_quality: int = 92,
|
||||
progress: Progress | None = None,
|
||||
) -> int:
|
||||
"""Save sampled frames of ``video_path`` into ``out_dir`` as JPGs.
|
||||
|
||||
``keyframes_only`` decodes only keyframes (fast). Otherwise keeps every
|
||||
``step``-th frame. ``max_dim`` (>0) caps the longest side. ``progress`` is
|
||||
called with (done_seconds, total_seconds); returning ``False`` cancels.
|
||||
Returns the number of frames written.
|
||||
"""
|
||||
out = Path(out_dir)
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
ffmpeg = _find_ffmpeg()
|
||||
if ffmpeg:
|
||||
return _extract_ffmpeg(ffmpeg, video_path, out, step, keyframes_only, max_dim, jpg_quality, progress)
|
||||
return _extract_cv2(video_path, out, step, max_dim, jpg_quality, progress)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- ffmpeg
|
||||
def _build_vf(step: int, keyframes_only: bool, max_dim: int) -> str | None:
|
||||
filters: list[str] = []
|
||||
if not keyframes_only and step > 1:
|
||||
filters.append(f"select=not(mod(n\\,{int(step)}))")
|
||||
if max_dim and max_dim > 0:
|
||||
# Cap the longest side to max_dim, preserve aspect, never upscale.
|
||||
filters.append(f"scale='min({max_dim},iw)':'min({max_dim},ih)':force_original_aspect_ratio=decrease")
|
||||
return ",".join(filters) if filters else None
|
||||
|
||||
|
||||
def _extract_ffmpeg(
|
||||
ffmpeg: str, video_path: str, out: Path, step: int,
|
||||
keyframes_only: bool, max_dim: int, jpg_quality: int, progress: Progress | None,
|
||||
) -> int:
|
||||
total = _video_duration_seconds(video_path)
|
||||
cmd = [ffmpeg, "-hide_banner", "-loglevel", "error", "-nostdin"]
|
||||
if keyframes_only:
|
||||
cmd += ["-skip_frame", "nokey"] # input option: decode only keyframes
|
||||
cmd += ["-i", video_path]
|
||||
vf = _build_vf(step, keyframes_only, max_dim)
|
||||
if vf:
|
||||
cmd += ["-vf", vf]
|
||||
cmd += ["-vsync", "0", "-q:v", str(_quality_to_qscale(jpg_quality)),
|
||||
"-progress", "pipe:1", str(out / "%06d.jpg")]
|
||||
|
||||
proc = subprocess.Popen(
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
|
||||
stdin=subprocess.DEVNULL, text=True, bufsize=1,
|
||||
)
|
||||
try:
|
||||
assert proc.stdout is not None
|
||||
for line in proc.stdout:
|
||||
if progress is None:
|
||||
continue
|
||||
line = line.strip()
|
||||
if line.startswith("out_time_us=") or line.startswith("out_time_ms="):
|
||||
raw = line.split("=", 1)[1]
|
||||
try:
|
||||
# out_time_us is microseconds; out_time_ms is *also* microseconds
|
||||
# in ffmpeg (historical misnomer). Both -> seconds via /1e6.
|
||||
done = int(raw) / 1_000_000 if raw.isdigit() else 0.0
|
||||
except ValueError:
|
||||
done = 0.0
|
||||
if progress(done, total) is False:
|
||||
proc.terminate()
|
||||
break
|
||||
finally:
|
||||
proc.wait()
|
||||
return len(list(out.glob("*.jpg")))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------- opencv
|
||||
def _extract_cv2(
|
||||
video_path: str, out: Path, step: int, max_dim: int, jpg_quality: int, progress: Progress | None,
|
||||
) -> int:
|
||||
cap = cv2.VideoCapture(str(video_path))
|
||||
if not cap.isOpened():
|
||||
raise RuntimeError(f"Не удалось открыть видео: {video_path}")
|
||||
step = max(1, int(step))
|
||||
fps = cap.get(cv2.CAP_PROP_FPS) or 0.0
|
||||
total = (cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0.0) / fps if fps > 0 else 0.0
|
||||
params = [cv2.IMWRITE_JPEG_QUALITY, int(jpg_quality)]
|
||||
idx = saved = 0
|
||||
try:
|
||||
while True:
|
||||
if not cap.grab():
|
||||
break
|
||||
if idx % step == 0:
|
||||
ok, frame = cap.retrieve()
|
||||
if ok:
|
||||
if max_dim and max(frame.shape[:2]) > max_dim:
|
||||
s = max_dim / max(frame.shape[:2])
|
||||
frame = cv2.resize(frame, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
|
||||
imwrite_unicode(str(out / f"{idx:06d}.jpg"), frame, params)
|
||||
saved += 1
|
||||
idx += 1
|
||||
if progress is not None and idx % 30 == 0:
|
||||
done = idx / fps if fps > 0 else 0.0
|
||||
if progress(done, total) is False:
|
||||
break
|
||||
finally:
|
||||
cap.release()
|
||||
return saved
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Decoded video frame passed from the reader to the detector."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class Frame:
|
||||
image: np.ndarray # BGR, HxWx3, uint8 (OpenCV convention)
|
||||
index: int # 0-based frame counter since the last open/seek
|
||||
pts: float # presentation timestamp, seconds
|
||||
Reference in New Issue
Block a user