Implement multi-model detection in HVideoTool: updated the detection system to support multiple YOLO models simultaneously, enhancing detection capabilities. Reflected changes in the UI with a new model selection menu and updated documentation in README and CLAUDE.md to guide users on model management and configuration.
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
"""Discover the YOLO models available for detection.
|
||||
|
||||
ADetailer-style layout: drop weights under ``models/yolo/<category>/*.pt``. The
|
||||
*category* (the sub-folder) becomes the detection label and its overlay colour, so
|
||||
e.g. ``models/yolo/mosaic/lada.pt`` tags its boxes "mosaic" and ``models/yolo/face/
|
||||
yolov8n-face.pt`` tags "face". The user ticks which discovered models are active; a
|
||||
detect runs every ticked model and merges the results (see ``multi.MultiYoloDetector``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
YOLO_SUBDIR = ("models", "yolo") # relative to the working directory
|
||||
_UNCATEGORIZED = "misc" # category for a .pt sitting directly under models/yolo
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelEntry:
|
||||
path: str # absolute path to the .pt
|
||||
category: str # sub-folder under models/yolo (the detection label / colour key)
|
||||
name: str # filename stem (display name)
|
||||
|
||||
|
||||
def yolo_root(root: Path | None = None) -> Path:
|
||||
return (root or Path.cwd()).joinpath(*YOLO_SUBDIR)
|
||||
|
||||
|
||||
def _category_of(pt: Path, root: Path) -> str:
|
||||
rel = pt.relative_to(root).parts
|
||||
return rel[0] if len(rel) > 1 else _UNCATEGORIZED
|
||||
|
||||
|
||||
def discover_models(root: Path | None = None) -> list[ModelEntry]:
|
||||
"""All ``*.pt`` under ``models/yolo/**``, sorted by (category, name)."""
|
||||
base = yolo_root(root)
|
||||
if not base.is_dir():
|
||||
return []
|
||||
out = [
|
||||
ModelEntry(path=str(p), category=_category_of(p, base), name=p.stem)
|
||||
for p in base.rglob("*.pt")
|
||||
]
|
||||
out.sort(key=lambda e: (e.category.lower(), e.name.lower()))
|
||||
return out
|
||||
|
||||
|
||||
def category_of(model_path: str, root: Path | None = None) -> str:
|
||||
"""Category (label) for a model path, derived from its folder under models/yolo."""
|
||||
base = yolo_root(root)
|
||||
p = Path(model_path)
|
||||
try:
|
||||
return _category_of(p, base)
|
||||
except ValueError: # outside models/yolo — fall back to the parent folder name
|
||||
return p.parent.name or _UNCATEGORIZED
|
||||
Reference in New Issue
Block a user