715 lines
30 KiB
Python
715 lines
30 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import logging
|
||
import math
|
||
import threading
|
||
from contextlib import contextmanager
|
||
from pathlib import Path
|
||
|
||
from PIL import Image
|
||
|
||
from .accelerator import AcceleratorService
|
||
from .config import Settings
|
||
from .model_validation import (
|
||
sentence_transformer_module_path,
|
||
validate_face_model_root,
|
||
validate_openvino_ir_directory,
|
||
validate_sentence_transformer_root,
|
||
)
|
||
from .text import search_tokens
|
||
from .vectors import normalize
|
||
|
||
logger = logging.getLogger(__name__)
|
||
_OPTIMUM_LOCAL_LOAD_LOCK = threading.Lock()
|
||
|
||
|
||
class ModelUnavailable(RuntimeError):
|
||
pass
|
||
|
||
|
||
class VisualModelStageError(RuntimeError):
|
||
"""Keep the failing visual tower visible without hiding its root cause."""
|
||
|
||
def __init__(self, stage: str, label: str, reason: BaseException):
|
||
self.stage = stage
|
||
self.reason = reason
|
||
detail = " ".join(str(reason).split())
|
||
suffix = f"：{detail}" if detail else ""
|
||
super().__init__(f"{label}加载失败：{type(reason).__name__}{suffix}")
|
||
|
||
|
||
@contextmanager
|
||
def _allow_existing_absolute_local_hf_paths():
|
||
"""Let Optimum/HF validators accept fnOS absolute model paths.
|
||
|
||
Several Optimum 2.x loaders still pass local ``/vol1/...`` paths through
|
||
Hugging Face repo-id validation, even when the caller has already set
|
||
``local_files_only=True`` and supplied a validated OpenVINO IR directory.
|
||
Keep the bypass scoped to this local load window and only accept paths that
|
||
already exist on disk.
|
||
"""
|
||
|
||
try:
|
||
from huggingface_hub.utils import _validators
|
||
except Exception:
|
||
yield
|
||
return
|
||
|
||
original = _validators.validate_repo_id
|
||
|
||
def validate_repo_id(repo_id):
|
||
if isinstance(repo_id, (str, Path)):
|
||
try:
|
||
candidate = Path(repo_id)
|
||
except TypeError:
|
||
candidate = None
|
||
if candidate is not None and candidate.is_absolute() and candidate.exists():
|
||
return None
|
||
return original(repo_id)
|
||
|
||
_validators.validate_repo_id = validate_repo_id
|
||
try:
|
||
yield
|
||
finally:
|
||
_validators.validate_repo_id = original
|
||
|
||
|
||
class _OpenVINOClipImageEncoder:
|
||
"""Small ``encode`` adapter around Optimum's CLIP OpenVINO model.
|
||
|
||
SentenceTransformers' generic OpenVINO backend is designed around text
|
||
feature-extraction inputs. A CLIP image tower needs the dedicated
|
||
zero-shot image model so that ``pixel_values`` and ``image_embeds`` are
|
||
preserved during export and inference.
|
||
"""
|
||
|
||
def __init__(self, path: Path, device: str, ov_config: dict[str, object]):
|
||
export = path / "openvino"
|
||
try:
|
||
validate_openvino_ir_directory(export, "画面语义图像")
|
||
except ValueError as exc:
|
||
raise ModelUnavailable(str(exc)) from exc
|
||
try:
|
||
from optimum.intel.openvino import OVModelForZeroShotImageClassification
|
||
from transformers import AutoConfig, AutoProcessor
|
||
except ImportError as exc:
|
||
raise ModelUnavailable("缺少 CLIP OpenVINO 图像运行依赖") from exc
|
||
processor_path = sentence_transformer_module_path(path, "CLIPModel")
|
||
with _allow_existing_absolute_local_hf_paths():
|
||
config = AutoConfig.from_pretrained(
|
||
str(processor_path.resolve()), local_files_only=True
|
||
)
|
||
# Optimum 2.x's public ``from_pretrained`` always calls
|
||
# ``TasksManager.infer_library_from_model`` before delegating to the
|
||
# OpenVINO loader. On fnOS an absolute ``/vol1/...`` model path can be
|
||
# handed to huggingface_hub as though it were a repository id, even
|
||
# when ``local_files_only`` is set. The IR directory has already been
|
||
# validated above and the Transformers config is explicit, so enter
|
||
# the local loader directly and make any Hub/library inference
|
||
# impossible.
|
||
self.model = OVModelForZeroShotImageClassification._from_pretrained(
|
||
str(export.resolve()),
|
||
config=config,
|
||
file_name="openvino_model.xml",
|
||
local_files_only=True,
|
||
device=device,
|
||
ov_config=ov_config,
|
||
)
|
||
self.processor = AutoProcessor.from_pretrained(
|
||
str(processor_path.resolve()), local_files_only=True
|
||
)
|
||
|
||
def encode(self, images, normalize_embeddings: bool = True):
|
||
import numpy as np
|
||
|
||
values = list(images)
|
||
inputs = self.processor(
|
||
text=[""] * len(values),
|
||
images=values,
|
||
return_tensors="pt",
|
||
padding=True,
|
||
)
|
||
output = self.model(**inputs).image_embeds
|
||
if hasattr(output, "detach"):
|
||
output = output.detach()
|
||
if hasattr(output, "cpu"):
|
||
output = output.cpu()
|
||
vectors = np.asarray(output, dtype=np.float32)
|
||
if normalize_embeddings:
|
||
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
|
||
vectors = vectors / np.maximum(norms, np.finfo(np.float32).eps)
|
||
return vectors
|
||
|
||
|
||
def _load_local_openvino_sentence_transformer(model_class, path: Path, kwargs: dict):
|
||
"""Load a persisted text IR without Optimum's Hub library inference.
|
||
|
||
SentenceTransformer delegates its Transformer module to
|
||
``OVModelForFeatureExtraction.from_pretrained``. Optimum 2.x performs Hub
|
||
library detection in that public method even for an existing absolute
|
||
fnOS ``/vol1/...`` directory. Intercept only that construction call and
|
||
enter the already validated local IR loader directly.
|
||
"""
|
||
|
||
try:
|
||
from optimum.intel.openvino import OVModelForFeatureExtraction
|
||
except ImportError as exc:
|
||
raise ModelUnavailable("缺少 CLIP OpenVINO 文本运行依赖") from exc
|
||
|
||
model_type = OVModelForFeatureExtraction
|
||
inherited = "from_pretrained" not in vars(model_type)
|
||
original = vars(model_type).get("from_pretrained")
|
||
|
||
@classmethod
|
||
def from_local_ir(cls, model_id, config=None, export=False, **model_kwargs):
|
||
if export:
|
||
raise ModelUnavailable("画面语义文本禁止在推理阶段重新导出")
|
||
return cls._from_pretrained(model_id=model_id, config=config, **model_kwargs)
|
||
|
||
# The override is process-global, so serialize the very short module
|
||
# construction window and restore the exact original descriptor.
|
||
with _OPTIMUM_LOCAL_LOAD_LOCK:
|
||
model_type.from_pretrained = from_local_ir
|
||
try:
|
||
with _allow_existing_absolute_local_hf_paths():
|
||
return model_class(str(path), **kwargs)
|
||
finally:
|
||
if inherited:
|
||
delattr(model_type, "from_pretrained")
|
||
else:
|
||
model_type.from_pretrained = original
|
||
|
||
|
||
class EmbeddingService:
|
||
"""Lazy visual/text encoder.
|
||
|
||
A release model bundle contains two SentenceTransformer-compatible directories.
|
||
The multilingual text tower is trained to match the CLIP image tower's space.
|
||
"""
|
||
|
||
dimensions = 512
|
||
|
||
def __init__(self, settings: Settings, accelerator: AcceleratorService | None = None):
|
||
self.settings = settings
|
||
self.accelerator = accelerator or AcceleratorService(settings)
|
||
self._image_model = None
|
||
self._text_model = None
|
||
|
||
@property
|
||
def image_path(self) -> Path:
|
||
return self.settings.models_dir / "visual" / "image"
|
||
|
||
@property
|
||
def text_path(self) -> Path:
|
||
return self.settings.models_dir / "visual" / "text"
|
||
|
||
def status(self) -> dict:
|
||
exported = False
|
||
if self.image_path.exists() and self.text_path.exists():
|
||
try:
|
||
validate_openvino_ir_directory(
|
||
self.image_path / "openvino", "画面语义图像"
|
||
)
|
||
text_module = sentence_transformer_module_path(self.text_path, "Transformer")
|
||
validate_openvino_ir_directory(
|
||
text_module / "openvino", "画面语义文本"
|
||
)
|
||
exported = True
|
||
except ValueError:
|
||
pass
|
||
return {
|
||
"backend": self.settings.embedding_backend,
|
||
"visual_ready": self.settings.embedding_backend == "hash"
|
||
or (self.image_path.exists() and self.text_path.exists()),
|
||
"image_model": self.settings.visual_image_model,
|
||
"text_model": self.settings.visual_text_model,
|
||
"openvino_exported": exported,
|
||
}
|
||
|
||
def _load(self) -> None:
|
||
if self.settings.embedding_backend == "hash":
|
||
return
|
||
if not self.image_path.exists() or not self.text_path.exists():
|
||
raise ModelUnavailable("视觉模型尚未下载，请先在设置中安装模型包")
|
||
try:
|
||
from sentence_transformers import SentenceTransformer
|
||
except ImportError as exc:
|
||
raise ModelUnavailable("缺少 AI 运行依赖，请安装 imagefind[ai]") from exc
|
||
if self._image_model is None:
|
||
validate_sentence_transformer_root(self.image_path, "画面语义图像")
|
||
validate_sentence_transformer_root(self.text_path, "画面语义文本")
|
||
backend = "openvino" if self.settings.embedding_backend in {"auto", "openvino"} else "torch"
|
||
device = self.accelerator.device_for("visual") if backend == "openvino" else "CPU"
|
||
try:
|
||
self._build_models(SentenceTransformer, backend, device)
|
||
except Exception as exc:
|
||
self._image_model = None
|
||
self._text_model = None
|
||
reason = exc.reason if isinstance(exc, VisualModelStageError) else exc
|
||
stage = exc.stage if isinstance(exc, VisualModelStageError) else "encoder_compile"
|
||
if (
|
||
backend == "openvino"
|
||
and device != "CPU"
|
||
and self.accelerator.is_device_error(reason)
|
||
and self.accelerator.fall_back("visual", reason, stage=stage)
|
||
):
|
||
logger.warning("visual GPU model load failed; retrying on CPU", exc_info=True)
|
||
try:
|
||
self._build_models(SentenceTransformer, backend, "CPU")
|
||
except Exception as cpu_exc:
|
||
self._image_model = None
|
||
self._text_model = None
|
||
cpu_reason = (
|
||
cpu_exc.reason if isinstance(cpu_exc, VisualModelStageError) else cpu_exc
|
||
)
|
||
cpu_stage = (
|
||
cpu_exc.stage
|
||
if isinstance(cpu_exc, VisualModelStageError)
|
||
else "encoder_compile"
|
||
)
|
||
self.accelerator.mark_unavailable(
|
||
"visual", cpu_reason, stage=cpu_stage
|
||
)
|
||
raise
|
||
else:
|
||
self.accelerator.mark_unavailable("visual", reason, stage=stage)
|
||
raise
|
||
|
||
def _build_models(self, model_class, backend: str, device: str) -> None:
|
||
image_path = self.image_path.resolve()
|
||
text_path = self.text_path.resolve()
|
||
# Optimum's public examples and Intel GPU plugin use the canonical GPU
|
||
# alias. Core discovery may return GPU.0, which remains useful for
|
||
# reporting but has caused inconsistent Python-binding errors when
|
||
# passed through higher-level model wrappers.
|
||
compile_device = "GPU" if device.upper().startswith("GPU") else device
|
||
|
||
def load_text(path: Path):
|
||
kwargs: dict[str, object] = {"backend": backend, "local_files_only": True}
|
||
if backend == "openvino":
|
||
transformer_path = sentence_transformer_module_path(path, "Transformer")
|
||
export = transformer_path / "openvino"
|
||
try:
|
||
validate_openvino_ir_directory(export, "画面语义文本")
|
||
except ValueError as exc:
|
||
raise ModelUnavailable(str(exc)) from exc
|
||
model_kwargs: dict[str, object] = {
|
||
"device": compile_device,
|
||
"ov_config": self.accelerator.ov_config(device),
|
||
# SentenceTransformers otherwise scans the wrapper and may
|
||
# silently ask Optimum to export again. Point it at the
|
||
# exact persisted IR and make runtime export impossible.
|
||
"export": False,
|
||
"file_name": "openvino_model.xml",
|
||
# SentenceTransformer instantiates this module with the
|
||
# resolved Transformer module directory (for example
|
||
# ``text/0_Transformer``), not with the wrapper root.
|
||
# Supplying ``0_Transformer/openvino`` here duplicates the
|
||
# module segment and makes Optimum fall through to Hub
|
||
# repository validation. The IR is always directly below
|
||
# the module directory.
|
||
"subfolder": "openvino",
|
||
}
|
||
kwargs["model_kwargs"] = model_kwargs
|
||
if backend == "openvino" and str(getattr(model_class, "__module__", "")).startswith(
|
||
"sentence_transformers"
|
||
):
|
||
return _load_local_openvino_sentence_transformer(model_class, path, kwargs)
|
||
return model_class(str(path), **kwargs)
|
||
|
||
if backend == "openvino":
|
||
try:
|
||
self._image_model = _OpenVINOClipImageEncoder(
|
||
image_path,
|
||
compile_device,
|
||
self.accelerator.ov_config(device),
|
||
)
|
||
except ModelUnavailable:
|
||
raise
|
||
except Exception as exc:
|
||
raise VisualModelStageError(
|
||
"image_encoder_compile", "画面语义图像 OpenVINO", exc
|
||
) from exc
|
||
else:
|
||
self._image_model = model_class(str(image_path), backend=backend, local_files_only=True)
|
||
try:
|
||
self._text_model = load_text(text_path)
|
||
except ModelUnavailable:
|
||
raise
|
||
except Exception as exc:
|
||
raise VisualModelStageError(
|
||
"text_encoder_compile", "画面语义文本 OpenVINO", exc
|
||
) from exc
|
||
reader = getattr(self.accelerator, "execution_devices_from", None)
|
||
execution_devices = reader(self._image_model, self._text_model) if callable(reader) else [device]
|
||
if device.upper().startswith("GPU") and not any(
|
||
str(actual).upper().startswith("GPU") for actual in execution_devices
|
||
):
|
||
raise RuntimeError("OpenVINO 未确认画面语义模型在 GPU 上执行")
|
||
self.accelerator.mark_ready("visual", device, execution_devices)
|
||
|
||
def _retry_cpu(self, kind: str, value, error: Exception):
|
||
if self.accelerator.device_for("visual") == "CPU":
|
||
raise error
|
||
self.accelerator.fall_back("visual", error, stage="inference")
|
||
self._image_model = None
|
||
self._text_model = None
|
||
self._load()
|
||
model = self._text_model if kind == "text" else self._image_model
|
||
return model.encode([value], normalize_embeddings=True)[0]
|
||
|
||
@staticmethod
|
||
def _hash_vector(data: bytes) -> list[float]:
|
||
values = [0.0] * 512
|
||
digest = hashlib.shake_256(data).digest(2048)
|
||
for index, byte in enumerate(digest):
|
||
values[index % 512] += (byte - 127.5) / 127.5
|
||
return normalize(values)
|
||
|
||
def encode_text(self, text: str) -> list[float]:
|
||
if self.settings.embedding_backend == "hash":
|
||
values = [0.0] * self.dimensions
|
||
for token in search_tokens(text):
|
||
digest = hashlib.sha256(token.encode()).digest()
|
||
index = int.from_bytes(digest[:2], "big") % self.dimensions
|
||
values[index] += 1 if digest[2] & 1 else -1
|
||
return normalize(values)
|
||
self._load()
|
||
try:
|
||
vector = self._text_model.encode([text], normalize_embeddings=True)[0]
|
||
except Exception as exc:
|
||
vector = self._retry_cpu("text", text, exc)
|
||
return [float(value) for value in vector]
|
||
|
||
def encode_image(self, image: Image.Image | Path) -> list[float]:
|
||
if isinstance(image, Path):
|
||
with Image.open(image) as opened:
|
||
return self.encode_image(opened.convert("RGB"))
|
||
if self.settings.embedding_backend == "hash":
|
||
sample = image.convert("RGB").resize((32, 32)).tobytes()
|
||
return self._hash_vector(sample)
|
||
self._load()
|
||
converted = image.convert("RGB")
|
||
try:
|
||
vector = self._image_model.encode([converted], normalize_embeddings=True)[0]
|
||
except Exception as exc:
|
||
vector = self._retry_cpu("image", converted, exc)
|
||
return [float(value) for value in vector]
|
||
|
||
def encode_images(self, images: list[Image.Image | Path]) -> list[list[float]]:
|
||
"""Encode a bounded batch while keeping input order.
|
||
|
||
OpenVINO and SentenceTransformers both benefit from a small batch. The
|
||
caller controls the batch size from the GPU profile; paths are opened
|
||
here so file handles never escape the call.
|
||
"""
|
||
if not images:
|
||
return []
|
||
if self.settings.embedding_backend == "hash":
|
||
return [self.encode_image(image) for image in images]
|
||
self._load()
|
||
opened: list[Image.Image] = []
|
||
try:
|
||
for image in images:
|
||
if isinstance(image, Path):
|
||
with Image.open(image) as source:
|
||
opened.append(source.convert("RGB"))
|
||
else:
|
||
opened.append(image.convert("RGB"))
|
||
try:
|
||
vectors = self._image_model.encode(opened, normalize_embeddings=True)
|
||
except Exception as exc:
|
||
if self.accelerator.device_for("visual") == "CPU":
|
||
raise
|
||
self.accelerator.fall_back("visual", exc, stage="inference")
|
||
self._image_model = None
|
||
self._text_model = None
|
||
self._load()
|
||
vectors = self._image_model.encode(opened, normalize_embeddings=True)
|
||
return [[float(value) for value in vector] for vector in vectors]
|
||
finally:
|
||
for image in opened:
|
||
image.close()
|
||
|
||
def verify_acceleration(self) -> dict:
|
||
"""Load both persisted IR towers and verify finite normalized vectors."""
|
||
|
||
self.reset()
|
||
vectors = (
|
||
self.encode_text("ImageFind Intel GPU verification"),
|
||
self.encode_image(Image.new("RGB", (32, 32), (103, 145, 244))),
|
||
)
|
||
for vector in vectors:
|
||
norm = math.sqrt(sum(value * value for value in vector))
|
||
if len(vector) != self.dimensions or not all(math.isfinite(value) for value in vector):
|
||
raise RuntimeError("画面语义模型输出不是有效的 512 维向量")
|
||
if not 0.98 <= norm <= 1.02:
|
||
raise RuntimeError("画面语义模型输出未正确归一化")
|
||
return self.accelerator.status()["components"]["visual"]
|
||
|
||
def reset(self) -> None:
|
||
self._image_model = None
|
||
self._text_model = None
|
||
self.accelerator.reset("visual")
|
||
|
||
|
||
class OcrService:
|
||
def __init__(self, settings: Settings, accelerator: AcceleratorService | None = None):
|
||
self.settings = settings
|
||
self.accelerator = accelerator or AcceleratorService(settings)
|
||
self._engine = None
|
||
self._engine_device: str | None = None
|
||
|
||
def ready(self) -> bool:
|
||
root = self.settings.models_dir / "ocr"
|
||
return (root / "det.onnx").is_file() and (root / "rec.onnx").is_file()
|
||
|
||
def recognize(self, image: Path) -> list[tuple[str, float]]:
|
||
if not self.ready():
|
||
return []
|
||
if self._engine is None:
|
||
device = self.accelerator.device_for("ocr")
|
||
try:
|
||
self._build_engine(device)
|
||
except Exception as exc:
|
||
if (
|
||
device != "CPU"
|
||
and self.accelerator.is_device_error(exc)
|
||
and self.accelerator.fall_back("ocr", exc, stage="encoder_compile")
|
||
):
|
||
logger.warning("OCR GPU load failed; retrying on CPU", exc_info=True)
|
||
try:
|
||
self._build_engine("CPU")
|
||
except Exception as retry_exc:
|
||
self.accelerator.mark_unavailable("ocr", retry_exc, stage="encoder_compile")
|
||
raise
|
||
else:
|
||
self.accelerator.mark_unavailable("ocr", exc, stage="encoder_compile")
|
||
raise
|
||
try:
|
||
result, _ = self._engine(str(image))
|
||
except Exception as exc:
|
||
if self._engine_device != "CPU" and self.accelerator.fall_back("ocr", exc, stage="inference"):
|
||
self._engine = None
|
||
self._build_engine("CPU")
|
||
result, _ = self._engine(str(image))
|
||
else:
|
||
self.accelerator.mark_unavailable("ocr", exc, stage="inference")
|
||
raise
|
||
if not result:
|
||
return []
|
||
return [(str(line[1]), float(line[2])) for line in result if len(line) >= 3]
|
||
|
||
def recognize_batch(self, images: list[Path]) -> list[list[tuple[str, float]]]:
|
||
"""Process a bounded batch, reusing one compiled OCR engine."""
|
||
return [self.recognize(image) for image in images]
|
||
|
||
def _build_engine(self, device: str) -> None:
|
||
root = self.settings.models_dir / "ocr"
|
||
paths = {
|
||
"det_model_path": str(root / "det.onnx"),
|
||
"rec_model_path": str(root / "rec.onnx"),
|
||
"cls_model_path": str(root / "cls.onnx") if (root / "cls.onnx").exists() else None,
|
||
}
|
||
if device != "CPU":
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
from rapidocr_onnxruntime.ch_ppocr_cls import text_cls
|
||
from rapidocr_onnxruntime.ch_ppocr_det import text_detect
|
||
from rapidocr_onnxruntime.ch_ppocr_rec import text_recognize
|
||
|
||
# Reuse RapidOCR's mature OCR pre/post-processing while replacing
|
||
# its three short-lived ONNX Runtime session constructors with
|
||
# explicit modern OpenVINO GPU compilations.
|
||
modules = (text_detect, text_cls, text_recognize)
|
||
originals = [module.OrtInferSession for module in modules]
|
||
|
||
def session_factory(config):
|
||
return _OpenVINOOrtSession(config["model_path"], device, self.accelerator.ov_config(device))
|
||
|
||
try:
|
||
for module in modules:
|
||
module.OrtInferSession = session_factory
|
||
self._engine = RapidOCR(**{name: value for name, value in paths.items() if value is not None})
|
||
finally:
|
||
for module, original in zip(modules, originals, strict=True):
|
||
module.OrtInferSession = original
|
||
else:
|
||
# The CPU fallback intentionally keeps the bounded ONNX runtime;
|
||
# it is predictable on older fnOS kernels and honours thread caps.
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
|
||
self._engine = RapidOCR(
|
||
intra_op_num_threads=self.settings.ai_cpu_threads,
|
||
inter_op_num_threads=1,
|
||
**paths,
|
||
)
|
||
self._engine_device = device
|
||
self.accelerator.mark_ready("ocr", device)
|
||
|
||
def reset(self) -> None:
|
||
self._engine = None
|
||
self._engine_device = None
|
||
self.accelerator.reset("ocr")
|
||
|
||
|
||
class _OpenVINOOrtSession:
|
||
"""Adapter matching RapidOCR's ONNX session contract on an OpenVINO device."""
|
||
|
||
def __init__(self, model_path: str, device: str, config: dict[str, object]):
|
||
import openvino as ov
|
||
|
||
core = ov.Core()
|
||
model = core.read_model(model_path)
|
||
self.characters: list[str] | None = None
|
||
try:
|
||
metadata = model.get_rt_info()["framework"]["character"]
|
||
value = metadata.value if hasattr(metadata, "value") else str(metadata)
|
||
self.characters = value.splitlines() or None
|
||
except (AttributeError, KeyError, TypeError):
|
||
pass
|
||
self.compiled = core.compile_model(model, device, config)
|
||
self.output = self.compiled.output(0)
|
||
|
||
def __call__(self, input_content):
|
||
result = self.compiled([input_content])
|
||
return [result[self.output]]
|
||
|
||
def have_key(self, key: str = "character") -> bool:
|
||
return key == "character" and bool(self.characters)
|
||
|
||
def get_character_list(self, key: str = "character") -> list[str] | None:
|
||
return self.characters if key == "character" else None
|
||
|
||
|
||
class FaceService:
|
||
def __init__(self, settings: Settings, accelerator: AcceleratorService | None = None):
|
||
self.settings = settings
|
||
self.accelerator = accelerator or AcceleratorService(settings)
|
||
self._compiled = None
|
||
self._compiled_device: str | None = None
|
||
|
||
def ready(self) -> bool:
|
||
root = self.settings.models_dir / "faces"
|
||
try:
|
||
validate_face_model_root(root)
|
||
except ValueError:
|
||
return False
|
||
return True
|
||
|
||
def detect_and_embed(self, image_path: Path) -> list[dict]:
|
||
if not self.ready():
|
||
return []
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
import openvino as ov
|
||
except ImportError:
|
||
logger.warning("face models exist but OpenVINO/OpenCV is unavailable")
|
||
return []
|
||
root = self.settings.models_dir / "faces"
|
||
if self._compiled is None:
|
||
device = self.accelerator.device_for("faces")
|
||
try:
|
||
self._compile(ov, root, device)
|
||
except Exception as exc:
|
||
if (
|
||
device != "CPU"
|
||
and self.accelerator.is_device_error(exc)
|
||
and self.accelerator.fall_back("faces", exc, stage="encoder_compile")
|
||
):
|
||
logger.warning("face GPU model load failed; retrying on CPU", exc_info=True)
|
||
try:
|
||
self._compile(ov, root, "CPU")
|
||
except Exception as retry_exc:
|
||
self.accelerator.mark_unavailable("faces", retry_exc, stage="encoder_compile")
|
||
raise
|
||
else:
|
||
self.accelerator.mark_unavailable("faces", exc, stage="encoder_compile")
|
||
raise
|
||
try:
|
||
return self._detect(image_path, cv2, np)
|
||
except Exception as exc:
|
||
if self._compiled_device != "CPU" and self.accelerator.fall_back("faces", exc, stage="inference"):
|
||
self._compiled = None
|
||
self._compile(ov, root, "CPU")
|
||
return self._detect(image_path, cv2, np)
|
||
self.accelerator.mark_unavailable("faces", exc, stage="inference")
|
||
raise
|
||
|
||
def detect_and_embed_batch(self, image_paths: list[Path]) -> list[list[dict]]:
|
||
"""Run face detection for a bounded batch without reloading models."""
|
||
return [self.detect_and_embed(image_path) for image_path in image_paths]
|
||
|
||
def _compile(self, ov, root: Path, device: str) -> None:
|
||
core = ov.Core()
|
||
config = self.accelerator.ov_config(device)
|
||
detector = core.compile_model(root / "detector.xml", device, config)
|
||
reid = core.compile_model(root / "reidentification.xml", device, config)
|
||
self._compiled = detector, reid
|
||
self._compiled_device = device
|
||
execution_devices = self.accelerator.execution_devices_from(detector, reid)
|
||
if device.upper().startswith("GPU") and not any(
|
||
actual.upper().startswith("GPU") for actual in execution_devices
|
||
):
|
||
raise RuntimeError("OpenVINO 未确认人物模型在 GPU 上执行")
|
||
self.accelerator.mark_ready("faces", device, execution_devices)
|
||
|
||
def _detect(self, image_path: Path, cv2, np) -> list[dict]:
|
||
detector, reid = self._compiled
|
||
image = cv2.imread(str(image_path))
|
||
if image is None:
|
||
return []
|
||
height, width = image.shape[:2]
|
||
|
||
def input_blob(compiled, frame):
|
||
shape = tuple(compiled.input(0).shape)
|
||
resized = cv2.resize(frame, (shape[3], shape[2]))
|
||
return resized.transpose(2, 0, 1)[None].astype(np.float32)
|
||
|
||
detections = detector([input_blob(detector, image)])[detector.output(0)]
|
||
results = []
|
||
for row in detections.reshape(-1, 7):
|
||
confidence = float(row[2])
|
||
if confidence < 0.65:
|
||
continue
|
||
x1 = max(0, min(width - 1, int(row[3] * width)))
|
||
y1 = max(0, min(height - 1, int(row[4] * height)))
|
||
x2 = max(x1 + 1, min(width, int(row[5] * width)))
|
||
y2 = max(y1 + 1, min(height, int(row[6] * height)))
|
||
crop = image[y1:y2, x1:x2]
|
||
if min(crop.shape[:2]) < 32:
|
||
continue
|
||
embedding = reid([input_blob(reid, crop)])[reid.output(0)].reshape(-1)
|
||
norm = float(np.linalg.norm(embedding))
|
||
if norm:
|
||
embedding = embedding / norm
|
||
results.append({"bbox": [x1, y1, x2, y2], "confidence": confidence, "vector": embedding.tolist()})
|
||
return results
|
||
|
||
def reset(self) -> None:
|
||
self._compiled = None
|
||
self._compiled_device = None
|
||
self.accelerator.reset("faces")
|
||
|
||
|
||
def difference_hash(image: Image.Image) -> str:
|
||
small = image.convert("L").resize((9, 8))
|
||
flattened = getattr(small, "get_flattened_data", None)
|
||
pixels = list(flattened() if flattened else small.getdata())
|
||
value = 0
|
||
for row in range(8):
|
||
for column in range(8):
|
||
value = (value << 1) | (pixels[row * 9 + column] > pixels[row * 9 + column + 1])
|
||
return f"{value:016x}"
|
||
|
||
|
||
def hash_similarity(left: str | None, right: str | None) -> float:
|
||
if not left or not right:
|
||
return 0.0
|
||
try:
|
||
distance = (int(left, 16) ^ int(right, 16)).bit_count()
|
||
except ValueError:
|
||
return 0.0
|
||
return 1.0 - distance / 64
|