Update app.py

app.py

@@ -38,13 +38,16 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            canvas = gr.Sketchpad(label="Draw Here")  # Updated: Removed 'shape' argument
+            # Create a sketchpad for drawing
+            canvas = gr.Sketchpad(label="Draw Here", tool="brush", brush_color="black")  
             clear_btn = gr.Button("Clear")
             generate_btn = gr.Button("Generate Image")
         
         with gr.Column():
-            # Add color buttons for terrains
-            color_btns = {name: gr.Button(name) for name in color_mapping.keys()}
+            # Create color buttons for different terrains
+            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)
+                
             output_image = gr.Image(label="Generated Image", type="numpy")
 
     # Define the actions for buttons
@@ -53,13 +56,6 @@ with gr.Blocks() as demo:
 
     clear_btn.click(fn=clear_canvas, inputs=None, outputs=canvas)
 
-    # Assign color to brush for each color button
-    def change_color(color):
-        return color
-
-    for color_name, color in color_mapping.items():
-        color_btns[color_name].click(fn=change_color, inputs=None, outputs=canvas)
-
     # Click to generate an image
     generate_btn.click(fn=predict, inputs=canvas, outputs=output_image)