139 lines
5.5 KiB
Python
139 lines
5.5 KiB
Python
"""
|
|
Face crop classifier.
|
|
|
|
Loads a MobileNetV3 model via ONNX/TensorRT FP32, applies temperature-scaled
|
|
softmax, and returns per-crop probabilities for the target class.
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any, List, Optional
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ImageNet normalisation constants
|
|
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
|
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
|
|
|
|
def calibrated_softmax(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
|
|
"""Softmax with temperature scaling; returns class probabilities."""
|
|
scaled = logits / max(temperature, 1e-8)
|
|
shifted = scaled - scaled.max(axis=-1, keepdims=True) # numerical stability
|
|
exp = np.exp(shifted)
|
|
return exp / exp.sum(axis=-1, keepdims=True)
|
|
|
|
|
|
class FaceClassifier:
|
|
"""Classify face crops and return calibrated target-class probabilities."""
|
|
|
|
def __init__(
|
|
self,
|
|
engine_path: str,
|
|
temperature: float = 1.0,
|
|
input_size: int = 224,
|
|
device: str = "cuda",
|
|
):
|
|
self.engine_path = engine_path
|
|
self.temperature = temperature
|
|
self.input_size = input_size
|
|
self.device = device
|
|
self._session: Optional[Any] = None
|
|
self._load_model()
|
|
|
|
def _load_model(self):
|
|
path = Path(self.engine_path)
|
|
if not path.exists():
|
|
logger.warning("Classifier engine not found at %s; using placeholder", self.engine_path)
|
|
return
|
|
|
|
suffix = path.suffix.lower()
|
|
if suffix == ".onnx":
|
|
try:
|
|
import onnxruntime as ort
|
|
providers = (["CUDAExecutionProvider"] if self.device.startswith("cuda")
|
|
else ["CPUExecutionProvider"])
|
|
self._session = ort.InferenceSession(str(path), providers=providers)
|
|
logger.info("Loaded ONNX classifier from %s", self.engine_path)
|
|
except Exception as exc:
|
|
logger.warning("Failed to load ONNX classifier: %s", exc)
|
|
elif suffix in (".trt", ".engine", ".plan"):
|
|
try:
|
|
import tensorrt as trt
|
|
with trt.Logger() as trt_logger, open(path, "rb") as f:
|
|
runtime = trt.Runtime(trt_logger)
|
|
self._session = runtime.deserialize_cuda_engine(f.read())
|
|
logger.info("Loaded TensorRT classifier from %s", self.engine_path)
|
|
except Exception as exc:
|
|
logger.warning("Failed to load TensorRT classifier: %s", exc)
|
|
else:
|
|
logger.warning("Unsupported classifier format: %s", suffix)
|
|
|
|
def classify(self, crop_paths: List[str], batch_size: int = 16) -> List[float]:
|
|
"""Return target-class probabilities (p ∈ [0,1]) for each crop path."""
|
|
results: List[float] = []
|
|
for i in range(0, len(crop_paths), batch_size):
|
|
batch = crop_paths[i : i + batch_size]
|
|
results.extend(self._classify_batch(batch))
|
|
return results
|
|
|
|
def _classify_batch(self, crop_paths: List[str]) -> List[float]:
|
|
preprocessed = []
|
|
for path in crop_paths:
|
|
try:
|
|
preprocessed.append(self._preprocess(path))
|
|
except Exception as exc:
|
|
logger.warning("Preprocessing failed for %s: %s", path, exc)
|
|
preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
|
|
|
|
batch = np.stack(preprocessed, axis=0) # (N, 3, H, W)
|
|
logits = self._infer(batch) # (N, 2)
|
|
probs = calibrated_softmax(logits, self.temperature)
|
|
return probs[:, 1].tolist() # target-class column
|
|
|
|
def _preprocess(self, crop_path: str) -> np.ndarray:
|
|
from PIL import Image
|
|
image = Image.open(crop_path).convert("RGB")
|
|
if image.size != (self.input_size, self.input_size):
|
|
image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
|
|
arr = np.array(image, dtype=np.float32) / 255.0
|
|
arr = (arr - _MEAN) / _STD
|
|
image.close()
|
|
return np.transpose(arr, (2, 0, 1)) # HWC → CHW
|
|
|
|
def _infer(self, batch: np.ndarray) -> np.ndarray:
|
|
if self._session is None:
|
|
return self._placeholder_logits(batch.shape[0])
|
|
|
|
try:
|
|
if hasattr(self._session, "run"):
|
|
input_name = self._session.get_inputs()[0].name
|
|
output = self._session.run(None, {input_name: batch})
|
|
return np.array(output[0])
|
|
|
|
# TensorRT path
|
|
import pycuda.autoinit # noqa: F401
|
|
import pycuda.driver as cuda
|
|
|
|
context = self._session.create_execution_context()
|
|
out_shape = (batch.shape[0], 2)
|
|
d_in = cuda.mem_alloc(batch.nbytes)
|
|
d_out = cuda.mem_alloc(np.prod(out_shape) * np.dtype(np.float32).itemsize)
|
|
stream = cuda.Stream()
|
|
cuda.memcpy_htod_async(d_in, batch, stream)
|
|
context.execute_async_v2(bindings=[int(d_in), int(d_out)], stream_handle=stream.handle)
|
|
output = np.empty(out_shape, dtype=np.float32)
|
|
cuda.memcpy_dtoh_async(output, d_out, stream)
|
|
stream.synchronize()
|
|
return output
|
|
except Exception as exc:
|
|
logger.warning("Classifier inference failed: %s", exc)
|
|
return self._placeholder_logits(batch.shape[0])
|
|
|
|
@staticmethod
|
|
def _placeholder_logits(n: int) -> np.ndarray:
|
|
"""Return neutral logits when no model is loaded."""
|
|
return np.zeros((n, 2), dtype=np.float32)
|