78 lines
3.7 KiB
Python
78 lines
3.7 KiB
Python
"""Application configuration and tunable defaults.
|
|
|
|
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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DetectionConfig:
|
|
"""Parameters for the YOLO detector."""
|
|
|
|
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")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OverlayConfig:
|
|
"""How detections are drawn over the image."""
|
|
|
|
# RGB per detection category/label (the models/yolo/<category> folder name, or the
|
|
# CensorType for legacy detections). Unknown categories get a stable palette colour.
|
|
colors: dict[str, tuple[int, int, int]] = field(
|
|
default_factory=lambda: {
|
|
"mosaic": (231, 76, 60), # red
|
|
"blur": (241, 196, 15), # yellow
|
|
"black_bar": (26, 188, 156), # teal
|
|
"unknown": (155, 89, 182), # purple
|
|
"face": (46, 204, 113), # green
|
|
"hand": (52, 152, 219), # blue
|
|
"person": (230, 126, 34), # orange
|
|
"eyes": (155, 89, 182), # purple
|
|
"text": (149, 165, 166), # grey
|
|
}
|
|
)
|
|
# Fallback colours cycled (deterministically) for categories not listed above.
|
|
palette: tuple[tuple[int, int, int], ...] = (
|
|
(231, 76, 60), (46, 204, 113), (52, 152, 219), (241, 196, 15),
|
|
(155, 89, 182), (26, 188, 156), (230, 126, 34), (149, 165, 166),
|
|
)
|
|
line_width: int = 2
|
|
fill_alpha: int = 48 # 0..255 translucency of the region fill
|
|
show_labels: bool = True
|
|
|
|
|
|
@dataclass
|
|
class AppConfig:
|
|
detection: DetectionConfig = field(default_factory=DetectionConfig)
|
|
overlay: OverlayConfig = field(default_factory=OverlayConfig)
|
|
detector: str = "yolo" # only "yolo"
|
|
# Active YOLO models (paths under models/yolo/<category>/). A detect runs every
|
|
# selected model and merges results — see core/detection/multi.MultiYoloDetector.
|
|
detector_models: list[str] = field(default_factory=list)
|
|
default_threshold: float = 0.20 # initial overlay confidence threshold
|
|
# Per-model overlay threshold overrides, keyed by the model file **stem** (e.g.
|
|
# "penis" -> 0.4). A detection from a model with no entry uses default_threshold.
|
|
# Display-only (filters what's drawn/counted as a hit), not the detection conf.
|
|
model_thresholds: dict[str, float] = field(default_factory=dict)
|
|
# Merge overlapping detections across models (greedy IoU NMS, keep higher score).
|
|
# Off by default — different categories are meant to coexist; on, it removes the
|
|
# duplicate boxes you get when overlapping models (e.g. penis + cockAndBall) fire.
|
|
cross_model_nms: bool = False
|
|
nms_iou: float = 0.6 # IoU above which two boxes are deemed duplicates
|
|
|
|
# --- restoration ("расцензурить") ---
|
|
restorer: str = "deepmosaics" # "deepmosaics" | "deepmosaics_video"
|
|
dm_dir: str | None = None # optional extra dir to search for mosaic_position.pth
|
|
dm_model: str | None = None # DeepMosaics clean weights (clean_*.pth)
|
|
dm_gpu: str = "0" # CUDA device id, "-1" for CPU
|
|
# Temporal engine only: feed already-restored PAST frames into the BVDNet window
|
|
# (instead of the original mosaic frames) for stronger temporal coherence. Slightly
|
|
# out-of-distribution for the net (trained on mosaic windows) — toggle in the dialog.
|
|
dm_feed_restored: bool = True
|