Update app.py

app.py

@@ -41,25 +41,22 @@ vit.load_state_dict(torch.load("vit_epoch_2.pth", map_location=torch.device("cpu
 ### 3. Prediction function ###
 
 def predict(img) -> Tuple[Dict[str, float], float]:
-    """Transforms and performs a prediction on img and returns prediction and time taken."""
     from PIL import UnidentifiedImageError
 
     try:
-        # Convert ndarray to PIL if needed
+        # Convert ndarray to PIL.Image if needed
         if isinstance(img, np.ndarray):
-            img = Image.fromarray(img)
+            img = Image.fromarray(img.astype("uint8"))  # Ensure correct dtype
 
-        # Catch bad image input
+        # Ensure image is in RGB mode
         if img.mode != "RGB":
             img = img.convert("RGB")
 
-        # Start timer
         start_time = timer()
 
-        # Transform and add batch dimension
+        # Apply transforms (expects a PIL image)
         img_tensor = vit_transforms(img).unsqueeze(0)
 
-        # Inference
         vit.eval()
         with torch.inference_mode():
             pred_probs = torch.softmax(vit(img_tensor), dim=1)
@@ -76,6 +73,7 @@ def predict(img) -> Tuple[Dict[str, float], float]:
     except (UnidentifiedImageError, TypeError, ValueError) as e:
         return {"Error": f"Invalid image input: {str(e)}"}, 0.0
 
+
 ### 4. Gradio app setup ###
 
 # Title, description, and article text