Behavision/behavision/model_assets.py

"""Model acquisition: download YuNet, copy reusable models from the old
projects on this machine when present. Idempotent — safe to re-run."""
from __future__ import annotations

import logging
import shutil
import urllib.request
from pathlib import Path

log = logging.getLogger(__name__)

YUNET_URL = ("https://github.com/opencv/opencv_zoo/raw/main/models/"
             "face_detection_yunet/face_detection_yunet_2023mar.onnx")
BUFFALO_SC_URL = ("https://github.com/deepinsight/insightface/releases/"
                  "download/v0.7/buffalo_sc.zip")
RECOGNIZERS = ["adaface_ir101.onnx", "adaface_ir50.onnx", "w600k_r50.onnx",
               "arcface_int8.onnx", "w600k_mbf.onnx", "arcface.onnx"]

# Known locations of reusable models from the previous projects.
_LEGACY_MODEL_DIRS = [
    Path(r"D:\NEARLE\WOrking now\RTSP_16072025\pattern_reg\models"),
]

BUFFALO_L_URL = ("https://github.com/deepinsight/insightface/releases/"
                 "download/v0.7/buffalo_l.zip")

# target filename -> legacy filename
_COPY_MAP = {
    "arcface.onnx": "arcface.onnx",
    "age_deploy.prototxt": "age_deploy.prototxt",
    "age_net.caffemodel": "age_net.caffemodel",
    "gender_deploy.prototxt": "gender_deploy.prototxt",
    "gender_net.caffemodel": "gender_net.caffemodel",
    "emotion-ferplus-8.onnx": "emotion-ferplus-8.onnx",
}


def setup_models(models_dir: Path) -> "list[str]":
    """Ensure all model files exist in models_dir. Returns missing ones."""
    models_dir = Path(models_dir)
    models_dir.mkdir(parents=True, exist_ok=True)

    yunet = models_dir / "face_detection_yunet_2023mar.onnx"
    if not yunet.exists():
        log.info("downloading YuNet face detector (~230 KB)...")
        tmp = yunet.with_suffix(".part")
        urllib.request.urlretrieve(YUNET_URL, tmp)
        tmp.rename(yunet)
        log.info("YuNet saved to %s", yunet)

    for target_name, legacy_name in _COPY_MAP.items():
        target = models_dir / target_name
        if target.exists():
            continue
        for legacy_dir in _LEGACY_MODEL_DIRS:
            src = legacy_dir / legacy_name
            if src.exists():
                log.info("copying %s from %s ...", legacy_name, legacy_dir)
                shutil.copy2(src, target)
                break

    # Any one recognizer is enough; get the lightweight MobileFaceNet if
    # none is present (13 MB, loads reliably on low-memory machines).
    if not any((models_dir / n).exists() for n in RECOGNIZERS):
        log.info("downloading MobileFaceNet recognizer (buffalo_sc, ~15 MB)...")
        import io
        import zipfile

        with urllib.request.urlopen(BUFFALO_SC_URL) as resp:
            payload = io.BytesIO(resp.read())
        with zipfile.ZipFile(payload) as zf, \
                zf.open("w600k_mbf.onnx") as src, \
                open(models_dir / "w600k_mbf.onnx", "wb") as dst:
            shutil.copyfileobj(src, dst)
        log.info("w600k_mbf.onnx saved")

    # buffalo_l carries both the modern gender+age net (1.3 MB) and the
    # ResNet50 recognizer (~166 MB, IJB-C 97.25 vs MobileFaceNet's 95.02).
    # One 275 MB download serves both, so fetch it once and take what is
    # missing. Optional: failure here must never block the pipeline.
    wanted = {name: models_dir / name
              for name in ("genderage.onnx", "w600k_r50.onnx")
              if not (models_dir / name).exists()}
    if wanted:
        try:
            log.info("downloading %s from the buffalo_l bundle (~275 MB "
                     "one-time download)...", ", ".join(wanted))
            import zipfile

            tmp = models_dir / "buffalo_l.zip.part"
            urllib.request.urlretrieve(BUFFALO_L_URL, tmp)
            with zipfile.ZipFile(tmp) as zf:
                for name, target in wanted.items():
                    member = next((n for n in zf.namelist()
                                   if n.endswith(name)), None)
                    if member is None:
                        log.warning("%s not found in bundle", name)
                        continue
                    part = target.with_suffix(".part")
                    with zf.open(member) as src, open(part, "wb") as dst:
                        shutil.copyfileobj(src, dst)
                    part.rename(target)  # never leave a half-written model
                    log.info("%s saved", name)
            tmp.unlink()
        except Exception:
            log.warning("buffalo_l download failed - falling back to the "
                        "models already present", exc_info=True)

    missing = []
    if not (models_dir / "face_detection_yunet_2023mar.onnx").exists():
        missing.append("face_detection_yunet_2023mar.onnx")
    if not any((models_dir / n).exists() for n in RECOGNIZERS):
        missing.append("a recognition model (any of: %s)" % ", ".join(RECOGNIZERS))
    optional_missing = [n for n in _COPY_MAP
                        if not (models_dir / n).exists() and n not in missing]
    if optional_missing:
        log.warning("optional attribute models missing (age/gender/emotion "
                    "will be skipped): %s", ", ".join(optional_missing))
    return missing