Fixed errors

app.py

@@ -39,22 +39,33 @@ transform = transforms.Compose([
 # Prediction function
 
 
-def predict(image_dict):
+
+def predict(image):
     try:
-        # Convert the dictionary to a PIL Image
-        image = Image.fromarray(np.uint8(image_dict["image"])).convert('RGB')
+        # The input image is already a 2D numpy array (grayscale)
+        # Ensure it's the right size and normalize it
+        if image.shape != (28, 28):
+            pil_image = Image.fromarray(image.squeeze().astype('uint8'))
+            pil_image = pil_image.resize((28, 28))
+            image = np.array(pil_image)
+        
+        # Convert to tensor and add batch and channel dimensions
+        tensor_image = torch.tensor(image, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
         
-        # Apply the transformation
-        image = transform(image).unsqueeze(0)  # Add batch dimension
+        # 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(image)
+            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)}
 
 # Create the Gradio interface