Upload inference.py

inference.py

@@ -28,6 +28,7 @@ import pathlib
 import platform
 from pathlib import Path
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -304,3 +305,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 = [
+                [self.classes[i], float(p[i])]
+                for i in range(len(self.classes))
+            ]
+            results.append(classifications)
+        return results