Upload inference.py

inference.py

@@ -24,6 +24,7 @@ import pathlib
 import platform
 from pathlib import Path
 
+import numpy as np
 import pandas as pd
 import torch
 import torch.nn as nn
@@ -311,3 +312,24 @@ class ModelInference:
 
         except Exception as e:
             raise RuntimeError(f"Failed to extract class names: {e}") from e
+
+    def get_tensor(self, crop: Image.Image):
+        """Preprocess a crop into a numpy array for batch inference."""
+        tensor = self.preprocess(crop)
+        return tensor.numpy()
+
+    def classify_batch(self, batch):
+        """Run inference on a batch of preprocessed numpy arrays."""
+        tensor = torch.from_numpy(batch).to(self.device)
+        with torch.no_grad():
+            output = self.model(tensor)
+            probs = F.softmax(output, dim=1).cpu().numpy()
+
+        results = []
+        for p in probs:
+            classifications = []
+            for i in range(len(p)):
+                pred_class = self.classes.iloc[i].values[1]  # 'Code' column
+                classifications.append([pred_class, float(p[i])])
+            results.append(classifications)
+        return results