Use seaborn color palette

Pipfile

@@ -13,6 +13,7 @@ torch = "==2.3.1"
 torchvision = "==0.18.1"
 lightning = "==2.3.0"
 gradio = "*"
+seaborn = "*"
 
 [dev-packages]
 

Pipfile.lock

@@ -1,7 +1,7 @@
 {
     "_meta": {
         "hash": {
-            "sha256": "89d094e3bcfe678655d839dafae5d0ca2cea12d8ff08ee98f9ec2c72bd372138"
+            "sha256": "6cd84f80a3c605c6c74d9229abcd62e631682c99fa656bdcb81e996e3bb9d715"
         },
         "pipfile-spec": 6,
         "requires": {
@@ -1772,6 +1772,15 @@
             "markers": "sys_platform != 'emscripten'",
             "version": "==0.4.9"
         },
+        "seaborn": {
+            "hashes": [
+                "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987",
+                "sha256:93e60a40988f4d65e9f4885df477e2fdaff6b73a9ded434c1ab356dd57eefff7"
+            ],
+            "index": "pypi",
+            "markers": "python_version >= '3.8'",
+            "version": "==0.13.2"
+        },
         "semantic-version": {
             "hashes": [
                 "sha256:bdabb6d336998cbb378d4b9db3a4b56a1e3235701dc05ea2690d9a997ed5041c",

app.py

@@ -2,6 +2,7 @@ import os
 import random
 
 import gradio as gr
+import seaborn as sns
 import matplotlib.pyplot as plt
 import numpy as np
 import polars as pl
@@ -12,6 +13,9 @@ from torchvision.transforms import v2
 
 from model import ScribbleItNet
 
+# Seaborn configuration
+sns.set_theme()
+
 # Matplotlib configuration
 plt.rc('font', size=16)
 plt.rc('axes', titlesize=24)
@@ -92,22 +96,34 @@ def process_image(image, current_word):
     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)
+        pl.col("prob").max().alias("prob"))
+    predictions_df = predictions_df.sort("prob", descending=True).head(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 == current_word else "tab:blue"
+        sns.color_palette()[2]
+        if word == current_word else sns.color_palette()[0]
         for word in predictions_df["word"]
     ]
 
     if current_word in predictions_df["word"]:
         gr.Info("AI guessed the word correctly! 🎉")
 
-    ax.barh(predictions_df["word"], predictions_df["prob"], color=colors)
+    sns.barplot(
+        data=predictions_df,
+        y="word",
+        x="prob",
+        hue="word",
+        orient="h",
+        palette=colors,
+        legend=False,
+        ax=ax,
+    )
     ax.set_title("Top 10 Predictions", pad=15)
     ax.set_xlabel("Probability")
+    ax.set_ylabel(None)
     plt.close(fig)
     return fig, current_word
 
@@ -118,7 +134,7 @@ def update_image(image):
 
 
 def create_initial_image():
-    data = np.full((500, 700, 3), 255, dtype=np.uint8)  # White image
+    data = np.full((520, 700, 3), 255, dtype=np.uint8)  # White image
     return Image.fromarray(data)
 
 
@@ -126,7 +142,7 @@ def create_initial_image():
 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!")
+    gr.Markdown("### Draw the word shown below and let the AI guess what it is!")
 
     with gr.Row():
         word_output = gr.Textbox(label="Your word to draw:",
@@ -144,7 +160,7 @@ with gr.Blocks(theme=gr.themes.Soft(),
             transforms=[],
             layers=False,
             value=create_initial_image,
-            brush=gr.Brush(colors=["#000000", "#FF0000", "#00FF00", "#0000FF"],
+            brush=gr.Brush(colors=["#000000", "#4c72b0", "#55a868", "#d62728"],
                            default_size=10))
         plot_output = gr.Plot(label="Model Guesses")
 

requirements.txt

@@ -2,6 +2,7 @@ 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'
+seaborn==0.13.2; python_version >= '3.8'
 torch==2.3.1; python_full_version >= '3.8.0'
 torchvision==0.18.1; python_version >= '3.8'
 lightning==2.3.0; python_version >= '3.8'