Update app.py

app.py

@@ -4,23 +4,36 @@ from PIL import Image
 from io import BytesIO
 from transformers import pipeline
 import torch
+from torchvision import transforms
 
 # Cache the model loading
 model_id = "Honey-Bee-Society/honeybee_bumblebee_vespidae_resnet50"
 classifier = pipeline("image-classification", model=model_id, device=0 if torch.cuda.is_available() else -1)
 
+# Define the same preprocessing steps as in the training script
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
 def classify_image_from_url(image_url: str):
     """
-    Downloads an image from a public URL and runs it through
+    Downloads an image from a public URL, preprocesses it, and runs it through
     the ResNet-50 fine-tuned image-classification pipeline, returning the top predictions.
     """
     try:
+        # Download the image
         response = requests.get(image_url)
         response.raise_for_status()
         image = Image.open(BytesIO(response.content)).convert("RGB")
 
+        # Apply the same preprocessing as in the training script
+        image_tensor = preprocess(image).unsqueeze(0)  # Add batch dimension
+
         # Run inference
-        results = classifier(image)
+        results = classifier(image_tensor)
 
         # Format scores to remove scientific notation
         for r in results: