change

app.py

@@ -42,30 +42,39 @@ vit.load_state_dict(torch.load("vit_epoch_2.pth", map_location=torch.device("cpu
 
 def predict(img) -> Tuple[Dict[str, float], float]:
     """Transforms and performs a prediction on img and returns prediction and time taken."""
-    # Ensure the image is a PIL image
-    if isinstance(img, np.ndarray):
-        img = Image.fromarray(img)
+    from PIL import UnidentifiedImageError
 
-    # Start the timer
-    start_time = timer()
+    try:
+        # Convert ndarray to PIL if needed
+        if isinstance(img, np.ndarray):
+            img = Image.fromarray(img)
 
-    # Transform the image and add batch dimension
-    img = vit_transforms(img).unsqueeze(0)
+        # Catch bad image input
+        if img.mode != "RGB":
+            img = img.convert("RGB")
 
-    # Run inference
-    vit.eval()
-    with torch.inference_mode():
-        pred_probs = torch.softmax(vit(img), dim=1)
+        # Start timer
+        start_time = timer()
 
-    # Create label and probability dict
-    pred_labels_and_probs = {
-        class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))
-    }
+        # Transform and add batch dimension
+        img_tensor = vit_transforms(img).unsqueeze(0)
 
-    # Calculate prediction time
-    pred_time = round(timer() - start_time, 5)
+        # Inference
+        vit.eval()
+        with torch.inference_mode():
+            pred_probs = torch.softmax(vit(img_tensor), dim=1)
 
-    return pred_labels_and_probs, pred_time
+        pred_labels_and_probs = {
+            class_names[i]: float(pred_probs[0][i])
+            for i in range(len(class_names))
+        }
+
+        pred_time = round(timer() - start_time, 5)
+
+        return pred_labels_and_probs, pred_time
+
+    except (UnidentifiedImageError, TypeError, ValueError) as e:
+        return {"Error": f"Invalid image input: {str(e)}"}, 0.0
 
 ### 4. Gradio app setup ###