Migrate to gradio

.streamlit/config.toml

@@ -1,7 +0,0 @@
-[client]
-showErrorDetails = false
-
-[theme]
-primaryColor = "#3a6ef3bd"
-backgroundColor = "#edbc41"
-textColor="#000000"

Pipfile

@@ -0,0 +1,20 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+streamlit = "==1.35.0"
+streamlit-drawable-canvas = "==0.9.3"
+huggingface-hub = "==0.23.4"
+polars = "==0.20.31"
+matplotlib = "==3.9.0"
+torch = "==2.3.1"
+torchvision = "==0.18.1"
+lightning = "==2.3.0"
+gradio = "*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.12"

README.md

@@ -3,8 +3,8 @@ title: "Scribble It! AI Demo"
 emoji: 🎨
 colorFrom: yellow
 colorTo: blue
-sdk: streamlit
-sdk_version: 1.35.0
+sdk: gradio
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,42 +1,25 @@
 import os
 import random
 
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
 import polars as pl
-import streamlit as st
 import torch
 from huggingface_hub import hf_hub_download
-from matplotlib import pyplot as plt
 from PIL import Image
-from streamlit_drawable_canvas import st_canvas
 from torchvision.transforms import v2
 
 from model import ScribbleItNet
 
-# Page configuration
-st.set_page_config(page_title="Scribble It! AI Demo 🎨")
-st.title("Scribble It! AI Demo 🎨")
-
-# Set the background image
-background_image = """
-<style>
-[data-testid="stAppViewContainer"] > .main {
-    background-image: url("https://detach-entertainment.com/img/clouds.ed95a9c8.svg");
-    background-color: #edbc41;
-}
-</style>
-"""
-
-st.markdown(background_image, unsafe_allow_html=True)
-
 # Matplotlib configuration
 plt.rc('font', size=16)
-plt.rc('axes', titlesize=16)
-plt.rc('xtick', labelsize=24)
+plt.rc('axes', titlesize=24)
+plt.rc('xtick', labelsize=20)
 plt.rc('ytick', labelsize=20)
 
 
 # Load the model
-@st.cache_resource
 def load_model():
     hf_hub_download("ScribbleItAI/efficientnet-b0",
                     token=os.getenv("HF_TOKEN"),
@@ -49,6 +32,8 @@ def load_model():
 
 
 model = load_model()
+
+# Transform configuration
 transform = v2.Compose([
     v2.Resize((224, 224)),
     v2.ToDtype(torch.float32, scale=True),
@@ -56,8 +41,7 @@ transform = v2.Compose([
 ])
 
 
-# Load the vocabulary
-@st.cache_data
+# Load vocabulary
 def load_vocabulary():
     hf_hub_download("ScribbleItAI/efficientnet-b0",
                     token=os.getenv("HF_TOKEN"),
@@ -68,49 +52,19 @@ def load_vocabulary():
 
 vocabulary = load_vocabulary()
 idx2vocab = {row["word_idx"]: row for row in vocabulary}
-vocabulary = {
-    f"{row['word_hash']}_{row['category_idx']}": row
-    for row in vocabulary
-}
-
-# Sidebar
-with st.sidebar:
-    drawing_mode = st.selectbox(
-        "Drawing tool:",
-        ("freedraw", "point", "line", "rect", "circle", "transform"))
-
-    stroke_width = st.slider("Stroke width: ", 1, 25, 5)
-    if drawing_mode == 'point':
-        point_display_radius = st.slider("Point display radius: ", 1, 25, 5)
-    stroke_color = st.color_picker("Stroke color hex: ")
-    bg_color = st.color_picker("Background color hex: ", "#ffffff")
-    realtime_update = st.checkbox("Update in realtime", True)
-
-if st.button("New word") or "sample" not in st.session_state:
-    st.session_state.sample = random.choice(list(vocabulary.values()))
-
-st.markdown(f" Draw a: **{st.session_state.sample['word']}**")
-
-# Canvas
-canvas_result = st_canvas(
-    stroke_width=stroke_width,
-    stroke_color=stroke_color,
-    background_color=bg_color,
-    update_streamlit=realtime_update,
-    height=500,
-    width=800,
-    drawing_mode=drawing_mode,
-    point_display_radius=point_display_radius
-    if drawing_mode == 'point' else 0,
-    key="canvas",
-)
-
-# Inference
-if canvas_result.image_data is not None:
-    img = canvas_result.image_data
-    img = torch.tensor(img)[:, :, :3].permute(2, 0, 1)
-    img = transform(img)
-    outputs = model(img.unsqueeze(0).to(model.device))
+vocab_list = [row["word"] for row in vocabulary]
+
+
+# Select a random word
+def get_random_word():
+    return random.choice(vocab_list)
+
+
+# Process the image drawn on canvas
+def process_image(image, current_word):
+    img_tensor = torch.tensor(image["composite"]).permute(2, 0, 1)
+    img_tensor = transform(img_tensor)
+    outputs = model(img_tensor.unsqueeze(0).to(model.device))
     outputs = torch.softmax(outputs, dim=1)
     preds, indices = outputs.topk(100, dim=1)
 
@@ -125,19 +79,71 @@ if canvas_result.image_data is not None:
             "category": vocab["category_name"],
             "prob": pred
         })
-    predictions = pl.DataFrame(predictions)
-    predictions = predictions.group_by("word").agg(
-        pl.col("prob").max().alias("prob"))
-    predictions = predictions.sort("prob")
-    predictions = predictions.tail(10)
-
-    # Plot the predictions
-    fig = plt.figure(figsize=(14, 10))
+
+    predictions_df = pl.DataFrame(predictions)
+    predictions_df = predictions_df.with_columns(
+        pl.col("word").str.to_lowercase())
+    predictions_df = predictions_df.group_by("word").agg(
+        pl.col("prob").max().alias("prob")).sort("prob").tail(10)
+
+    # Visualizing predictions
+    fig, ax = plt.subplots(figsize=(10, 8))
+    plt.subplots_adjust(left=0.25, top=0.9, right=0.9, bottom=0.1)
     colors = [
-        "green" if word == st.session_state.sample["word"] else "tab:blue"
-        for word in predictions["word"]
+        "green" if word == current_word else "tab:blue"
+        for word in predictions_df["word"]
     ]
-    plt.barh(predictions["word"], predictions["prob"], color=colors)
-    plt.xlabel("Probability")
-    plt.ylabel("Word")
-    st.pyplot(fig)
+    ax.barh(predictions_df["word"], predictions_df["prob"], color=colors)
+    ax.set_title("Top 10 Predictions", pad=15)
+    ax.set_xlabel("Probability")
+    plt.close(fig)
+    return fig, current_word
+
+
+def update_image(image):
+    image = Image.fromarray(image["composite"])
+    return image
+
+
+def create_initial_image():
+    data = np.full((500, 700, 3), 255, dtype=np.uint8)  # White image
+    return Image.fromarray(data)
+
+
+# Create a white image with the dimensions for the ImageEditor
+initial_image = create_initial_image
+
+# Interface definition
+with gr.Blocks(theme=gr.themes.Soft(),
+               css="input {font-size: 24px; font-weight: 600;}") as demo_app:
+    gr.Markdown("# Scribble It! AI Demo 🎨")
+    gr.Markdown("### Draw the word shown and let the AI guess what it is!")
+
+    with gr.Row():
+        word_output = gr.Textbox(label="Your word to draw:",
+                                 value=get_random_word(),
+                                 scale=1,
+                                 max_lines=1)
+        new_word_button = gr.Button("New Word", scale=0, variant="primary")
+
+    with gr.Row():
+        image_editor = gr.ImageEditor(
+            label="Draw Here!",
+            image_mode="RGB",
+            sources=[],
+            transforms=[],
+            layers=False,
+            value=initial_image,
+            brush=gr.Brush(colors=["#000000", "#FF0000", "#00FF00", "#0000FF"],
+                           default_size=10))
+        plot_output = gr.Plot(label="Model Guesses")
+
+    image_editor.clear(create_initial_image, outputs=image_editor)
+    image_editor.change(process_image,
+                        inputs=[image_editor, word_output],
+                        outputs=[plot_output, word_output])
+
+    new_word_button.click(get_random_word, outputs=word_output)
+    new_word_button.click(create_initial_image, outputs=image_editor)
+
+demo_app.launch()

requirements.txt

@@ -1,5 +1,4 @@
-streamlit==1.35.0; python_version >= '3.8' and python_full_version != '3.9.7'
-streamlit-drawable-canvas==0.9.3; python_version >= '3.6'
+gradio==4.36.1; python_version >= '3.8'
 huggingface-hub==0.23.4; python_version >= '3.8'
 polars==0.20.31; python_version >= '3.8'
 matplotlib==3.9.0; python_version >= '3.9'