Upload inference.py

inference.py

@@ -20,6 +20,7 @@ import pickle
 import platform
 import pathlib
 
+import numpy as np
 import torch
 from torchvision import transforms
 from PIL import Image, ImageFile
@@ -186,3 +187,24 @@ class ModelInference:
             class_names[class_id_str] = class_label
 
         return class_names
+
+    def get_tensor(self, crop: Image.Image):
+        """Preprocess a crop into a numpy array for batch inference."""
+        tensor = self.transform(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 = torch.nn.functional.softmax(output, dim=1).cpu().numpy()
+
+        results = []
+        for p in probs:
+            classifications = [
+                [self.class_labels[i], float(p[i])]
+                for i in range(len(self.class_labels))
+            ]
+            results.append(classifications)
+        return results