Refactor anomaly algorithms and add Hailo/KV260 backends

anomavision/__init__.py

@@ -8,11 +8,11 @@
 Provides functions for performing anomaly detection in images.
 """
 
+from .algorithm.common.feature_extraction import ResnetEmbeddingsExtractor
+from .algorithm.padim import Padim
+from .algorithm.patchcore import PatchCore
 from .datasets.dataset import AnodetDataset
 from .datasets.mvtec_dataset import MVTecDataset
-from .feature_extraction import ResnetEmbeddingsExtractor
-from .padim import Padim
-from .patchcore import PatchCore
 from .sampling_methods.kcenter_greedy import kCenterGreedy
 from .test import optimal_threshold, visualize_eval_data, visualize_eval_pair
 from .utils import get_logger  # Export for users

anomavision/algorithm/__init__.py

@@ -0,0 +1 @@
+"""Anomaly detection algorithm implementations."""

anomavision/algorithm/common/__init__.py

@@ -0,0 +1,10 @@
+"""Shared algorithm components used by anomaly detection implementations."""
+
+from .feature_extraction import ResnetEmbeddingsExtractor, concatenate_layers
+from .mahalanobis import MahalanobisDistance
+
+__all__ = [
+    "MahalanobisDistance",
+    "ResnetEmbeddingsExtractor",
+    "concatenate_layers",
+]

anomavision/algorithm/padim/__init__.py

@@ -0,0 +1,6 @@
+"""PaDiM anomaly detection algorithms."""
+
+from .padim import Padim
+from .padim_lite import PadimLite, load_padim_lite
+
+__all__ = ["Padim", "PadimLite", "load_padim_lite"]

anomavision/padim.py → anomavision/algorithm/padim/padim.py

@@ -10,9 +10,9 @@
 import torch
 import torch.nn.functional as F
 
-from .feature_extraction import ResnetEmbeddingsExtractor
-from .mahalanobis import MahalanobisDistance
-from .utils import pytorch_cov, split_tensor_and_run_function
+from ...utils import pytorch_cov, split_tensor_and_run_function
+from ..common.feature_extraction import ResnetEmbeddingsExtractor
+from ..common.mahalanobis import MahalanobisDistance
 
 BACKBONE_FEATURE_SIZES = {
     "resnet18": OrderedDict([(0, [64]), (1, [128]), (2, [256]), (3, [512])]),

anomavision/padim_lite.py → anomavision/algorithm/padim/padim_lite.py

@@ -4,8 +4,8 @@
 import torch
 import torch.nn.functional as F
 
-from .feature_extraction import ResnetEmbeddingsExtractor
-from .mahalanobis import MahalanobisDistance
+from ..common.feature_extraction import ResnetEmbeddingsExtractor
+from ..common.mahalanobis import MahalanobisDistance
 
 
 class PadimLite(torch.nn.Module):

anomavision/algorithm/patchcore/__init__.py

@@ -0,0 +1,10 @@
+"""Lightweight PatchCore anomaly detection algorithms."""
+
+from ..common.feature_extraction import ResnetEmbeddingsExtractor
+from .patchcore import PatchCore, build_patchcore_from_stats
+
+__all__ = [
+    "PatchCore",
+    "ResnetEmbeddingsExtractor",
+    "build_patchcore_from_stats",
+]

anomavision/patchcore.py → anomavision/algorithm/patchcore/patchcore.py

@@ -1,7 +1,7 @@
 """Lightweight PatchCore anomaly detection.
 
 This module provides a bounded-memory PatchCore implementation that follows the
-public design of :mod:`anomavision.padim`: fit on a normal-image DataLoader, predict
+public design of :mod:`anomavision.algorithm.padim`: fit on a normal-image DataLoader, predict
 image scores and spatial maps, and save a compact deployment artifact.
 """
 
@@ -12,7 +12,7 @@
 import torch
 import torch.nn.functional as F
 
-from .feature_extraction import ResnetEmbeddingsExtractor
+from ..common.feature_extraction import ResnetEmbeddingsExtractor
 
 
 class PatchCore(torch.nn.Module):
@@ -24,7 +24,7 @@ class PatchCore(torch.nn.Module):
     patch distance; the pixel map is the patch-distance grid upsampled to the input
     resolution.
 
-    The public methods intentionally mirror :class:`anomavision.padim.Padim`, so the
+    The public methods intentionally mirror :class:`anomavision.algorithm.padim.Padim`, so the
     model can be selected by the existing CLI training, inference, evaluation, and
     export workflows.
 

anomavision/export.py

@@ -30,12 +30,12 @@
 )
 from PIL import Image
 
-from anomavision.config import load_config
-from anomavision.general import determine_device
-from anomavision.padim_lite import (  # stats-only .pth → runtime module
+from anomavision.algorithm.padim.padim_lite import (  # stats-only .pth → runtime module
     build_padim_from_stats,
 )
-from anomavision.patchcore import build_patchcore_from_stats
+from anomavision.algorithm.patchcore import build_patchcore_from_stats
+from anomavision.config import load_config
+from anomavision.general import determine_device
 from anomavision.utils import (
     create_image_transform,
     get_logger,

anomavision/inference/model/backends/__init__.py

@@ -5,6 +5,8 @@
 """
 
 from .base import InferenceBackend, ScoresMaps
+from .hailo_backend import HailoAnomalyRuntime, HailoBackend
+from .k260_backend import K260Backend, KV260Backend
 from .onnx_backend import OnnxBackend
 from .openvino_backend import OpenVinoBackend
 from .tensorrt_backend import TensorRTBackend
@@ -19,4 +21,8 @@
     "TensorRTBackend",
     "OpenVinoBackend",
     "TorchScriptBackend",
+    "HailoBackend",
+    "HailoAnomalyRuntime",
+    "K260Backend",
+    "KV260Backend",
 ]

anomavision/inference/model/backends/hailo_backend.py

@@ -0,0 +1,194 @@
+"""HailoRT runtime for complete AnomaVision anomaly HEFs.
+
+The HEF is expected to expose two outputs generated by ``hailo_export``:
+``image_scores`` and ``score_map``. Feature extraction and distance calculation
+must already be compiled into the HEF. This runtime intentionally contains no
+fallback CNN or CPU distance implementation, which prevents accidental partial
+quantization on Kria.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, Tuple
+
+import numpy as np
+from PIL import Image
+
+from .base import InferenceBackend
+
+
+class HailoAnomalyRuntime:
+    """Run a complete PaDiM or PatchCore HEF through HailoRT."""
+
+    def __init__(
+        self,
+        hef_path: str | Path,
+        input_size: Tuple[int, int] = (224, 224),
+        input_dtype: np.dtype = np.float32,
+    ) -> None:
+        """Load and configure a complete Hailo-8 HEF.
+
+        Args:
+            hef_path: Path to a HEF exposing ``image_scores`` and ``score_map``.
+            input_size: Fixed ``(height, width)`` expected by the HEF.
+            input_dtype: Host input dtype passed to HailoRT.
+
+        Raises:
+            RuntimeError: If HailoRT is unavailable or has no network group.
+            FileNotFoundError: If ``hef_path`` does not exist.
+            ValueError: If required anomaly outputs are missing.
+        """
+        try:
+            from hailo_platform import (
+                HEF,
+                ConfigureParams,
+                FormatType,
+                HailoStreamInterface,
+                InferVStreams,
+                InputVStreamParams,
+                OutputVStreamParams,
+                VDevice,
+            )
+        except ImportError as exc:  # pragma: no cover - depends on Kria image
+            raise RuntimeError(
+                "HailoRT is not installed. Install the HailoRT Python package on "
+                "the Kria K26 image before loading a HEF."
+            ) from exc
+
+        self._api = {
+            "ConfigureParams": ConfigureParams,
+            "FormatType": FormatType,
+            "HEF": HEF,
+            "HailoStreamInterface": HailoStreamInterface,
+            "InputVStreamParams": InputVStreamParams,
+            "InferVStreams": InferVStreams,
+            "OutputVStreamParams": OutputVStreamParams,
+            "VDevice": VDevice,
+        }
+        self.hef_path = Path(hef_path)
+        if not self.hef_path.exists():
+            raise FileNotFoundError(self.hef_path)
+        self.input_size = tuple(int(v) for v in input_size)
+        self.input_dtype = input_dtype
+        self.device = VDevice()
+        self.hef = HEF(str(self.hef_path))
+        self.network_groups = self.device.configure(self.hef)
+        if not self.network_groups:
+            raise RuntimeError(f"No network group found in {self.hef_path}")
+        self.network_group = self.network_groups[0]
+        self.network_group_params = self.network_group.create_params()
+        self.input_name = self.hef.get_input_vstream_infos()[0].name
+        output_names = [info.name for info in self.hef.get_output_vstream_infos()]
+        required = {"image_scores", "score_map"}
+        missing = sorted(required.difference(output_names))
+        if missing:
+            raise ValueError(
+                "The HEF is not a complete AnomaVision anomaly graph; missing "
+                f"outputs: {', '.join(missing)}"
+            )
+        self.output_names = output_names
+
+    def _preprocess(self, image: Image.Image | np.ndarray | str | Path) -> np.ndarray:
+        """Convert an image path, PIL image, or RGB array to NCHW input."""
+        if isinstance(image, (str, Path)):
+            image = Image.open(image)
+        if isinstance(image, Image.Image):
+            image = np.asarray(image.convert("RGB"))
+        image = np.asarray(image)
+        if image.ndim != 3 or image.shape[2] != 3:
+            raise ValueError("image must be an HxWx3 RGB image")
+        image = np.asarray(
+            Image.fromarray(image.astype(np.uint8), mode="RGB").resize(
+                (self.input_size[1], self.input_size[0]), Image.Resampling.BILINEAR
+            ),
+            dtype=np.float32,
+        )
+        # Match AnomaVision's tensor contract: NCHW float RGB in [0, 1].
+        return np.transpose(image / 255.0, (2, 0, 1))[None].astype(self.input_dtype)
+
+    def predict(
+        self, image: Image.Image | np.ndarray | str | Path
+    ) -> Dict[str, np.ndarray]:
+        """Run one image and return complete image and localization outputs.
+
+        Args:
+            image: An image path, PIL RGB image, or HxWx3 RGB array.
+
+        Returns:
+            A mapping containing ``image_scores`` and ``score_map`` arrays.
+        """
+        api = self._api
+        input_params = api["InputVStreamParams"].make(
+            self.network_group, quantized=False, format_type=api["FormatType"].FLOAT32
+        )
+        output_params = api["OutputVStreamParams"].make(
+            self.network_group, quantized=False, format_type=api["FormatType"].FLOAT32
+        )
+        tensor = self._preprocess(image)
+        with self.network_group.activate(self.network_group_params):
+            with api["InferVStreams"](
+                self.network_group, input_params, output_params
+            ) as infer_pipeline:
+                outputs = infer_pipeline.infer({self.input_name: tensor})
+        return {
+            "image_scores": np.asarray(outputs["image_scores"]).squeeze(),
+            "score_map": np.asarray(outputs["score_map"]).squeeze(),
+        }
+
+    def close(self) -> None:
+        """Release the Hailo virtual device."""
+        release = getattr(self.device, "release", None)
+        if callable(release):
+            release()
+
+    def __enter__(self) -> "HailoAnomalyRuntime":
+        return self
+
+    def __exit__(self, exc_type, exc, traceback) -> None:
+        self.close()
+
+
+class HailoBackend(InferenceBackend):
+    """AnomaVision inference backend for a complete Hailo-8 HEF."""
+
+    def __init__(
+        self,
+        model_path: str | Path,
+        device: str = "hailo",
+        input_size: Tuple[int, int] = (224, 224),
+    ) -> None:
+        del device
+        self.runtime = HailoAnomalyRuntime(model_path, input_size=input_size)
+
+    def predict(self, batch) -> Tuple[np.ndarray, np.ndarray]:
+        """Run one image through the common backend contract.
+
+        Args:
+            batch: An HxWx3 RGB image or a single-image 1x3xHxW/1xHxWx3 batch.
+
+        Returns:
+            A tuple ``(image_scores, score_maps)`` as NumPy arrays.
+        """
+        array = np.asarray(batch)
+        if array.ndim == 4:
+            if array.shape[0] != 1:
+                raise ValueError("HailoBackend currently supports batch size 1")
+            array = (
+                np.transpose(array[0], (1, 2, 0)) if array.shape[1] == 3 else array[0]
+            )
+        elif array.ndim != 3:
+            raise ValueError("batch must be an HxWx3 or 1x3xHxW image")
+        result = self.runtime.predict(array)
+        return result["image_scores"], result["score_map"]
+
+    def warmup(self, batch=None, runs: int = 2) -> None:
+        """Warm up the device with a supplied image batch."""
+        if batch is None:
+            raise ValueError("Hailo warmup requires a sample image batch")
+        for _ in range(max(1, int(runs))):
+            self.predict(batch)
+
+    def close(self) -> None:
+        """Release HailoRT resources through the shared backend lifecycle."""
+        self.runtime.close()