Update app.py

app.py

@@ -54,31 +54,25 @@ def predict(img):
         {label: confidence, label: confidence, ...}
     """
 
-    try:
-        if img is None or not isinstance(img, dict) or 'image' not in img:
-            return {"Error": 1.0}
-    
-        img_data = img['composite']
-        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
-        return {LABELS[i]: float(v) for i, v in zip(indices, probs)}
-    except Exception as e:
-        print(f"Error in prediction: {str(e)}")
-        return {"Error": 1.0}
+    img_data = img['composite']
+    print(img_data.sum())
+    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
+    return {LABELS[i]: float(v) for i, v in zip(indices, probs)}
 
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Sketchpad(
         label="Draw Here",
-        brush=gr.Brush(default_size=20, default_color="#FFFFFF", colors=["#FFFFFF"]),
+        brush=gr.Brush(default_size=20, default_color="#000000", colors=["#000000"]),
         image_mode="L",
         layers=False
     ),