Update app.py

app.py

@@ -6,24 +6,26 @@ import logging
 logging.basicConfig(level=logging.DEBUG)
 
 def display_sketch(sketch):
-    try:
-        image_data = np.array(sketch)  # Convert the dict data to a numpy array
-        logging.debug(f"Image data type: {type(image_data)}")  # Log the type of image data
-        logging.debug(f"Image data shape: {image_data.shape}")  # Log the shape of image data
+    logging.debug(f"Received sketch data: {sketch}")
+    
+    if isinstance(sketch, dict) and "image" in sketch:
+        image_data = sketch["image"]
+        logging.debug(f"Image data type: {type(image_data)}")
+        logging.debug(f"Image data shape: {np.array(image_data).shape}")
         
-        # Plot the image
-        fig, ax = plt.subplots()
-        ax.imshow(image_data, cmap='gray')
+        plt.imshow(image_data, cmap='gray')
         plt.axis('off')
         plt.savefig("/mnt/data/output.png", bbox_inches='tight')
         return "/mnt/data/output.png"
-    except Exception as e:
-        logging.error(f"Error processing sketch: {e}")
-        return str(e)
+    else:
+        error_message = f"Unexpected sketch data format: {type(sketch)}"
+        logging.error(error_message)
+        return error_message
 
 with gr.Blocks() as demo:
-    sketchpad = gr.Sketchpad(label="Draw Something", tool="sketch")
+    sketchpad = gr.Sketchpad(label="Draw Something")
     output_image = gr.Image(label="Your Sketch")
-    sketchpad.submit(display_sketch, inputs=sketchpad, outputs=output_image)
+    submit_btn = gr.Button("Submit")
+    submit_btn.click(display_sketch, inputs=sketchpad, outputs=output_image)
 
 demo.launch()