feat(vision): add detection local zoo CLI

jhlu · jhlu · commit d7e7ca09bd14 · 2026-03-05T08:57:05.000+08:00
diff --git a/dlhub/vision/detection_zoo.py b/dlhub/vision/detection_zoo.py
@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+
+from torch import nn
+
+
+@dataclass(frozen=True)
+class BuildConfig:
+    in_channels: int
+    num_classes: int
+    width_mult: float = 1.0
+
+
+class UnknownLocalArch(KeyError):
+    pass
+
+
+def _split_arch_id(arch_id: str) -> tuple[str, str]:
+    arch_id = str(arch_id).strip()
+    if ":" not in arch_id:
+        return "dldet", arch_id
+
+    prefix, name = arch_id.split(":", 1)
+    prefix = prefix.strip().lower()
+    name = name.strip()
+    if not prefix or not name:
+        raise ValueError(f"Invalid arch id: {arch_id!r}")
+    return prefix, name
+
+
+Builder = Callable[[BuildConfig], nn.Module]
+
+
+def _extract_variants_from_source(src: str) -> list[str] | None:
+    """Extract `_VARIANTS` keys from a module source without importing it."""
+
+    import ast
+
+    try:
+        tree = ast.parse(src)
+    except SyntaxError:
+        return None
+
+    for node in tree.body:
+        target_name: str | None = None
+        value = None
+
+        if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
+            target_name = node.target.id
+            value = node.value
+        elif isinstance(node, ast.Assign):
+            for t in node.targets:
+                if isinstance(t, ast.Name):
+                    target_name = t.id
+                    break
+            value = node.value
+
+        if target_name != "_VARIANTS" or not isinstance(value, ast.Dict):
+            continue
+
+        keys: list[str] = []
+        for k in value.keys:
+            if isinstance(k, ast.Constant) and isinstance(k.value, str):
+                keys.append(k.value)
+        return keys or None
+
+    return None
+
+
+def _extract_builder_name_from_source(src: str) -> str | None:
+    """Extract the first `build_*_detector` function name without importing."""
+
+    import ast
+
+    try:
+        tree = ast.parse(src)
+    except SyntaxError:
+        return None
+
+    for node in tree.body:
+        if not isinstance(node, ast.FunctionDef):
+            continue
+        name = str(node.name)
+        if name.startswith("build_") and name.endswith("_detector"):
+            return name
+    return None
+
+
+def _make_lazy_detector_builder(module_name: str, *, builder_name: str, variant: str) -> Builder:
+    module_name = str(module_name).strip()
+    builder_name = str(builder_name).strip()
+    variant = str(variant).strip()
+
+    def _builder(cfg: BuildConfig) -> nn.Module:
+        import importlib
+        import inspect
+
+        mod = importlib.import_module(f"dlhub.vision.detection.{module_name}")
+        fn = getattr(mod, builder_name, None)
+        if fn is None:
+            raise RuntimeError(f"Detection module {module_name!r} missing {builder_name}()")
+
+        kwargs: dict[str, object] = {
+            "in_channels": int(cfg.in_channels),
+            "num_classes": int(cfg.num_classes),
+            "variant": str(variant),
+            "width_mult": float(cfg.width_mult),
+        }
+
+        try:
+            sig = inspect.signature(fn)
+        except (TypeError, ValueError):
+            sig = None
+
+        if sig is not None:
+            params = set(sig.parameters)
+            kwargs = {k: v for k, v in kwargs.items() if k in params}
+
+        return fn(**kwargs)
+
+    return _builder
+
+
+def _extend_registry_with_discovered_detectors(r: dict[str, Builder]) -> None:
+    """Discover detector variants under `dlhub/vision/detection/*.py`."""
+
+    here = Path(__file__).resolve().parent
+    det_dir = here / "detection"
+
+    if not det_dir.exists():
+        return
+
+    hidden = {"__init__"}
+
+    for py in sorted(det_dir.glob("*.py")):
+        module_name = py.stem
+        if module_name in hidden or module_name.startswith("_"):
+            continue
+
+        try:
+            src = py.read_text(encoding="utf-8")
+        except OSError:
+            continue
+
+        if "_VARIANTS" not in src or "def build_" not in src:
+            continue
+
+        variants = _extract_variants_from_source(src)
+        if not variants:
+            continue
+
+        builder_name = _extract_builder_name_from_source(src)
+        if builder_name is None:
+            continue
+
+        for v in variants:
+            name = str(v).lower().strip()
+            if not name or name in r:
+                continue
+            r[name] = _make_lazy_detector_builder(module_name, builder_name=builder_name, variant=name)
+
+
+def _registry() -> dict[str, Builder]:
+    r: dict[str, Builder] = {}
+    _extend_registry_with_discovered_detectors(r)
+    return r
+
+
+_REGISTRY = _registry()
+
+
+def list_local_arches() -> list[str]:
+    """List all available local detection architecture ids (e.g. `dldet:ssd_tiny`)."""
+
+    return [f"dldet:{name}" for name in sorted(_REGISTRY)]
+
+
+def build_local_model(
+    arch_id: str,
+    *,
+    in_channels: int,
+    num_classes: int,
+    width_mult: float = 1.0,
+) -> nn.Module:
+    """Build a local detection model by architecture id.
+
+    Architecture ids are variants extracted from each detector module's `_VARIANTS`,
+    namespaced with `dldet:` (e.g. `dldet:ssd_tiny`).
+    """
+
+    prefix, name = _split_arch_id(arch_id)
+    if prefix == "det":
+        prefix = "dldet"
+    if prefix not in {"dldet", "local"}:
+        raise ValueError(f"Unsupported detection prefix: {prefix!r} (arch_id={arch_id!r})")
+
+    builder = _REGISTRY.get(str(name).lower().strip())
+    if builder is None:
+        raise UnknownLocalArch(f"Unknown detection arch: {arch_id!r}. Tip: run `python scripts/detection_zoo.py --list`.")
+
+    return builder(
+        BuildConfig(
+            in_channels=int(in_channels),
+            num_classes=int(num_classes),
+            width_mult=float(width_mult),
+        )
+    )
+
+
+__all__ = [
+    "BuildConfig",
+    "UnknownLocalArch",
+    "build_local_model",
+    "list_local_arches",
+]
+
diff --git a/scripts/detection_zoo.py b/scripts/detection_zoo.py
@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from collections.abc import Iterable
+from pathlib import Path
+
+
+def _ensure_repo_root_on_path() -> None:
+    repo_root = Path(__file__).resolve().parents[1]
+    sys.path.insert(0, str(repo_root))
+
+
+def _summarize(obj) -> str:
+    try:
+        import torch
+    except Exception:
+        return f"{type(obj).__name__}"
+
+    if isinstance(obj, torch.Tensor):
+        return f"Tensor(shape={tuple(obj.shape)}, dtype={obj.dtype}, device={obj.device})"
+    if isinstance(obj, dict):
+        keys = ", ".join(sorted(map(str, obj.keys())))
+        return f"dict(keys=[{keys}])"
+    if isinstance(obj, list | tuple):
+        head = ", ".join(_summarize(x) for x in obj[:2])
+        tail = "" if len(obj) <= 2 else f", ... (+{len(obj) - 2})"
+        return f"{type(obj).__name__}([{head}{tail}])"
+    return f"{type(obj).__name__}"
+
+
+def _print_lines(lines: Iterable[str], *, limit: int = 80) -> None:
+    lines = list(lines)
+    if len(lines) <= limit:
+        for line in lines:
+            print(line)
+        return
+
+    head = max(10, limit - 10)
+    for line in lines[:head]:
+        print(line)
+    print(f"... ({len(lines) - limit} more) ...")
+    for line in lines[-10:]:
+        print(line)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Detection local model zoo utilities (no downloads).")
+
+    parser.add_argument("--list", action="store_true", help="List available architecture ids.")
+    parser.add_argument("--search", type=str, default=None, help="Filter list by substring (case-insensitive).")
+    parser.add_argument("--limit", type=int, default=80, help="Max lines to print when listing.")
+
+    parser.add_argument("--smoke", type=str, default=None, metavar="ARCH_ID", help="Run a forward smoke on an arch id.")
+    parser.add_argument("--batch-size", type=int, default=2, help="Batch size for smoke inputs.")
+    parser.add_argument("--image-size", type=int, default=64, help="Image size for smoke inputs.")
+    parser.add_argument("--in-channels", type=int, default=3, help="Input channels for local detectors.")
+    parser.add_argument("--num-classes", type=int, default=2, help="Detector classes.")
+    parser.add_argument("--width-mult", type=float, default=1.0, help="Width multiplier for local detectors.")
+
+    return parser.parse_args()
+
+
+def main() -> int:
+    _ensure_repo_root_on_path()
+
+    from dlhub.vision.detection_zoo import build_local_model, list_local_arches
+
+    args = parse_args()
+
+    if not args.list and args.smoke is None:
+        print("Nothing to do. Try one of:")
+        print("- python scripts/detection_zoo.py --list")
+        print("- python scripts/detection_zoo.py --smoke dldet:ssd_tiny")
+        return 2
+
+    arches = list_local_arches()
+    if args.search:
+        needle = str(args.search).lower()
+        arches = [a for a in arches if needle in a.lower()]
+
+    if args.list:
+        print("Detection local zoo")
+        print(f"- total_arches={len(arches)}")
+        print("")
+        _print_lines(arches, limit=int(args.limit))
+
+    if args.smoke is not None:
+        arch_id = str(args.smoke).strip()
+        if ":" not in arch_id:
+            arch_id = f"dldet:{arch_id}"
+
+        import torch
+
+        x = torch.randn(int(args.batch_size), int(args.in_channels), int(args.image_size), int(args.image_size))
+        model = build_local_model(
+            arch_id,
+            in_channels=int(args.in_channels),
+            num_classes=int(args.num_classes),
+            width_mult=float(args.width_mult),
+        )
+        model.eval()
+        with torch.no_grad():
+            out = model(x)
+
+        print("")
+        print(f"smoke: {arch_id}")
+        print(f"- output={_summarize(out)}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+
diff --git a/tests/test_dlhub_vision_detection_zoo.py b/tests/test_dlhub_vision_detection_zoo.py
@@ -0,0 +1,31 @@
+import pytest
+
+
+torch = pytest.importorskip("torch")
+
+
+def _sum_tensor_means(x):
+    if torch.is_tensor(x):
+        return x.to(torch.float32).mean()
+    if isinstance(x, dict):
+        return sum((_sum_tensor_means(v) for v in x.values()), start=torch.tensor(0.0))
+    if isinstance(x, (list, tuple)):
+        return sum((_sum_tensor_means(v) for v in x), start=torch.tensor(0.0))
+    raise TypeError(f"Unsupported output type in detection zoo smoke: {type(x)!r}")
+
+
+def test_detection_zoo_list_and_build_smoke() -> None:
+    from dlhub.vision.detection_zoo import build_local_model, list_local_arches
+
+    arches = list_local_arches()
+    assert "dldet:ssd_tiny" in arches
+    assert "dldet:detr_tiny" in arches
+
+    for arch_id in ["dldet:ssd_tiny", "dldet:detr_tiny", "dldet:yolo_v1_tiny"]:
+        model = build_local_model(arch_id, in_channels=3, num_classes=2, width_mult=0.5)
+        x = torch.randn(2, 3, 64, 64)
+        out = model(x)
+        loss = _sum_tensor_means(out)
+        assert torch.isfinite(loss)
+        loss.backward()
+
