Fixed errors

app.py

@@ -39,47 +39,19 @@ transform = transforms.Compose([
 # Prediction function
 
 
-def predict(sketch):
-    try:
-        # Extract the image data from the dictionary
-        image = sketch['image']
-        
-        # Convert the image to a numpy array
-        image_array = np.array(image)
-        
-        # Ensure the image is 2D (grayscale)
-        if len(image_array.shape) == 3:
-            image_array = image_array[:,:,0]  # Take the first channel if it's RGB
-        
-        # Resize to 28x28 if necessary
-        if image_array.shape != (28, 28):
-            pil_image = Image.fromarray(image_array.astype('uint8'))
-            pil_image = pil_image.resize((28, 28))
-            image_array = np.array(pil_image)
-        
-        # Convert to tensor and add batch and channel dimensions
-        tensor_image = torch.tensor(image_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
-        
-        # Normalize the image
-        tensor_image = tensor_image / 255.0  # Scale to [0, 1]
-        tensor_image = (tensor_image - 0.5) / 0.5  # Normalize to [-1, 1]
-        
-        with torch.no_grad():
-            output = model(tensor_image)
-            probabilities = nn.Softmax(dim=1)(output)
-        
-        # Return the probabilities as a dictionary
-        return {str(i): float(probabilities[0][i]) for i in range(10)}
-    except Exception as e:
-        print(f"Error in predict function: {e}")
-        import traceback
-        traceback.print_exc()
-        return {"error": str(e)}
+# Prediction function
+def predict(image):
+    image = transform(image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        output = model(image)
+        probabilities = nn.Softmax(dim=1)(output)
+        predicted_class = torch.argmax(probabilities, dim=1)
+    return {str(i): probabilities[0][i].item() for i in range(10)}
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Sketchpad(),
+    inputs=gr.Sketchpad(type='pil'),
     outputs=gr.Label(num_top_classes=10)
 )