src/utils/gpu_utils.py

"""
GPU / backend detection utilities for OpenCV 5 DNN.

OpenCV 5 engine selection strategy:
  ENGINE_AUTO    → new graph engine (CPU-optimised, >80% ONNX coverage)
  ENGINE_CLASSIC → legacy 4.x engine (needed for CUDA backend today)
  ENGINE_ORT     → delegate to ONNX Runtime (optional, enables TensorRT etc.)
"""

from __future__ import annotations

import cv2

from src.utils.logger import get_logger

log = get_logger(__name__)


def get_dnn_backend_target(backend_name: str = "cpu") -> tuple[int, int]:
    """
    Return the (backend, target) pair for cv2.dnn.Net.

    Args:
        backend_name: ``"cpu"`` | ``"cuda"`` | ``"opencl"``

    Returns:
        Tuple of (cv2.dnn.DNN_BACKEND_*, cv2.dnn.DNN_TARGET_*).
    """
    backend_name = backend_name.lower()

    if backend_name == "cuda":
        if _cuda_available():
            log.info("CUDA backend selected")
            return cv2.dnn.DNN_BACKEND_CUDA, cv2.dnn.DNN_TARGET_CUDA
        else:
            log.warning("CUDA requested but not available, falling back to CPU")

    if backend_name == "opencl":
        log.info("OpenCL backend selected")
        return cv2.dnn.DNN_BACKEND_DEFAULT, cv2.dnn.DNN_TARGET_OPENCL

    log.info("CPU backend selected (OpenCV 5 graph engine)")
    return cv2.dnn.DNN_BACKEND_DEFAULT, cv2.dnn.DNN_TARGET_CPU


def _cuda_available() -> bool:
    """Check whether OpenCV was compiled with CUDA support."""
    try:
        info = cv2.getBuildInformation()
        return "CUDA" in info and "YES" in info[info.index("CUDA"):]
    except Exception:
        return False


def get_engine_flag(engine_name: str = "auto") -> int | None:
    """
    Return the ENGINE_* constant for cv2.dnn.readNet() if OpenCV 5 supports it.
    Falls back to None (which means OpenCV 4-style call without engine param).

    Args:
        engine_name: ``"auto"`` | ``"new"`` | ``"classic"`` | ``"ort"``
    """
    engine_name = engine_name.lower()
    engine_map = {
        "auto": getattr(cv2.dnn, "ENGINE_AUTO", None),
        "new": getattr(cv2.dnn, "ENGINE_NEW", None),
        "classic": getattr(cv2.dnn, "ENGINE_CLASSIC", None),
        "ort": getattr(cv2.dnn, "ENGINE_ORT", None),
    }
    flag = engine_map.get(engine_name)
    if flag is None and engine_name != "auto":
        log.warning(
            "ENGINE_%s not available in this OpenCV build — using ENGINE_AUTO",
            engine_name.upper(),
        )
    return flag


def log_system_info() -> None:
    """Log OpenCV build information relevant to DNN inference."""
    log.info("OpenCV version: %s", cv2.__version__)
    build = cv2.getBuildInformation()
    for line in build.splitlines():
        if any(kw in line for kw in ("CUDA", "ONNX", "Inference", "GPU", "OpenCL")):
            log.debug("  %s", line.strip())