Source code

src/models/inference.py

@@ -0,0 +1,58 @@
+import cv2
+import numpy as np
+import torch
+import torchvision.transforms.v2 as T
+from PIL import Image
+from typing import List, Tuple
+
+from src.models.regression_model import SpindaRegressionModel
+
+
+class SpindaInference:
+    """Loads the trained model and predicts spot coordinates from an image crop."""
+
+    def __init__(self, model_path: str = "models/best_spinda_model.pth"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = SpindaRegressionModel(pretrained=False)
+        self.model.load_state_dict(
+            torch.load(model_path, map_location=self.device)
+        )
+        self.model.to(self.device).eval()
+
+        self.transform = T.Compose([
+            T.Resize((128, 128)),
+            T.ToDtype(torch.float32, scale=True),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def predict(self, image_path: str) -> Tuple[List[int], str]:
+        """Predict the 8 grid coordinates and return them with a fingerprint string.
+
+        Returns:
+            grid_coords: list of 8 integers in [0, 15]
+            fingerprint: "X1-Y1-X2-Y2-X3-Y3-X4-Y4"
+        """
+        img_bgr = cv2.imread(image_path)
+        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        img_tensor = torch.from_numpy(np.array(img_rgb)).permute(2, 0, 1)
+        img_tensor = self.transform(img_tensor).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(img_tensor)   # (1, 8, 16)
+
+        grid_coords = logits.argmax(dim=2).squeeze(0).cpu().tolist()  # [8]
+        fingerprint = "-".join(f"{c:02d}" for c in grid_coords)
+
+        return grid_coords, fingerprint
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python src/models/inference.py <image_path>")
+    else:
+        inf = SpindaInference()
+        coords, fingerprint = inf.predict(sys.argv[1])
+        print(f"Predicted Grid Coordinates: {coords}")
+        print(f"Visual Fingerprint: {fingerprint}")