Update app.py

app.py

@@ -8,16 +8,9 @@ model = tf.keras.models.load_model("cnn_model.h5")
 
 def predict(image_array):
     try:
-        # Debug the input type
-        print(f"Type of input: {type(image_array)}")
-        
-        # Handle if input is a dictionary (which appears to be happening)
-        if isinstance(image_array, dict):
-            # If it's a dict, try to get the image data
-            if "image" in image_array:
-                image_array = image_array["image"]
-            else:
-                return "Error: Received dictionary input but couldn't find image data"
+        if isinstance(image_array, dict) and 'composite' in image_array:
+            # Extract the image data from the 'composite' key
+            image_array = image_array['composite']
         
         if image_array is None or np.sum(image_array) == 0:
             return "Please draw a digit."
@@ -31,7 +24,7 @@ def predict(image_array):
         confidence = float(tf.nn.softmax(logits)[0][prediction])
         return f"Digit: {prediction} (confidence: {confidence:.2%})"
     except Exception as err:
-        return f"Runtime error: {str(err)}\nInput type: {type(image_array)}"
+        return f"Runtime error: {str(err)}"
 
 gr.Interface(
     fn=predict,