56 lines
2.0 KiB
Python
56 lines
2.0 KiB
Python
"""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
|