File size: 2,220 Bytes
704a77e
ce971f1
 
704a77e
00935ea
704a77e
 
 
 
 
c5b0335
 
 
 
 
 
 
 
 
704a77e
1e1e794
ce971f1
 
 
 
704a77e
ce971f1
 
 
704a77e
c5b0335
 
 
 
ce971f1
1e1e794
 
 
 
ce971f1
 
 
 
 
 
55f7e53
09a050f
ce971f1
70720ab
ce971f1
 
09a050f
 
 
 
 
 
1e1e794
c5b0335
 
ce971f1
 
fa2d0fb
c5b0335
70720ab
ce971f1
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import numpy as np
import base64
import gradio as gr
from tensorflow import keras
import cv2

# Load the model
model_path = 'sketch2draw_model.h5'  # Update with your model path
model = keras.models.load_model(model_path)

# 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
}

# 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)
    image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)

    # Resize and normalize the 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

    # 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

# 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>")
    
    with gr.Row():
        with gr.Column():
            # Create a sketchpad for drawing
            canvas = gr.Sketchpad(label="Draw Here")
            clear_btn = gr.Button("Clear")
            generate_btn = gr.Button("Generate Image")
        
        with gr.Column():
            # Create a dropdown for different terrain colors
            color_dropdown = gr.Dropdown(
                label="Select Brush Color",
                choices=list(color_mapping.keys()),
                value="Water"  # Default color
            )

            output_image = gr.Image(label="Generated Image", type="numpy")

    # Define the actions for buttons
    clear_btn.click(fn=clear_canvas, inputs=None, outputs=canvas)

    # Click to generate an image
    generate_btn.click(fn=predict, inputs=canvas, outputs=output_image)

# Launch the Gradio app
demo.launch()