Update app.py

app.py

@@ -53,21 +53,27 @@ def predict(img):
     as a dictionary:
         {label: confidence, label: confidence, ...}
     """
-    if img is None:
-        return None
+    print("Input type:", type(img))
 
-    img_data = img['image']
-    img_gray = Image.fromarray(img_data).convert('L').resize((28, 28))
-    img_tensor = transforms.ToTensor()(img_gray).unsqueeze(0)
-
-    # Make prediction
-    with torch.no_grad():
-        probs = model(img_tensor).softmax(dim=1).squeeze()
-
-    probs, indices = torch.topk(probs, 5)  # select top 5
-    probs, indices = probs.tolist(), indices.tolist()  # transform to list
-    confidences = {LABELS[i]: float(v) for i, v in zip(indices, probs)}
-    return confidences
+    try:
+        if img is None or not isinstance(img, dict) or 'image' not in img:
+            return {"Error": "Invalid input"}
+    
+        img_data = img['image']
+        img_gray = Image.fromarray(img_data).convert('L').resize((28, 28))
+        img_tensor = transforms.ToTensor()(img_gray).unsqueeze(0)
+    
+        # Make prediction
+        with torch.no_grad():
+            probs = model(img_tensor).softmax(dim=1).squeeze()
+    
+        probs, indices = torch.topk(probs, 5)  # select top 5
+        probs, indices = probs.tolist(), indices.tolist()  # transform to list
+        confidences = {LABELS[i]: float(v) for i, v in zip(indices, probs)}
+        return confidences
+    except Exception as e:
+        print(f"Error in prediction: {str(e)}")
+        return {"Error": str(e)}
 
 
 demo = gr.Interface(