RFC: metadata.json in exported models

[ashwinvaidya17]

📝 Description

• Add additional inputs to the model graph to make the threshold configurable while retaining the post-processing logic

Proposal reproducer

Details

"""Reproducer: Sensitivity as ONNX graph input (Strategy C — re-center normalization).

Validates that:
1. A simplified PostProcessor with sensitivity as forward() args traces to ONNX
2. Sensitivity shifts the normalization center (effective threshold), not the decision boundary
3. Runtime override of sensitivity changes both normalized scores AND labels
4. ONNX initializers make sensitivity inputs optional (defaults baked in)
5. OpenVINO conversion works (if openvino is available)
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch
from onnx import numpy_helper
from torch import nn


class MiniPostProcessor(nn.Module):
    """Minimal PostProcessor using Strategy C: sensitivity re-centers normalization."""

    def __init__(self, image_sensitivity: float = 0.5) -> None:
        super().__init__()
        self.image_sensitivity = image_sensitivity

        self.register_buffer("_image_threshold", torch.tensor(0.42))
        self.register_buffer("image_min", torch.tensor(0.02))
        self.register_buffer("image_max", torch.tensor(0.91))

    def forward(
        self,
        pred_score: torch.Tensor,
        image_sensitivity: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        eff_threshold = self._effective_threshold(
            self._image_threshold, self.image_min, self.image_max, image_sensitivity,
        )
        normalized = self._normalize(
            pred_score, self.image_min, self.image_max, eff_threshold,
        )
        pred_label = (normalized > 0.5).float()
        return normalized, pred_label

    @staticmethod
    def _effective_threshold(
        threshold: torch.Tensor,
        norm_min: torch.Tensor,
        norm_max: torch.Tensor,
        sensitivity: torch.Tensor,
    ) -> torch.Tensor:
        return threshold + (norm_max - norm_min) * (0.5 - sensitivity)

    @staticmethod
    def _normalize(
        preds: torch.Tensor,
        norm_min: torch.Tensor,
        norm_max: torch.Tensor,
        threshold: torch.Tensor,
    ) -> torch.Tensor:
        preds = ((preds - threshold) / (norm_max - norm_min)) + 0.5
        return preds.clamp(min=0, max=1)


class MiniModel(nn.Module):
    """Minimal model: identity backbone + post-processor."""

    def __init__(self) -> None:
        super().__init__()
        self.backbone = nn.Identity()
        self.post_processor = MiniPostProcessor(image_sensitivity=0.5)

    def forward(
        self,
        x: torch.Tensor,
        image_sensitivity: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        raw_score = self.backbone(x)
        return self.post_processor(raw_score, image_sensitivity)


def add_initializers(onnx_path: Path, defaults: dict[str, float]) -> None:
    """Post-process ONNX graph: add initializers so sensitivity inputs are optional."""
    model = onnx.load(str(onnx_path))
    for name, default_value in defaults.items():
        init = numpy_helper.from_array(
            np.array(default_value, dtype=np.float32),
            name=name,
        )
        model.graph.initializer.append(init)
    onnx.save(model, str(onnx_path))


def test_onnx_export_and_override() -> None:
    model = MiniModel()
    model.eval()

    raw_score = torch.tensor([0.6])
    default_sens = torch.tensor(0.5)

    with torch.no_grad():
        torch_norm, torch_label = model(raw_score, default_sens)
    print(f"PyTorch (sens=0.5): normalized={torch_norm.item():.4f}, label={torch_label.item():.0f}")

    with tempfile.TemporaryDirectory() as tmpdir:
        onnx_path = Path(tmpdir) / "model.onnx"

        # --- Export to ONNX with 2 inputs ---
        torch.onnx.export(
            model,
            (raw_score, default_sens),
            str(onnx_path),
            input_names=["input", "image_sensitivity"],
            output_names=["pred_score", "pred_label"],
            dynamic_axes={"input": {0: "batch"}},
            opset_version=14,
            dynamo=False,
        )
        print(f"\n[OK] ONNX export succeeded: {onnx_path.name}")

        # --- Add initializer for image_sensitivity (makes it optional) ---
        add_initializers(onnx_path, {"image_sensitivity": 0.5})
        print("[OK] Added ONNX initializer for image_sensitivity (default=0.5)")

        # --- Validate ONNX model ---
        onnx.checker.check_model(onnx.load(str(onnx_path)))
        print("[OK] ONNX model validation passed")

        # --- Test 1: Provide both inputs (override sensitivity) ---
        sess = ort.InferenceSession(str(onnx_path))

        def sens_arr(val: float) -> np.ndarray:
            return np.array(val, dtype=np.float32)

        result = sess.run(None, {
            "input": np.array([0.6], dtype=np.float32),
            "image_sensitivity": sens_arr(0.5),
        })
        print(f"\n--- Test 1: Provide sensitivity=0.5 (same as default) ---")
        print(f"  pred_score={result[0].item():.4f}, pred_label={result[1].item():.0f}")
        assert np.isclose(result[0].item(), torch_norm.item(), atol=1e-5), "Score mismatch!"
        print("  [OK] Matches PyTorch output")

        # --- Test 2: Override with higher sensitivity (score shifts up) ---
        result_high = sess.run(None, {
            "input": np.array([0.6], dtype=np.float32),
            "image_sensitivity": sens_arr(0.9),
        })
        print(f"\n--- Test 2: Override sensitivity=0.9 (more sensitive) ---")
        print(f"  pred_score={result_high[0].item():.4f}, pred_label={result_high[1].item():.0f}")
        # eff_threshold = 0.42 + 0.89*(-0.4) = 0.064, norm ≈ 1.0 (clamped)
        assert result_high[0].item() > torch_norm.item(), "Score should increase with higher sensitivity"
        assert result_high[1].item() == 1.0, "Expected label=1 with high sensitivity"
        print("  [OK] Higher sensitivity -> higher score, label=1")

        # --- Test 3: Override with lower sensitivity (score shifts down) ---
        result_low = sess.run(None, {
            "input": np.array([0.6], dtype=np.float32),
            "image_sensitivity": sens_arr(0.1),
        })
        print(f"\n--- Test 3: Override sensitivity=0.1 (less sensitive) ---")
        print(f"  pred_score={result_low[0].item():.4f}, pred_label={result_low[1].item():.0f}")
        # eff_threshold = 0.42 + 0.89*(0.4) = 0.776, norm ≈ 0.302
        assert result_low[0].item() < torch_norm.item(), "Score should decrease with lower sensitivity"
        assert result_low[1].item() == 0.0, "Expected label=0 with low sensitivity"
        print("  [OK] Lower sensitivity -> lower score, label=0")

        # --- Test 4: Omit sensitivity (use baked-in default from initializer) ---
        result_default = sess.run(None, {
            "input": np.array([0.6], dtype=np.float32),
        })
        print(f"\n--- Test 4: Omit sensitivity (use ONNX initializer default=0.5) ---")
        print(f"  pred_score={result_default[0].item():.4f}, pred_label={result_default[1].item():.0f}")
        assert np.isclose(result_default[0].item(), torch_norm.item(), atol=1e-5), "Default mismatch!"
        print("  [OK] Omitted input uses initializer default, matches PyTorch")

        # --- Test 5: OpenVINO with initializers (optional inputs) ---
        try:
            import openvino as ov

            core = ov.Core()
            ov_model = core.read_model(str(onnx_path))
            compiled = core.compile_model(ov_model, "CPU")
            input_names = [inp.any_name for inp in compiled.inputs]
            print(f"\n--- Test 5a: OpenVINO (ONNX with initializers) ---")
            print(f"  Model inputs: {input_names}")

            if "image_sensitivity" in input_names:
                ov_result = compiled({
                    "input": np.array([0.6], dtype=np.float32),
                    "image_sensitivity": np.array(0.9, dtype=np.float32),
                })
                ov_score = list(ov_result.values())[0].item()
                print(f"  sens=0.9: pred_score={ov_score:.4f}")
                print("  [OK] OpenVINO keeps sensitivity as overridable input")
            else:
                print("  [INFO] OpenVINO folded initializer into constant (expected).")
                print("  Sensitivity is baked in, not overridable via this ONNX graph.")
                ov_result = compiled({"input": np.array([0.6], dtype=np.float32)})
                ov_score = list(ov_result.values())[0].item()
                print(f"  Default result: pred_score={ov_score:.4f}")
                assert np.isclose(ov_score, torch_norm.item(), atol=1e-3)
                print("  [OK] Default matches PyTorch output")

        except ImportError:
            print("\n--- Test 5a: OpenVINO (SKIPPED - not installed) ---")

        # --- Test 6: OpenVINO with NO initializers (always-required inputs) ---
        print(f"\n--- Test 6: OpenVINO without initializers (required inputs) ---")
        onnx_no_init_path = Path(tmpdir) / "model_no_init.onnx"

        torch.onnx.export(
            model,
            (raw_score, default_sens),
            str(onnx_no_init_path),
            input_names=["input", "image_sensitivity"],
            output_names=["pred_score", "pred_label"],
            dynamic_axes={"input": {0: "batch"}},
            opset_version=14,
            dynamo=False,
        )

        try:
            import openvino as ov

            core = ov.Core()
            ov_model = core.read_model(str(onnx_no_init_path))
            compiled = core.compile_model(ov_model, "CPU")
            input_names = [inp.any_name for inp in compiled.inputs]
            print(f"  Model inputs: {input_names}")

            # Default sensitivity
            ov_result = compiled({
                "input": np.array([0.6], dtype=np.float32),
                "image_sensitivity": np.array(0.5, dtype=np.float32),
            })
            ov_score = list(ov_result.values())[0].item()
            ov_label = list(ov_result.values())[1].item()
            print(f"  sens=0.5: pred_score={ov_score:.4f}, pred_label={ov_label:.0f}")
            assert np.isclose(ov_score, torch_norm.item(), atol=1e-3)
            print("  [OK] Matches PyTorch")

            # Override sensitivity
            ov_result_high = compiled({
                "input": np.array([0.6], dtype=np.float32),
                "image_sensitivity": np.array(0.9, dtype=np.float32),
            })
            ov_label_high = list(ov_result_high.values())[1].item()
            print(f"  sens=0.9: pred_label={ov_label_high:.0f}")
            assert ov_label_high == 1.0
            print("  [OK] Override works in OpenVINO")

            ov_result_low = compiled({
                "input": np.array([0.6], dtype=np.float32),
                "image_sensitivity": np.array(0.1, dtype=np.float32),
            })
            ov_label_low = list(ov_result_low.values())[1].item()
            print(f"  sens=0.1: pred_label={ov_label_low:.0f}")
            assert ov_label_low == 0.0
            print("  [OK] Lower sensitivity -> label=0")

        except ImportError:
            print("  (SKIPPED - OpenVINO not installed)")

    print("\n========================================")
    print("All tests passed!")
    print("========================================")


if __name__ == "__main__":
    test_onnx_export_and_override()

[Copilot]

Pull request overview

Adds an RFC to anomalib describing a deploy-time configurable threshold/sensitivity design by passing sensitivity as an exported graph input and introducing a versioned metadata.json sidecar + migration strategy.

Changes:

• Proposes adding image_sensitivity/pixel_sensitivity as required inputs to exported ONNX/OpenVINO graphs (while keeping post-processing inside the graph).
• Proposes a metadata.json contract to document preprocess transforms and store default sensitivities for inferencers.
• Proposes a schema versioning + migration mechanism (Alembic-inspired) for exported metadata.