Upload inference.py

inference.py

@@ -215,3 +215,32 @@ class ModelInference:
         return {
             str(i + 1): name for i, name in enumerate(self.class_names)
         }
+
+    def get_tensor(self, crop: Image.Image):
+        """Preprocess a crop into a numpy array for batch inference."""
+        if crop.mode != "RGB":
+            crop = crop.convert("RGB")
+
+        img_tensor = TF.pil_to_tensor(crop)
+        img_tensor = TF.convert_image_dtype(img_tensor, torch.float32)
+        img_tensor = TF.resize(
+            img_tensor, [IMG_SIZE, IMG_SIZE], antialias=False
+        )
+        img_tensor = TF.convert_image_dtype(img_tensor, torch.uint8)
+        return img_tensor.permute(1, 2, 0).numpy().astype("float32") / 255.0
+
+    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():
+            logits = self.model(tensor)
+            probs = F.softmax(logits, dim=1).cpu().numpy()
+
+        results = []
+        for p in probs:
+            classifications = [
+                [self.class_names[i], float(p[i])]
+                for i in range(len(self.class_names))
+            ]
+            results.append(classifications)
+        return results