Spaces:
Sleeping
Sleeping
File size: 2,787 Bytes
65728b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
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())
|