Initial commit: Animal Doodle Classifier

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,28 @@
+---
+title: Animal Doodle Classifier
+emoji: 🎨
+colorFrom: blue
+colorTo: purple
+sdk: streamlit
+sdk_version: "1.28.0"
+app_file: app.py
+pinned: false
+---
+
+# 🎨 Animal Doodle Classifier
+
+An RNN-based classifier that recognizes hand-drawn animal doodles in real-time!
+
+## Supported Animals
+- butterfly, cow, elephant, giraffe, monkey
+- octopus, scorpion, shark, snake, spider
+
+## Model
+- **Architecture:** Bidirectional GRU
+- **Accuracy:** 97.47% Top-1, 99.75% Top-3
+- **Training Data:** Google Quick Draw dataset
+
+## How It Works
+1. Draw an animal on the canvas
+2. Your strokes are captured and preprocessed to match Quick Draw format
+3. The RNN model predicts which animal you drew

app.py

@@ -0,0 +1,265 @@
+"""
+RNN Animal Doodle Classifier - Hugging Face Spaces
+Self-contained Streamlit app with embedded model class
+"""
+import ast
+import json
+from pathlib import Path
+import numpy as np
+import streamlit as st
+from streamlit_drawable_canvas import st_canvas
+import torch
+from torch import nn
+
+# ============================================================================
+# Model Definition (embedded from training-doodle.py)
+# ============================================================================
+
+class GRUClassifier(nn.Module):
+    """Bidirectional GRU classifier for sequence classification."""
+    def __init__(self, input_size: int, hidden_size: int, num_layers: int,
+                 bidirectional: bool, dropout: float, num_classes: int, use_packing: bool = True):
+        super().__init__()
+        self.use_packing = use_packing
+        self.gru = nn.GRU(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+            dropout=dropout if num_layers > 1 else 0.0,
+        )
+        out_dim = hidden_size * (2 if bidirectional else 1)
+        self.norm = nn.LayerNorm(out_dim)
+        self.fc = nn.Linear(out_dim, num_classes)
+
+    def forward(self, x: torch.Tensor, lengths: torch.Tensor):
+        if self.use_packing:
+            packed = nn.utils.rnn.pack_padded_sequence(
+                x, lengths.cpu(), batch_first=True, enforce_sorted=False
+            )
+            _, h_n = self.gru(packed)
+        else:
+            _, h_n = self.gru(x)
+        if self.gru.bidirectional:
+            h = torch.cat([h_n[-2], h_n[-1]], dim=1)
+        else:
+            h = h_n[-1]
+        h = self.norm(h)
+        return self.fc(h)
+
+def parse_drawing_to_seq(drawing_str: str) -> np.ndarray:
+    """Convert drawing JSON to sequence of [dx, dy, pen_lift]."""
+    try:
+        strokes = json.loads(drawing_str)
+    except Exception:
+        try:
+            strokes = ast.literal_eval(drawing_str)
+        except Exception:
+            return np.zeros((0, 3), dtype=np.float32)
+    
+    seq_parts = []
+    for stroke in strokes:
+        if not isinstance(stroke, (list, tuple)) or len(stroke) != 2:
+            continue
+        x, y = stroke
+        n = min(len(x), len(y))
+        if n < 2:
+            continue
+        x = np.asarray(x[:n], dtype=np.int16)
+        y = np.asarray(y[:n], dtype=np.int16)
+        dx = np.diff(x).astype(np.float32) / 255.0
+        dy = np.diff(y).astype(np.float32) / 255.0
+        if dx.size == 0:
+            continue
+        pen = np.zeros_like(dx, dtype=np.float32)
+        pen[-1] = 1.0
+        seq_parts.append(np.stack([dx, dy, pen], axis=1))
+    
+    if not seq_parts:
+        return np.zeros((0, 3), dtype=np.float32)
+    
+    seq = np.concatenate(seq_parts, axis=0)
+    seq[:, :2] = np.clip(seq[:, :2], -1.0, 1.0)
+    return seq.astype(np.float32)
+
+# ============================================================================
+# Constants
+# ============================================================================
+
+CANVAS_SIZE = 400
+STROKE_WIDTH = 3
+ANIMALS = ["butterfly", "cow", "elephant", "giraffe", "monkey",
+           "octopus", "scorpion", "shark", "snake", "spider"]
+
+CALIB_TARGET_MEAN = 0.04
+CALIB_MAX_GAIN = 12.0
+CALIB_MIN_GAIN = 0.5
+
+def _calibrate_seq(seq: np.ndarray) -> np.ndarray:
+    """Scale (dx, dy) so the mean step magnitude matches training data."""
+    if seq is None or seq.ndim != 2 or seq.shape[1] < 2 or seq.shape[0] == 0:
+        return seq
+    steps = np.sqrt((seq[:, 0] ** 2) + (seq[:, 1] ** 2))
+    curr = float(steps.mean()) if steps.size else 0.0
+    if curr <= 1e-6:
+        return seq
+    gain = float(np.clip(CALIB_TARGET_MEAN / curr, CALIB_MIN_GAIN, CALIB_MAX_GAIN))
+    out = seq.astype(np.float32).copy()
+    out[:, 0:2] = np.clip(out[:, 0:2] * gain, -1.0, 1.0)
+    return out
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
+@st.cache_resource
+def load_model():
+    """Load the trained RNN model."""
+    model_path = Path(__file__).parent / "rnn_animals_best.pt"
+    if not model_path.exists():
+        st.error(f"Model file not found: {model_path}")
+        return None, None
+    
+    checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+    cfg = checkpoint.get("config", {})
+    
+    model = GRUClassifier(
+        input_size=3,
+        hidden_size=cfg.get("hidden_size", 512),
+        num_layers=cfg.get("num_layers", 2),
+        bidirectional=cfg.get("bidirectional", True),
+        dropout=cfg.get("dropout", 0.3),
+        num_classes=len(ANIMALS),
+        use_packing=True
+    )
+    model.load_state_dict(checkpoint["model_state"])
+    model.eval()
+    
+    class_to_idx = checkpoint.get("class_to_idx", {a: i for i, a in enumerate(ANIMALS)})
+    idx_to_class = {v: k for k, v in class_to_idx.items()}
+    return model, idx_to_class
+
+# ============================================================================
+# Stroke Processing
+# ============================================================================
+
+def canvas_strokes_to_quickdraw(canvas_json):
+    """Convert canvas to QuickDraw format with preprocessing."""
+    if canvas_json is None:
+        return []
+    
+    objects = canvas_json.get("objects", [])
+    raw_strokes = []
+    
+    for obj in objects:
+        if obj.get("type") != "path":
+            continue
+        path = obj.get("path", [])
+        xs, ys = [], []
+        for cmd in path:
+            if len(cmd) < 3:
+                continue
+            if cmd[0] == "M":
+                xs.append(float(cmd[1]))
+                ys.append(float(cmd[2]))
+            elif cmd[0] == "Q" and len(cmd) >= 5:
+                xs.append(float(cmd[3]))
+                ys.append(float(cmd[4]))
+            elif cmd[0] == "L":
+                xs.append(float(cmd[1]))
+                ys.append(float(cmd[2]))
+        if len(xs) >= 2:
+            raw_strokes.append((xs, ys))
+    
+    if not raw_strokes:
+        return []
+    
+    # Downsample
+    downsampled = []
+    for xs, ys in raw_strokes:
+        if len(xs) > 25:
+            step = max(1, len(xs) // 25)
+            xs, ys = xs[::step], ys[::step]
+        downsampled.append((xs, ys))
+    
+    # Smooth
+    smoothed = []
+    for xs, ys in downsampled:
+        if len(xs) >= 3:
+            xs_s = [xs[0]] + [(xs[i-1]+xs[i]+xs[i+1])/3 for i in range(1, len(xs)-1)] + [xs[-1]]
+            ys_s = [ys[0]] + [(ys[i-1]+ys[i]+ys[i+1])/3 for i in range(1, len(ys)-1)] + [ys[-1]]
+            smoothed.append((xs_s, ys_s))
+        else:
+            smoothed.append((xs, ys))
+    
+    # Center and scale
+    all_x = [x for xs, _ in smoothed for x in xs]
+    all_y = [y for _, ys in smoothed for y in ys]
+    min_x, max_x = min(all_x), max(all_x)
+    min_y, max_y = min(all_y), max(all_y)
+    
+    scale = 235 / max(max(1, max_x - min_x), max(1, max_y - min_y))
+    cx, cy = (min_x + max_x) / 2, (min_y + max_y) / 2
+    ox, oy = 127.5 - cx * scale, 127.5 - cy * scale
+    
+    result = []
+    for xs, ys in smoothed:
+        xs_n = [int(np.clip(x * scale + ox, 0, 255)) for x in xs]
+        ys_n = [int(np.clip(y * scale + oy, 0, 255)) for y in ys]
+        result.append([xs_n, ys_n])
+    return result
+
+def predict(model, idx_to_class, strokes):
+    """Make prediction from strokes."""
+    if not strokes or model is None:
+        return None
+    try:
+        seq = parse_drawing_to_seq(json.dumps(strokes))
+        if seq is None or len(seq) < 6:
+            return None
+        seq = _calibrate_seq(seq)
+        seq_t = torch.tensor(seq, dtype=torch.float32).unsqueeze(0)
+        lengths = torch.tensor([seq.shape[0]], dtype=torch.long)
+        with torch.no_grad():
+            probs = torch.softmax(model(seq_t, lengths), dim=1)
+            top_p, top_i = torch.topk(probs, k=5, dim=1)
+        return [(idx_to_class.get(top_i[0,i].item()), top_p[0,i].item()) for i in range(5)]
+    except Exception as e:
+        st.error(f"Error: {e}")
+        return None
+
+# ============================================================================
+# Main App
+# ============================================================================
+
+def main():
+    st.set_page_config(page_title="Animal Doodle Classifier", page_icon="🎨", layout="wide")
+    st.title("🎨 Animal Doodle Classifier")
+    st.caption("Draw: butterfly, cow, elephant, giraffe, monkey, octopus, scorpion, shark, snake, spider")
+    
+    model, idx_to_class = load_model()
+    if model is None:
+        return
+    
+    col1, col2 = st.columns([1, 1])
+    with col1:
+        canvas = st_canvas(
+            stroke_width=STROKE_WIDTH, stroke_color="#000000",
+            background_color="#FFFFFF", height=CANVAS_SIZE, width=CANVAS_SIZE,
+            drawing_mode="freedraw", key="canvas"
+        )
+    
+    with col2:
+        st.subheader("Predictions")
+        if canvas.json_data:
+            strokes = canvas_strokes_to_quickdraw(canvas.json_data)
+            if strokes:
+                results = predict(model, idx_to_class, strokes)
+                if results:
+                    st.success(f"**{results[0][0].upper()}** ({results[0][1]*100:.1f}%)")
+                    for name, prob in results:
+                        st.progress(prob, text=f"{name}: {prob*100:.1f}%")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,4 @@
+streamlit>=1.28.0
+streamlit-drawable-canvas>=0.9.3
+torch>=2.0.0
+numpy>=1.24.0

rnn_animals_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3045301cb82537f7ccdaa2271c0d0944470d6a5079cc633a1bfcafb2198ac895
+size 44206972