Fixed errors

app.py

@@ -39,18 +39,26 @@ transform = transforms.Compose([
 # Prediction function
 
 
-
-def predict(image):
+def predict(sketch):
     try:
-        # The input image is already a 2D numpy array (grayscale)
-        # Ensure it's the right size and normalize it
+        # 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
         
-        pil_image = Image.fromarray(image.squeeze().astype('uint8'))
-        pil_image = pil_image.resize((28, 28))
-        image = np.array(pil_image)
+        # 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, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
+        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]