Update app.py

app.py

@@ -8,8 +8,15 @@ import cv2
 model_path = 'sketch2draw_model.h5'  # Update with your model path
 model = keras.models.load_model(model_path)
 
-# Load class names for predictions
-class_names = ['grass', 'dirt', 'wood', 'water', 'sky', 'clouds']
+# Define color mapping for different terrains
+color_mapping = {
+    "Water": (0, 0, 255),   # Blue
+    "Grass": (0, 255, 0),   # Green
+    "Dirt": (139, 69, 19),  # Brown
+    "Clouds": (255, 255, 255),  # White
+    "Wood": (160, 82, 45),  # Sienna
+    "Sky": (135, 206, 235)   # Sky Blue
+}
 
 def predict(image):
     # Decode the image from base64
@@ -20,11 +27,10 @@ def predict(image):
     image = cv2.resize(image, (128, 128))  # Resize as per your model's input
     image = np.expand_dims(image, axis=0) / 255.0  # Normalize if needed
 
-    # Make prediction
-    predictions = model.predict(image)
-    predicted_class = class_names[np.argmax(predictions)]
-    
-    return predicted_class
+    # Generate image based on the model
+    generated_image = model.predict(image)[0]  # Generate image
+    generated_image = (generated_image * 255).astype(np.uint8)  # Rescale to 0-255
+    return generated_image
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -32,21 +38,25 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            canvas = gr.Sketchpad(label="Draw Here")  # Removed tool argument
-            brush_color = gr.ColorPicker(value="black", label="Brush Color")
+            canvas = gr.Sketchpad(label="Draw Here")
             clear_btn = gr.Button("Clear")
         
         with gr.Column():
-            predict_btn = gr.Button("Predict")
-            output_label = gr.Textbox(label="Predicted Texture")
-    
+            # Add color buttons for terrains
+            color_btns = {name: gr.Button(name) for name in color_mapping.keys()}
+            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)
     
-    predict_btn.click(fn=predict, inputs=canvas, outputs=output_label)
+    for color_name, color in color_mapping.items():
+        color_btns[color_name].click(fn=lambda color=color: canvas.update(value=color), inputs=None, outputs=None)
+
+    # Click to generate an image
+    canvas.submit(fn=predict, inputs=canvas, outputs=output_image)
 
 # Launch the Gradio app
 demo.launch()