Update app.py

app.py

@@ -1,10 +1,10 @@
 import os
-import time
 import json
+import time
 import threading
 
 import numpy as np
-from PIL import Image, ImageOps
+from PIL import Image
 
 import torch
 import torch.nn as nn
@@ -21,9 +21,9 @@ import gradio as gr
 class MnistCNN(nn.Module):
     def __init__(self, num_classes: int = 10, dropout: float = 0.25):
         super().__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)    # 28x28 -> 28x28
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)   # 28x28 -> 28x28
-        self.pool = nn.MaxPool2d(2, 2)                             # 28x28 -> 14x14
+        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)  # 28x28 -> 14x14
         self.dropout = nn.Dropout(dropout)
         self.fc1 = nn.Linear(64 * 14 * 14, 128)
         self.fc2 = nn.Linear(128, num_classes)
@@ -46,7 +46,7 @@ MODEL = MnistCNN().to(DEVICE)
 WEIGHTS_PATH = "mnist_cnn.pth"
 CONFIG_PATH = "mnist_config.json"
 
-DEFAULT_CONFIG = {
+CFG_DEFAULT = {
     "num_classes": 10,
     "dropout": 0.25,
     "normalize_mean": 0.1307,
@@ -54,30 +54,28 @@ DEFAULT_CONFIG = {
     "image_size": 28
 }
 
-# Use deterministic-ish behavior for demos (not perfect determinism on all systems)
 torch.manual_seed(42)
 np.random.seed(42)
 
 
-def save_config():
-    with open(CONFIG_PATH, "w") as f:
-        json.dump(DEFAULT_CONFIG, f, indent=2)
-
-
-def load_config():
+def load_or_init_config():
     if os.path.exists(CONFIG_PATH):
         with open(CONFIG_PATH, "r") as f:
             return json.load(f)
-    save_config()
-    return DEFAULT_CONFIG
+    with open(CONFIG_PATH, "w") as f:
+        json.dump(CFG_DEFAULT, f, indent=2)
+    return CFG_DEFAULT
 
 
-CFG = load_config()
+CFG = load_or_init_config()
+
+
+def blank_editor_value(size=280):
+    """Initial blank canvas for ImageEditor."""
+    img = Image.new("RGBA", (size, size), (255, 255, 255, 255))
+    return {"background": img, "layers": [], "composite": img}
 
 
-# -----------------------------
-# Utilities
-# -----------------------------
 def maybe_load_weights():
     global MODEL
     if os.path.exists(WEIGHTS_PATH):
@@ -91,36 +89,25 @@ def maybe_load_weights():
 
 def preprocess_pil(img: Image.Image) -> torch.Tensor:
     """
-    Converts a PIL image to MNIST-like tensor: (1,1,28,28), normalized.
-    Also attempts to handle "black ink on white background" by auto-inverting.
+    Convert PIL image to MNIST tensor (1,1,28,28), normalized like training.
+    Auto-invert if the background is bright.
     """
     if img is None:
         raise ValueError("No image provided.")
 
-    # Convert to grayscale
-    img = img.convert("L")
-
-    # Resize to 28x28
-    img = img.resize((CFG["image_size"], CFG["image_size"]))
-
-    # Convert to numpy [0..1]
+    img = img.convert("L").resize((CFG["image_size"], CFG["image_size"]))
     arr = np.array(img).astype(np.float32) / 255.0
 
-    # Auto-invert if background looks white-ish (common with sketch tools)
-    # MNIST digits are typically bright strokes on darker background.
+    # If background is mostly white, invert so digit becomes bright on dark
     if arr.mean() > 0.5:
         arr = 1.0 - arr
 
-    # Normalize like training
     arr = (arr - CFG["normalize_mean"]) / CFG["normalize_std"]
-
-    # Shape to (1,1,28,28)
-    x = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)
+    x = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)  # (1,1,28,28)
     return x.to(DEVICE)
 
 
-def predict_digit(img: Image.Image):
-    global MODEL
+def predict_from_pil(img: Image.Image):
     if img is None:
         return "No image", {}
 
@@ -137,6 +124,13 @@ def predict_digit(img: Image.Image):
     return pred, prob_dict
 
 
+def predict_from_editor(editor_value):
+    # ImageEditor returns a dict with keys: background, layers, composite
+    if editor_value is None or "composite" not in editor_value:
+        return "No drawing", {}
+    return predict_from_pil(editor_value["composite"])
+
+
 # -----------------------------
 # Training
 # -----------------------------
@@ -149,13 +143,11 @@ def get_dataloaders(batch_size: int, max_train_samples: int, max_test_samples: i
     train_ds = datasets.MNIST(root="data", train=True, download=True, transform=transform)
     test_ds  = datasets.MNIST(root="data", train=False, download=True, transform=transform)
 
-    # Subset for faster training on Spaces (optional)
     if max_train_samples and max_train_samples < len(train_ds):
         train_ds = Subset(train_ds, range(max_train_samples))
     if max_test_samples and max_test_samples < len(test_ds):
         test_ds = Subset(test_ds, range(max_test_samples))
 
-    # num_workers=0 is safest in Spaces
     train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0)
     test_dl  = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=0)
     return train_dl, test_dl
@@ -163,49 +155,38 @@ def get_dataloaders(batch_size: int, max_train_samples: int, max_test_samples: i
 
 def evaluate(model: nn.Module, test_dl: DataLoader):
     model.eval()
-    correct = 0
-    total = 0
-    loss_sum = 0.0
     criterion = nn.CrossEntropyLoss()
+    loss_sum, correct, total = 0.0, 0, 0
 
     with torch.no_grad():
         for x, y in test_dl:
             x, y = x.to(DEVICE), y.to(DEVICE)
             logits = model(x)
-            loss = criterion(logits, y)
-            loss_sum += loss.item()
-
+            loss_sum += criterion(logits, y).item()
             preds = logits.argmax(dim=1)
             correct += (preds == y).sum().item()
             total += y.numel()
 
-    avg_loss = loss_sum / max(1, len(test_dl))
-    acc = correct / max(1, total)
-    return avg_loss, acc
+    return loss_sum / max(1, len(test_dl)), correct / max(1, total)
 
 
-def train_mnist(epochs: int, lr: float, batch_size: int, max_train_samples: int, max_test_samples: int, progress=gr.Progress()):
+def train_mnist(epochs: int, lr: float, batch_size: int, max_train_samples: int, max_test_samples: int):
     global MODEL
 
+    start = time.time()
     train_dl, test_dl = get_dataloaders(batch_size, max_train_samples, max_test_samples)
 
-    # Re-init model each time you train (simple + predictable)
     model = MnistCNN(num_classes=CFG["num_classes"], dropout=CFG["dropout"]).to(DEVICE)
     optimizer = torch.optim.Adam(model.parameters(), lr=lr)
     criterion = nn.CrossEntropyLoss()
 
     logs = []
-    start = time.time()
-
     for epoch in range(1, epochs + 1):
         model.train()
-        running_loss = 0.0
-        correct = 0
-        total = 0
+        running_loss, correct, total = 0.0, 0, 0
 
-        for step, (x, y) in enumerate(progress.tqdm(train_dl, desc=f"Epoch {epoch}/{epochs}")):
+        for x, y in train_dl:
             x, y = x.to(DEVICE), y.to(DEVICE)
-
             optimizer.zero_grad()
             logits = model(x)
             loss = criterion(logits, y)
@@ -219,7 +200,6 @@ def train_mnist(epochs: int, lr: float, batch_size: int, max_train_samples: int,
 
         train_loss = running_loss / max(1, len(train_dl))
         train_acc = correct / max(1, total)
-
         test_loss, test_acc = evaluate(model, test_dl)
 
         logs.append(
@@ -228,100 +208,95 @@ def train_mnist(epochs: int, lr: float, batch_size: int, max_train_samples: int,
             f"test loss {test_loss:.4f} acc {test_acc:.4f}"
         )
 
-    # Save weights locally
     torch.save(model.state_dict(), WEIGHTS_PATH)
-    save_config()
 
-    # Swap global model
     with MODEL_LOCK:
         MODEL.load_state_dict(model.state_dict())
         MODEL.eval()
 
     elapsed = time.time() - start
-    header = f"Done. Saved weights to `{WEIGHTS_PATH}`. Device: {DEVICE}. Time: {elapsed:.1f}s\n"
-    return header + "\n".join(logs)
+    status = f"✅ Done. Saved `{WEIGHTS_PATH}`. Device: {DEVICE}. Time: {elapsed:.1f}s"
+    return status, "\n".join(logs)
 
 
-def load_saved_weights_ui():
+def load_weights_ui():
     ok = maybe_load_weights()
-    if ok:
-        return f"Loaded saved weights from `{WEIGHTS_PATH}`."
-    return f"No saved weights found at `{WEIGHTS_PATH}`. Train first."
+    return f"✅ Loaded `{WEIGHTS_PATH}`." if ok else f"⚠️ No `{WEIGHTS_PATH}` found yet. Train first."
 
 
-# Try to load weights at startup (if present)
-_ = maybe_load_weights()
+# Try load at startup
+maybe_load_weights()
 
 
 # -----------------------------
-# Gradio UI
+# Gradio UI (Gradio 6+)
 # -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# MNIST (Custom `nn.Module`) — Train + Predict (PyTorch + Gradio)")
-    gr.Markdown(
-        "Use **Train** to fit a small CNN on MNIST. Then **draw** or **upload** a digit to predict.\n\n"
-        f"- Running on: `{DEVICE}`\n"
-        f"- Weights file: `{WEIGHTS_PATH}`"
-    )
+    gr.Markdown("# MNIST — Train + Predict (PyTorch custom `nn.Module`)")
+    gr.Markdown(f"- Running on: `{DEVICE}`  \n- Weights file: `{WEIGHTS_PATH}`")
 
     with gr.Row():
         with gr.Column():
             gr.Markdown("## 1) Train (optional)")
-            epochs = gr.Slider(1, 5, value=1, step=1, label="Epochs (start with 1)")
+            epochs = gr.Slider(1, 5, value=1, step=1, label="Epochs")
             lr = gr.Number(value=1e-3, label="Learning rate", precision=6)
             batch = gr.Slider(32, 256, value=128, step=32, label="Batch size")
 
-            gr.Markdown("### Speed controls (use smaller values for faster training)")
+            gr.Markdown("### Speed controls (smaller = faster)")
             max_train = gr.Slider(1000, 60000, value=10000, step=1000, label="Max train samples")
             max_test = gr.Slider(500, 10000, value=2000, step=500, label="Max test samples")
 
             train_btn = gr.Button("Train model")
             load_btn = gr.Button("Load saved weights")
 
-            train_log = gr.Textbox(label="Training log", lines=10)
             status = gr.Textbox(label="Status", lines=2)
+            train_log = gr.Textbox(label="Training log", lines=10)
 
         with gr.Column():
             gr.Markdown("## 2) Predict")
+
             with gr.Tab("Draw"):
-                draw_img = gr.Image(source="canvas", tool="sketch", type="pil", label="Draw a digit (0-9)")
+                # ImageEditor is the Gradio 6 way to draw/paint
+                draw_editor = gr.ImageEditor(
+                    value=blank_editor_value,
+                    type="pil",
+                    canvas_size=(280, 280),
+                    fixed_canvas=True,
+                    label="Draw a digit (0–9)"
+                )
                 draw_btn = gr.Button("Predict from drawing")
+
             with gr.Tab("Upload"):
-                up_img = gr.Image(source="upload", type="pil", label="Upload an image of a digit")
+                up_img = gr.Image(type="pil", label="Upload a digit image")
                 up_btn = gr.Button("Predict from upload")
 
             pred_out = gr.Number(label="Prediction")
             prob_out = gr.Label(num_top_classes=3, label="Probabilities (top 3)")
 
-    # Wiring
     train_btn.click(
         fn=train_mnist,
         inputs=[epochs, lr, batch, max_train, max_test],
-        outputs=[train_log],
-    ).then(
-        fn=lambda: "Training complete. You can now predict.",
-        inputs=[],
-        outputs=[status],
+        outputs=[status, train_log],
     )
 
     load_btn.click(
-        fn=load_saved_weights_ui,
+        fn=load_weights_ui,
         inputs=[],
         outputs=[status],
     )
 
     draw_btn.click(
-        fn=predict_digit,
-        inputs=[draw_img],
+        fn=predict_from_editor,
+        inputs=[draw_editor],
         outputs=[pred_out, prob_out],
     )
 
     up_btn.click(
-        fn=predict_digit,
+        fn=predict_from_pil,
         inputs=[up_img],
         outputs=[pred_out, prob_out],
     )
 
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()