Update app.py

app.py

@@ -18,6 +18,7 @@ color_mapping = {
     "Sky": (135, 206, 235)   # Sky Blue
 }
 
+# Function to predict and generate image
 def predict(image):
     # Decode the image from base64
     image_data = np.frombuffer(base64.b64decode(image.split(",")[1]), np.uint8)
@@ -32,6 +33,10 @@ def predict(image):
     generated_image = (generated_image * 255).astype(np.uint8)  # Rescale to 0-255
     return generated_image
 
+# Function to clear the canvas
+def clear_canvas():
+    return np.zeros((400, 400, 3), dtype=np.uint8)
+
 # Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("<h1>Sketch to Draw Model</h1>")
@@ -39,21 +44,23 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             # Create a sketchpad for drawing
-            canvas = gr.Sketchpad(label="Draw Here", brush_color="black")  
+            canvas = gr.Sketchpad(label="Draw Here")  
             clear_btn = gr.Button("Clear")
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column():
             # Create color buttons for different terrains
+            brush_color = gr.Variable(default="black")  # Variable to store brush color
+
+            def change_color(color):
+                brush_color.set(color)  # Update the brush color
+
             for color_name, color in color_mapping.items():
-                gr.Button(color_name, variant="primary").click(fn=lambda c=color: canvas.update(brush_color=c), inputs=None, outputs=None)
-                
+                gr.Button(color_name, variant="primary").click(fn=change_color, inputs=color, outputs=None)
+
             output_image = gr.Image(label="Generated Image", type="numpy")
 
     # Define the actions for buttons
-    def clear_canvas():
-        return np.zeros((400, 400, 3), dtype=np.uint8)
-
     clear_btn.click(fn=clear_canvas, inputs=None, outputs=canvas)
 
     # Click to generate an image