Upload folder using huggingface_hub

model/config.json

@@ -0,0 +1,12 @@
+{
+  "architecture": "VAAS",
+  "version": "v1",
+  "alpha": 0.5,
+  "input_size": [
+    224,
+    224
+  ],
+  "px_checkpoint": "px_model.pth",
+  "fx_backbone": "google/vit-base-patch16-224",
+  "px_backbone": "nvidia/segformer-b1"
+}

model/px_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c8f0aea456a5175db54de8c8483ddd5b001e816fcac249d3968dcd7549603fb
+size 54798133

model/ref_stats.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:09287fa16965a465e7b71a19d43c9eca95f2a086af4428d47e963ff230da432e
+size 1845

vaas/inference/pipeline.py

@@ -0,0 +1,173 @@
+import os
+from typing import Dict, Union
+
+import torch
+from PIL import Image
+import torchvision.transforms as T
+
+from vaas.fx.fx_model import FxViT
+from vaas.px.px_model import PatchConsistencySegformer
+from vaas.fusion.hybrid_score import compute_scores
+from vaas.inference.utils import load_ref_stats, load_px_checkpoint
+
+import warnings
+warnings.filterwarnings("ignore")
+
+from transformers.utils import logging as hf_logging
+hf_logging.set_verbosity_error()
+
+from huggingface_hub import hf_hub_download
+
+
+class VAASPipeline:
+    def __init__(
+        self,
+        model_px,
+        model_fx,
+        mu_ref,
+        sigma_ref,
+        device,
+        transform,
+        alpha=0.5,
+    ):
+        self.device = device
+
+        self.model_px = model_px.to(device)
+        self.model_fx = model_fx.to(device)
+
+        self.mu_ref = (
+            mu_ref.to(device) if torch.is_tensor(mu_ref)
+            else torch.tensor(mu_ref, device=device)
+        )
+
+        self.sigma_ref = (
+            sigma_ref.to(device) if torch.is_tensor(sigma_ref)
+            else torch.tensor(sigma_ref, device=device)
+        )
+
+        self.transform = transform
+        self.alpha = alpha
+
+        self.model_px.eval()
+        self.model_fx.eval()
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint_dir: str,
+        device: Union[str, torch.device] = "cpu",
+        alpha: float = 0.5,
+    ):
+        if isinstance(device, str):
+            device = torch.device(device)
+
+        model_px = PatchConsistencySegformer()
+        model_fx = FxViT()
+
+        model_fx.eval()
+        model_px.eval()
+
+        load_px_checkpoint(model_px, checkpoint_dir)
+        model_px = model_px.to(device)
+        model_fx = model_fx.to(device)
+
+        mu_ref, sigma_ref = load_ref_stats(checkpoint_dir)
+
+        transform = T.Compose(
+            [
+                T.Resize((224, 224)),
+                T.ToTensor(),
+                T.Normalize(
+                    mean=(0.485, 0.456, 0.406),
+                    std=(0.229, 0.224, 0.225),
+                ),
+            ]
+        )
+
+        return cls(
+            model_px=model_px,
+            model_fx=model_fx,
+            mu_ref=mu_ref,
+            sigma_ref=sigma_ref,
+            device=device,
+            transform=transform,
+            alpha=alpha,
+        )
+
+    
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        device: str = "cpu",
+        alpha: float = 0.5,
+    ):
+        px_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="model/px_model.pth",
+        )
+        ref_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="model/ref_stats.pth",
+        )
+
+        model_px = PatchConsistencySegformer()
+        state = torch.load(px_path, map_location="cpu")
+        model_px.load_state_dict(state)
+
+        ref = torch.load(ref_path, map_location="cpu")
+        mu_ref = ref["mu_ref"]
+        sigma_ref = ref["sigma_ref"]
+
+        model_fx = FxViT()
+
+        transform = T.Compose(
+            [
+                T.Resize((224, 224)),
+                T.ToTensor(),
+                T.Normalize(
+                    mean=(0.485, 0.456, 0.406),
+                    std=(0.229, 0.224, 0.225),
+                ),
+            ]
+        )
+
+        return cls(
+            model_px=model_px,
+            model_fx=model_fx,
+            mu_ref=mu_ref,
+            sigma_ref=sigma_ref,
+            device=device,
+            transform=transform,
+            alpha=alpha,
+        )
+
+
+    @torch.no_grad()
+    def __call__(self, image: Union[str, Image.Image]) -> Dict[str, Union[float, "np.ndarray"]]:
+        if isinstance(image, str):
+            image = Image.open(image).convert("RGB")
+
+        s_f, s_p, s_h, anomaly_map = compute_scores(
+            img=image,
+            mask=None,
+            model_px=self.model_px,
+            vit_model=self.model_fx,
+            mu_ref=self.mu_ref,
+            sigma_ref=self.sigma_ref,
+            transform=self.transform,
+            alpha=self.alpha,
+        )
+
+        if torch.is_tensor(anomaly_map):
+            anomaly_map = anomaly_map.detach().cpu().numpy()
+
+
+        return {
+            "S_F": float(s_f),
+            "S_P": float(s_p),
+            "S_H": float(s_h),
+            "anomaly_map": anomaly_map,
+        }
+

vaas/inference/utils.py

@@ -0,0 +1,21 @@
+import os
+import torch
+import json
+
+
+def load_px_checkpoint(model, checkpoint_dir):
+    ckpt_path = os.path.join(checkpoint_dir, "best_model_px.pth")
+    if not os.path.exists(ckpt_path):
+        raise FileNotFoundError(f"Missing checkpoint: {ckpt_path}")
+
+    state = torch.load(ckpt_path, map_location="cpu")
+    model.load_state_dict(state["model_state_dict"])
+
+
+def load_ref_stats(checkpoint_dir):
+    ref_path = os.path.join(checkpoint_dir, "ref_stats.pth")
+    if not os.path.exists(ref_path):
+        raise FileNotFoundError(f"Missing reference stats: {ref_path}")
+
+    stats = torch.load(ref_path, map_location="cpu")
+    return stats["mu_ref"], stats["sigma_ref"]