Upload inference.py

inference.py

@@ -27,6 +27,7 @@ import platform
 import sys
 from pathlib import Path
 
+import numpy as np
 import pandas as pd
 import torch
 import torch.nn as nn
@@ -308,3 +309,24 @@ class ModelInference:
             class_names[class_id_str] = class_name
 
         return class_names
+
+    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]['Code']
+                classifications.append([pred_class, float(p[i])])
+            results.append(classifications)
+        return results