# -*- coding: utf-8 -*-
|
import json
|
import logging
|
import os
|
from typing import List, Optional, Tuple
|
|
import numpy as np
|
import onnxruntime as ort
|
from PIL import Image
|
|
logger = logging.getLogger(__name__)
|
|
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
|
|
class OnnxClassifier:
|
def __init__(self, model_path: str, image_size: int = 224):
|
self.image_size = image_size
|
self.session: Optional[ort.InferenceSession] = None
|
if model_path and os.path.isfile(model_path):
|
self.session = ort.InferenceSession(
|
model_path,
|
providers=["CPUExecutionProvider"],
|
)
|
logger.info("加载 ONNX: %s", model_path)
|
|
@property
|
def loaded(self) -> bool:
|
return self.session is not None
|
|
def preprocess(self, image_path: str) -> np.ndarray:
|
img = Image.open(image_path).convert("RGB").resize((self.image_size, self.image_size))
|
arr = np.array(img).astype(np.float32) / 255.0
|
arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
|
return arr.transpose(2, 0, 1)[None]
|
|
def predict_batch(self, image_paths: List[str], batch_size: int = 16) -> List[float]:
|
if not self.session:
|
return [0.0] * len(image_paths)
|
scores = []
|
for i in range(0, len(image_paths), batch_size):
|
batch_paths = image_paths[i : i + batch_size]
|
batch = np.concatenate([self.preprocess(p) for p in batch_paths], axis=0)
|
logits = self.session.run(None, {"input": batch})[0]
|
for logit in logits:
|
scores.append(float(1.0 / (1.0 + np.exp(-logit[0]))))
|
return scores
|
|
|
class ModelRegistry:
|
def __init__(self, model_dir: str):
|
self.model_dir = model_dir
|
self.version = "unknown"
|
self.image_size = 224
|
self.storefront: Optional[OnnxClassifier] = None
|
self.handover: Optional[OnnxClassifier] = None
|
self._load()
|
|
def _load(self):
|
version_path = os.path.join(self.model_dir, "version.json")
|
sf_name = ho_name = None
|
if os.path.isfile(version_path):
|
with open(version_path, encoding="utf-8") as f:
|
meta = json.load(f)
|
self.version = meta.get("model_version", "1.0.0")
|
self.image_size = int(meta.get("image_size", 224))
|
sf_name = meta.get("storefront_model")
|
ho_name = meta.get("handover_model")
|
|
def resolve(name: str, fallback: str) -> str:
|
if name:
|
path = os.path.join(self.model_dir, name)
|
if os.path.isfile(path):
|
return path
|
for suffix in (f"{fallback}_int8.onnx", f"{fallback}.onnx"):
|
path = os.path.join(self.model_dir, suffix)
|
if os.path.isfile(path):
|
return path
|
return ""
|
|
sf_path = resolve(sf_name, "storefront")
|
ho_path = resolve(ho_name, "handover")
|
self.storefront = OnnxClassifier(sf_path, self.image_size) if sf_path else OnnxClassifier("", self.image_size)
|
self.handover = OnnxClassifier(ho_path, self.image_size) if ho_path else OnnxClassifier("", self.image_size)
|
|
@property
|
def ready(self) -> bool:
|
return self.storefront.loaded and self.handover.loaded
|
|
|
def score_frames(
|
registry: ModelRegistry,
|
frames: List[Tuple[float, str]],
|
) -> Tuple[List[Tuple[float, float]], List[Tuple[float, float]]]:
|
times = [f[0] for f in frames]
|
paths = [f[1] for f in frames]
|
sf_scores = registry.storefront.predict_batch(paths)
|
ho_scores = registry.handover.predict_batch(paths)
|
return list(zip(times, sf_scores)), list(zip(times, ho_scores))
|
