initial commit

.gitignore

@@ -0,0 +1 @@
+draw_venv/

.streamlit/config.toml

@@ -0,0 +1,2 @@
+[theme]
+backgroundColor="rgb(150, 150, 150)"

app.py

@@ -0,0 +1,53 @@
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image
+import streamlit as st
+import torch
+from streamlit_drawable_canvas import st_canvas
+
+st.set_page_config(page_title="Draw Something!", layout="centered")
+
+if "prediction" not in st.session_state:
+    st.session_state["prediction"] = "Draw something!"
+
+st.markdown(f"<h1 style='text-align: center;'>{st.session_state['prediction']}</h1>", unsafe_allow_html=True)
+
+processor = AutoImageProcessor.from_pretrained("kmewhort/resnet34-sketch-classifier")
+model = AutoModelForImageClassification.from_pretrained("kmewhort/resnet34-sketch-classifier")
+
+canvas = st_canvas(
+    stroke_width=5,
+    stroke_color="#000000",
+    background_color="#FFFFFF",
+    height=700,
+    width=700,
+    drawing_mode="freedraw",
+)
+
+def predict_drawing():
+    if canvas.image_data is not None:
+        drawing = canvas.image_data.astype("uint8")
+        image = Image.fromarray(drawing).convert("L")
+        image = image.convert("RGB")
+
+        inputs = processor(images=image, return_tensors="pt")
+        
+        with torch.no_grad():
+            logits = model(**inputs).logits
+
+        predicted_class_idx = logits.argmax(-1).item()
+        st.session_state["prediction"] = model.config.id2label[predicted_class_idx]
+    else:
+        st.session_state["prediction"] = "Draw something!"
+
+if canvas.image_data is not None:
+    predict_drawing()
+
+css = '''
+<style>
+section.stMain {
+    overflow: hidden;
+}
+</style>
+'''
+st.markdown(css, unsafe_allow_html=True)
+    

requirements.txt

@@ -0,0 +1,6 @@
+Pillow==11.1.0
+streamlit==1.42.0
+streamlit_drawable_canvas==0.9.3
+transformers==4.48.3
+torch==2.6.0
+torchvision==0.21.0