diff --git a/.gitignore b/.gitignore index f3f4082..8a809d0 100644 --- a/.gitignore +++ b/.gitignore @@ -369,3 +369,5 @@ tests/__pycache__/* compiled_patchcore_kv260/* quantize_result/* compiled_padim_kv260/* +*.har +*.hef diff --git a/anomavision/inference/model/backends/hailo_backend.py b/anomavision/inference/model/backends/hailo_backend.py index b97f280..4079d97 100644 --- a/anomavision/inference/model/backends/hailo_backend.py +++ b/anomavision/inference/model/backends/hailo_backend.py @@ -1,14 +1,19 @@ """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. +The HEF is expected to expose ``image_scores`` and ``score_map``. Feature +extraction and anomaly scoring are compiled into the HEF; this runtime only +adapts the common AnomaVision inference input contract to HailoRT. + +The public backend contract matches the ONNX backend: input is a single +ImageNet-normalized RGB tensor in NCHW float32. Hailo receives the same values +in NHWC layout. No second resize or normalization is applied to tensors that +have already passed through the AnomaVision dataset preprocessing pipeline. """ from __future__ import annotations +import ctypes +import os from pathlib import Path from typing import Dict, Tuple @@ -18,6 +23,79 @@ from .base import InferenceBackend +def _load_hailort(): + """Load HailoRT and preload its native shared library when necessary.""" + version = "5.3.0" + library_name = f"libhailort.so.{version}" + candidates = [] + + configured = os.environ.get("HAILORT_LIB_PATH") + if configured: + configured_path = Path(configured) + if configured_path.is_file(): + candidates.append(configured_path) + elif configured_path.is_dir(): + candidates.append(configured_path / library_name) + + for directory in ( + "/usr/lib", + "/usr/lib/aarch64-linux-gnu", + "/usr/lib/x86_64-linux-gnu", + "/usr/local/lib", + ): + candidates.append(Path(directory) / library_name) + + for candidate in candidates: + if candidate.is_file(): + try: + ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL) + break + except OSError as exc: + raise RuntimeError( + f"Found {candidate}, but it could not be loaded: {exc}. " + "Install the matching HailoRT native runtime and PCIe driver." + ) from exc + + try: + from hailo_platform import ( + HEF, + ConfigureParams, + FormatType, + HailoStreamInterface, + InferVStreams, + InputVStreamParams, + OutputVStreamParams, + VDevice, + ) + except ModuleNotFoundError as exc: + raise RuntimeError( + "HailoRT Python bindings are not installed. Install the matching " + "HailoRT Python wheel (5.3.0) on the target system." + ) from exc + except ImportError as exc: + message = str(exc) + if "libhailort.so" in message: + raise RuntimeError( + "HailoRT Python bindings are installed, but the native " + f"{library_name} library is not available to the dynamic linker. " + "Install the matching HailoRT runtime package, or set " + "HAILORT_LIB_PATH to the directory/file containing " + f"{library_name}. The Python wheel alone is not sufficient." + ) from exc + raise RuntimeError(f"HailoRT could not be imported: {exc}") from exc + + return { + "ConfigureParams": ConfigureParams, + "FormatType": FormatType, + "HEF": HEF, + "HailoStreamInterface": HailoStreamInterface, + "InputVStreamParams": InputVStreamParams, + "InferVStreams": InferVStreams, + "OutputVStreamParams": OutputVStreamParams, + "VDevice": VDevice, + } + + class HailoAnomalyRuntime: """Run a complete PaDiM or PatchCore HEF through HailoRT.""" @@ -26,59 +104,30 @@ def __init__( hef_path: str | Path, input_size: Tuple[int, int] = (224, 224), input_dtype: np.dtype = np.float32, + mean: Tuple[float, float, float] = (0.485, 0.456, 0.406), + std: Tuple[float, float, float] = (0.229, 0.224, 0.225), ) -> 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, - } + api = _load_hailort() + self._api = api 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.mean = np.asarray(mean, dtype=np.float32).reshape(1, 1, 3) + self.std = np.asarray(std, dtype=np.float32).reshape(1, 1, 3) + self.device = api["VDevice"]() + self.hef = api["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 + + input_infos = self.hef.get_input_vstream_infos() + if not input_infos: + raise ValueError(f"The HEF has no input stream: {self.hef_path}") + self.input_name = input_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)) @@ -89,43 +138,64 @@ def __init__( ) 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 _prepare_input(self, batch) -> np.ndarray: + """Convert AnomaVision NCHW input to the HEF's NHWC input.""" + if isinstance(batch, Image.Image): + array = np.asarray(batch.convert("RGB"), dtype=np.uint8) + return self._preprocess_raw_hwc(array)[None] - def predict( - self, image: Image.Image | np.ndarray | str | Path - ) -> Dict[str, np.ndarray]: - """Run one image and return complete image and localization outputs. + array = np.asarray(batch) + if array.ndim == 4: + if array.shape[0] != 1: + raise ValueError("HailoBackend currently supports batch size 1") + array = array[0] + + if array.ndim != 3: + raise ValueError("batch must be an HxWx3, 3xHxW, or single-image batch") + + if array.shape[0] == 3: + if array.shape[1:] != self.input_size: + raise ValueError( + f"Hailo input must be {self.input_size}, got {array.shape[1:]}" + ) + return np.ascontiguousarray( + np.transpose(array, (1, 2, 0)), dtype=self.input_dtype + ) - Args: - image: An image path, PIL RGB image, or HxWx3 RGB array. + if array.shape[2] == 3: + if np.issubdtype(array.dtype, np.integer): + return self._preprocess_raw_hwc(array) + if array.shape[:2] != self.input_size: + raise ValueError( + f"Hailo input must be {self.input_size}, got {array.shape[:2]}" + ) + return np.ascontiguousarray(array, dtype=self.input_dtype) + + raise ValueError("batch must be RGB with three channels") + + def _preprocess_raw_hwc(self, array: np.ndarray) -> np.ndarray: + """Preprocess a raw uint8 HWC RGB image exactly once.""" + image = Image.fromarray(array.astype(np.uint8), mode="RGB").resize( + (self.input_size[1], self.input_size[0]), Image.Resampling.BILINEAR + ) + normalized = np.asarray(image, dtype=np.float32) / 255.0 + normalized = (normalized - self.mean) / self.std + return np.ascontiguousarray(normalized, dtype=self.input_dtype) - Returns: - A mapping containing ``image_scores`` and ``score_map`` arrays. - """ + def predict(self, image) -> Dict[str, np.ndarray]: + """Run one image and return the HEF's complete anomaly outputs.""" api = self._api input_params = api["InputVStreamParams"].make( - self.network_group, quantized=False, format_type=api["FormatType"].FLOAT32 + 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 + self.network_group, + quantized=False, + format_type=api["FormatType"].FLOAT32, ) - tensor = self._preprocess(image) + tensor = self._prepare_input(image) with self.network_group.activate(self.network_group_params): with api["InferVStreams"]( self.network_group, input_params, output_params @@ -162,28 +232,12 @@ def __init__( 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) + """Run a single image through the HEF using the common backend contract.""" + result = self.runtime.predict(batch) 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.""" + """Warm up the device with a supplied preprocessed image batch.""" if batch is None: raise ValueError("Hailo warmup requires a sample image batch") for _ in range(max(1, int(runs))): diff --git a/anomavision/quantize/model/backends/hef/efficientad.py b/anomavision/quantize/model/backends/hef/efficientad.py new file mode 100644 index 0000000..4629bd7 --- /dev/null +++ b/anomavision/quantize/model/backends/hef/efficientad.py @@ -0,0 +1,48 @@ +"""Fixed-shape EfficientAD graph for Hailo end-to-end deployment.""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class EfficientADHailoGraph(nn.Module): + """Export-friendly EfficientAD inference graph for Hailo.""" + + def __init__(self, model: nn.Module, input_size: Tuple[int, int] = (224, 224)) -> None: + super().__init__() + if tuple(input_size) != (224, 224): + raise ValueError("Hailo EfficientAD export currently requires input_size=(224, 224)") + + self.input_size = (224, 224) + self.teacher = model.teacher.eval() + self.student = model.student.eval() + self.register_buffer("map_mean", model.map_mean.detach().float().clone()) + map_std = model.map_std.detach().float().clone().clamp_min(1e-6) + self.register_buffer("map_inv_std", map_std.reciprocal()) + + self.channel_mean = nn.Conv2d(112, 1, kernel_size=1, bias=False) + with torch.no_grad(): + self.channel_mean.weight.fill_(1.0 / 112.0) + self.channel_mean.weight.requires_grad_(False) + + def forward(self, image: torch.Tensor): + teacher_features = self.teacher(image) + student_features = self.student(image) + diff = student_features - teacher_features + squared = diff * diff + + # Reduce the 112 feature channels with a fixed 1x1 convolution. + raw = self.channel_mean(squared) + raw = F.interpolate( + raw, size=(224, 224), mode="bilinear", align_corners=False + ) + normalized = (raw - self.map_mean.unsqueeze(0)) * self.map_inv_std.unsqueeze(0) + + image_scores = F.max_pool2d( + normalized, kernel_size=(224, 224), stride=(224, 224) + ).flatten(1) + return image_scores, normalized.squeeze(1) diff --git a/anomavision/quantize/model/backends/hef/exporter.py b/anomavision/quantize/model/backends/hef/exporter.py index ec0fbac..9ec76cf 100644 --- a/anomavision/quantize/model/backends/hef/exporter.py +++ b/anomavision/quantize/model/backends/hef/exporter.py @@ -1,11 +1,9 @@ """Export complete AnomaVision anomaly graphs for Hailo Dataflow Compiler. -This module produces an ONNX graph containing the entire selected algorithm. The -actual INT8 quantization and HEF generation are delegated to the Hailo SDK when -it is installed on the development host. On a machine without the Hailo SDK, the -command still creates the graph and a calibration manifest, then exits with a -clear hardware-toolchain instruction instead of silently producing a partial -CPU artifact. +This module produces a fixed-shape, end-to-end ONNX graph for PaDiM, PatchCore, +or EfficientAD and prepares representative calibration tensors in the format +expected by the Hailo Dataflow Compiler. The Hailo SDK performs the actual INT8 +optimization and HEF compilation. """ from __future__ import annotations @@ -15,28 +13,53 @@ import shlex import subprocess from pathlib import Path -from typing import Any, Dict, Iterable, List, Tuple +from typing import Any, Tuple +import numpy as np import torch from PIL import Image -from .graphs import ( - PadimEndToEndGraph, - PatchCoreEndToEndGraph, - exportable_output_names, -) +from .efficientad import EfficientADHailoGraph +from .graphs import PadimEndToEndGraph, exportable_output_names +from .patchcore import PatchCoreHailoGraph -def _load_artifact(path: Path) -> Dict[str, Any]: - artifact = torch.load(path, map_location="cpu", weights_only=False) - if not isinstance(artifact, dict): - raise ValueError(f"Expected a dictionary artifact, got {type(artifact)!r}") - return artifact +def _load_artifact(path: Path) -> Any: + """Load an AnomaVision deployment artifact from disk.""" + return torch.load(path, map_location="cpu", weights_only=False) -def _build_graph(algorithm: str, artifact: Dict[str, Any], input_size: Tuple[int, int]): +def _build_efficientad_graph(artifact: Any, input_size: Tuple[int, int]): + """Build the fixed-shape EfficientAD Hailo graph.""" + from anomavision.algorithm.efficientad.efficientad import EfficientAD + + if isinstance(artifact, EfficientAD): + model = artifact + elif isinstance(artifact, dict) and artifact.get("algorithm") == "efficientad": + model = EfficientAD( + device=torch.device("cpu"), + model_size=artifact.get("model_size", "s"), + pretrained_teacher=False, + threshold_quantile=artifact.get("threshold_quantile", 0.995), + ) + model.load_state_dict(artifact["model_state"], strict=True) + else: + raise ValueError( + "EfficientAD Hailo export requires an EfficientAD model or " + "an EfficientAD statistics artifact" + ) + + if not bool(model.trained.item()): + raise ValueError("EfficientAD artifact is not trained/calibrated") + return EfficientADHailoGraph(model, input_size=input_size) + + +def _build_graph(algorithm: str, artifact: Any, input_size: Tuple[int, int]): + """Build the fixed-shape end-to-end graph for the selected algorithm.""" algorithm = algorithm.lower() if algorithm == "padim": + if not isinstance(artifact, dict): + raise ValueError("PaDiM Hailo export requires a statistics artifact dictionary") required = {"backbone", "layer_indices", "channel_indices", "mean", "cov_inv"} missing = sorted(required.difference(artifact)) if missing: @@ -49,53 +72,104 @@ def _build_graph(algorithm: str, artifact: Dict[str, Any], input_size: Tuple[int cov_inv=artifact["cov_inv"], input_size=input_size, ) + if algorithm == "patchcore": - required = {"backbone", "layer_indices", "memory_bank"} - missing = sorted(required.difference(artifact)) - if missing: - raise ValueError( - f"PatchCore artifact is missing keys: {', '.join(missing)}" - ) - return PatchCoreEndToEndGraph( - backbone=str(artifact["backbone"]), - layer_indices=list(artifact["layer_indices"]), - memory_bank=artifact["memory_bank"], - patch_grid=artifact.get("patch_grid", 14), + if isinstance(artifact, dict): + required = {"backbone", "layer_indices", "memory_bank"} + missing = sorted(required.difference(artifact)) + if missing: + raise ValueError(f"PatchCore artifact is missing keys: {', '.join(missing)}") + backbone = str(artifact["backbone"]) + layer_indices = list(artifact["layer_indices"]) + memory_bank = artifact["memory_bank"] + patch_grid = int(artifact.get("patch_grid", 14)) + else: + try: + backbone = str(artifact.backbone) + layer_indices = list(artifact.layer_indices) + memory_bank = artifact.memory_bank + patch_grid = int(getattr(artifact, "patch_grid", 14)) + except AttributeError as exc: + raise ValueError( + "PatchCore artifact must be a PatchCore model or artifact dictionary" + ) from exc + return PatchCoreHailoGraph( + backbone=backbone, + layer_indices=layer_indices, + memory_bank=memory_bank, + patch_grid=patch_grid, input_size=input_size, ) - raise ValueError("algorithm must be 'padim' or 'patchcore'") + + if algorithm == "efficientad": + return _build_efficientad_graph(artifact, input_size) + + raise ValueError("algorithm must be 'padim', 'patchcore' or 'efficientad'") def _write_calibration_manifest( image_dir: Path, output_dir: Path, input_size: Tuple[int, int] ) -> Path: + """Create normalized calibration tensors and a JSON manifest.""" suffixes = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} paths = sorted(p for p in image_dir.rglob("*") if p.suffix.lower() in suffixes) if not paths: raise ValueError(f"No calibration images found in {image_dir}") + + calibration_dir = output_dir / "calibration_npy" + calibration_dir.mkdir(parents=True, exist_ok=True) manifest = output_dir / "calibration_manifest.json" + mean = (0.485, 0.456, 0.406) + std = (0.229, 0.224, 0.225) records = [] - for path in paths: + + for stale in calibration_dir.glob("*.npy"): + stale.unlink() + + for index, path in enumerate(paths[:1024]): with Image.open(path) as image: - image.convert("RGB").resize((input_size[1], input_size[0])) + image = image.convert("RGB").resize( + (input_size[1], input_size[0]), Image.Resampling.BILINEAR + ) + array = np.asarray(image, dtype=np.float32) / 255.0 + array = (array - np.asarray(mean, dtype=np.float32)) / np.asarray( + std, dtype=np.float32 + ) + np.save(calibration_dir / f"sample_{index:04d}.npy", array) records.append( { "path": str(path.resolve()), - "width": input_size[1], - "height": input_size[0], + "calibration": str((calibration_dir / f"sample_{index:04d}.npy").resolve()), + "shape": list(array.shape), + "normalized": True, + "normalization": { + "mean": list(mean), + "std": list(std), + "scale": "1/255 before mean/std", + }, } ) + manifest.write_text(json.dumps(records, indent=2), encoding="utf-8") return manifest +def _prepare_calibration( + image_dir: Path, output_dir: Path, input_size: Tuple[int, int] +) -> Tuple[Path, Path]: + """Create Hailo calibration tensors and a JSON manifest.""" + manifest = _write_calibration_manifest(image_dir, output_dir, input_size) + return manifest.parent / "calibration_npy", manifest + + def export_onnx( algorithm: str, artifact_path: Path, output_dir: Path, input_size: Tuple[int, int], - opset: int = 17, + opset: int = 13, ) -> Path: + """Export the complete anomaly detector as a fixed-shape ONNX graph.""" output_dir.mkdir(parents=True, exist_ok=True) artifact = _load_artifact(artifact_path) graph = _build_graph(algorithm, artifact, input_size).eval() @@ -117,6 +191,7 @@ def export_onnx( def _run_hailo_command(command: str, onnx_path: Path, output_dir: Path) -> None: + """Run an optional user-supplied Hailo SDK command template.""" rendered = command.format( onnx=shlex.quote(str(onnx_path)), output=shlex.quote(str(output_dir)) ) @@ -124,48 +199,33 @@ def _run_hailo_command(command: str, onnx_path: Path, output_dir: Path) -> None: def main() -> None: + """Export an end-to-end Hailo graph and prepare calibration data.""" parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--algorithm", choices=["padim", "patchcore"], required=True) + parser.add_argument( + "--algorithm", choices=["padim", "patchcore", "efficientad"], required=True + ) parser.add_argument("--artifact", type=Path, required=True) parser.add_argument("--calibration-dir", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--height", type=int, default=224) parser.add_argument("--width", type=int, default=224) parser.add_argument("--opset", type=int, default=13) - parser.add_argument( - "--hailo-command", - help=( - "Optional installed Hailo SDK command template. Use {onnx} and {output}; " - "the command must perform parse, calibration/optimization, and compile." - ), - ) + parser.add_argument("--hailo-command") args = parser.parse_args() + input_size = (args.height, args.width) onnx_path = export_onnx( args.algorithm, args.artifact, args.output_dir, input_size, args.opset ) - manifest = _write_calibration_manifest( + calibration_dir, manifest = _prepare_calibration( args.calibration_dir, args.output_dir, input_size ) - metadata = { - "algorithm": args.algorithm, - "quantization_scope": "end_to_end", - "graph_outputs": exportable_output_names(), - "input_size": list(input_size), - "onnx": str(onnx_path), - "calibration_manifest": str(manifest), - "hailo_compile_invoked": bool(args.hailo_command), - } - (args.output_dir / "hailo_export.json").write_text( - json.dumps(metadata, indent=2), encoding="utf-8" - ) + print(f"ONNX: {onnx_path}") + print(f"Calibration: {calibration_dir}") + print(f"Manifest: {manifest}") + if args.hailo_command: _run_hailo_command(args.hailo_command, onnx_path, args.output_dir) - else: - print("ONNX graph and calibration manifest created.") - print( - "No Hailo compiler was invoked; install the Hailo SDK and provide --hailo-command to create a quantized HEF." - ) if __name__ == "__main__": diff --git a/anomavision/quantize/model/backends/hef/patchcore.py b/anomavision/quantize/model/backends/hef/patchcore.py new file mode 100644 index 0000000..4f0bb78 --- /dev/null +++ b/anomavision/quantize/model/backends/hef/patchcore.py @@ -0,0 +1,88 @@ +"""Fixed-shape PatchCore graph for Hailo end-to-end deployment.""" + +from __future__ import annotations + +from typing import List, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from anomavision.algorithm.common.feature_extraction import ResnetEmbeddingsExtractor + + +class PatchCoreHailoGraph(nn.Module): + """Export-friendly PatchCore graph with the complete scoring pipeline. + + The graph keeps the patch scores in their native 14x14 spatial layout. + Hailo DFC does not support the global ``ReduceMax`` generated by + ``tensor.amax(dim=(1, 2, 3))``, so the image score is produced with a + 14x14 max-pooling operation instead. This is mathematically equivalent + for the fixed PatchCore score map and avoids the unsupported reduction. + """ + + def __init__( + self, + backbone: str, + layer_indices: List[int], + memory_bank: torch.Tensor, + patch_grid: int = 14, + input_size: Tuple[int, int] = (224, 224), + ) -> None: + super().__init__() + if patch_grid < 1: + raise ValueError("patch_grid must be positive") + if tuple(input_size) != (224, 224) or patch_grid != 14: + raise ValueError( + "Hailo PatchCore export currently requires " + "input_size=(224, 224) and patch_grid=14" + ) + + self.input_size = (224, 224) + self.patch_grid = 14 + self.layer_indices = list(layer_indices) + self.extractor = ResnetEmbeddingsExtractor(backbone, torch.device("cpu")) + + # Each normalized memory-bank vector becomes a 1x1 convolution filter. + # This computes feature-to-memory-bank cosine similarity while preserving + # the native 14x14 spatial layout and avoiding a patch-score reshape. + self.register_buffer("memory_bank", F.normalize(memory_bank.float(), dim=-1)) + + def forward(self, image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + embeddings, _, _ = self.extractor(image, layer_indices=self.layer_indices) + + # Fixed ResNet feature layout for the K26 PatchCore artifact. + # Keep this spatial instead of flattening to 196 patches. + features = embeddings.reshape(1, 56, 56, 64).permute(0, 3, 1, 2) + features = F.avg_pool2d(features, kernel_size=4, stride=4) + features = F.normalize(features, dim=1) + + # [1, 64, 14, 14] x [memory_bank, 64, 1, 1] + # -> [1, memory_bank_size, 14, 14]. + similarity = F.conv2d( + features, + self.memory_bank.unsqueeze(-1).unsqueeze(-1), + ) + + # Best memory-bank match for every spatial patch. + best_similarity = similarity.amax(dim=1, keepdim=True) + distances = torch.sqrt(torch.clamp(2.0 - 2.0 * best_similarity, min=0.0)) + + # Native spatial score map: [1, 1, 14, 14] -> [1, 224, 224]. + score_map = F.interpolate( + distances, + size=(224, 224), + mode="bilinear", + align_corners=False, + ).squeeze(1) + + # PatchCore image score is the maximum patch distance. Using a fixed + # 14x14 MaxPool is equivalent to a global max over the fixed score map, + # while exporting as a Hailo-supported pooling operation. + image_scores = F.max_pool2d( + distances, + kernel_size=(14, 14), + stride=(14, 14), + ) + + return image_scores, score_map diff --git a/config.yml b/config.yml index 9f7960e..d3f36a3 100644 --- a/config.yml +++ b/config.yml @@ -1,7 +1,7 @@ # ========================= # Dataset / preprocessing (shared by train, detect, eval, stream) # ========================= -dataset_path: "D:/01-DATA" +dataset_path: "/root/dataset" class_name: "bottle" resize: [224, 224] crop_size: diff --git a/docs/hailo_quant.md b/docs/hailo_quant.md new file mode 100644 index 0000000..02984dd --- /dev/null +++ b/docs/hailo_quant.md @@ -0,0 +1,268 @@ +# Hailo Quantization + +
+
+