Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import json
 import time
-import math
+import random
 from datetime import datetime
 from typing import List, Tuple
 
@@ -13,13 +13,15 @@ from torch.utils.data import DataLoader, random_split
 from torchvision import datasets, transforms
 from PIL import Image
 
+
 # ============================================================
-# Configuration
+# Paths / Device
 # ============================================================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
 DATA_DIR = os.path.join(BASE_DIR, "data")
 MODEL_DIR = os.path.join(BASE_DIR, "saved_models")
 META_DIR = os.path.join(BASE_DIR, "saved_models_meta")
+
 os.makedirs(DATA_DIR, exist_ok=True)
 os.makedirs(MODEL_DIR, exist_ok=True)
 os.makedirs(META_DIR, exist_ok=True)
@@ -32,8 +34,14 @@ CLASS_NAMES = [str(i) for i in range(10)]
 # Model
 # ============================================================
 class SimpleCNN(nn.Module):
-    def __init__(self, conv1_channels: int = 16, conv2_channels: int = 32, kernel_size: int = 3,
-                 dropout: float = 0.2, fc_dim: int = 128):
+    def __init__(
+        self,
+        conv1_channels: int = 16,
+        conv2_channels: int = 32,
+        kernel_size: int = 3,
+        dropout: float = 0.2,
+        fc_dim: int = 128,
+    ):
         super().__init__()
         padding = kernel_size // 2
 
@@ -41,13 +49,13 @@ class SimpleCNN(nn.Module):
             nn.Conv2d(1, conv1_channels, kernel_size=kernel_size, padding=padding),
             nn.ReLU(),
             nn.MaxPool2d(2),
+
             nn.Conv2d(conv1_channels, conv2_channels, kernel_size=kernel_size, padding=padding),
             nn.ReLU(),
             nn.MaxPool2d(2),
         )
 
-        # MNIST input = 1 x 28 x 28
-        # after two 2x2 poolings => 7 x 7
+        # 28x28 -> 14x14 -> 7x7
         flattened_dim = conv2_channels * 7 * 7
 
         self.classifier = nn.Sequential(
@@ -65,13 +73,15 @@ class SimpleCNN(nn.Module):
 
 
 # ============================================================
-# Data utilities
+# Dataset helpers
 # ============================================================
 def get_datasets(dataset_name: str):
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.5,), (0.5,))
-    ])
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ]
+    )
 
     if dataset_name == "MNIST":
         train_dataset = datasets.MNIST(DATA_DIR, train=True, download=True, transform=transform)
@@ -90,23 +100,34 @@ def make_loaders(dataset_name: str, batch_size: int, val_ratio: float = 0.1):
 
     val_size = int(len(train_dataset) * val_ratio)
     train_size = len(train_dataset) - val_size
+
     train_subset, val_subset = random_split(train_dataset, [train_size, val_size])
 
     train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
     val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False)
     test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
     return train_loader, val_loader, test_loader
 
 
 # ============================================================
-# Model registry helpers
+# Model save/load helpers
 # ============================================================
+def model_weight_path(model_name: str) -> str:
+    return os.path.join(MODEL_DIR, f"{model_name}.pt")
+
+
 def model_meta_path(model_name: str) -> str:
     return os.path.join(META_DIR, f"{model_name}.json")
 
 
-def model_weight_path(model_name: str) -> str:
-    return os.path.join(MODEL_DIR, f"{model_name}.pt")
+def list_saved_models() -> List[str]:
+    names = []
+    for fn in os.listdir(META_DIR):
+        if fn.endswith(".json"):
+            names.append(fn[:-5])
+    names.sort(reverse=True)
+    return names
 
 
 def save_model(model: nn.Module, model_name: str, config: dict, training_summary: dict):
@@ -121,15 +142,6 @@ def save_model(model: nn.Module, model_name: str, config: dict, training_summary
         json.dump(payload, f, indent=2, ensure_ascii=False)
 
 
-def list_saved_models() -> List[str]:
-    models = []
-    for filename in os.listdir(META_DIR):
-        if filename.endswith(".json"):
-            models.append(filename[:-5])
-    models.sort(reverse=True)
-    return models
-
-
 def load_model(model_name: str) -> Tuple[nn.Module, dict]:
     meta_file = model_meta_path(model_name)
     weight_file = model_weight_path(model_name)
@@ -142,13 +154,14 @@ def load_model(model_name: str) -> Tuple[nn.Module, dict]:
     with open(meta_file, "r", encoding="utf-8") as f:
         meta = json.load(f)
 
-    config = meta["config"]
+    cfg = meta["config"]
+
     model = SimpleCNN(
-        conv1_channels=config["conv1_channels"],
-        conv2_channels=config["conv2_channels"],
-        kernel_size=config["kernel_size"],
-        dropout=config["dropout"],
-        fc_dim=config["fc_dim"],
+        conv1_channels=cfg["conv1_channels"],
+        conv2_channels=cfg["conv2_channels"],
+        kernel_size=cfg["kernel_size"],
+        dropout=cfg["dropout"],
+        fc_dim=cfg["fc_dim"],
     )
     state_dict = torch.load(weight_file, map_location=DEVICE)
     model.load_state_dict(state_dict)
@@ -158,17 +171,18 @@ def load_model(model_name: str) -> Tuple[nn.Module, dict]:
 
 
 # ============================================================
-# Training / evaluation
+# Train / Eval
 # ============================================================
 def evaluate(model: nn.Module, loader: DataLoader, criterion: nn.Module):
     model.eval()
     total_loss = 0.0
-    correct = 0
     total = 0
+    correct = 0
 
     with torch.no_grad():
         for images, labels in loader:
             images, labels = images.to(DEVICE), labels.to(DEVICE)
+
             outputs = model(images)
             loss = criterion(outputs, labels)
 
@@ -177,14 +191,23 @@ def evaluate(model: nn.Module, loader: DataLoader, criterion: nn.Module):
             correct += (preds == labels).sum().item()
             total += labels.size(0)
 
-    avg_loss = total_loss / total if total > 0 else 0.0
-    acc = correct / total if total > 0 else 0.0
+    avg_loss = total_loss / total if total else 0.0
+    acc = correct / total if total else 0.0
     return avg_loss, acc
 
 
-def train_model(dataset_name: str, conv1_channels: int, conv2_channels: int, kernel_size: int,
-                dropout: float, fc_dim: int, learning_rate: float, batch_size: int,
-                epochs: int, model_tag: str):
+def train_model(
+    dataset_name: str,
+    conv1_channels: int,
+    conv2_channels: int,
+    kernel_size: int,
+    dropout: float,
+    fc_dim: int,
+    learning_rate: float,
+    batch_size: int,
+    epochs: int,
+    model_tag: str,
+):
     train_loader, val_loader, test_loader = make_loaders(dataset_name, batch_size)
 
     model = SimpleCNN(
@@ -198,21 +221,15 @@ def train_model(dataset_name: str, conv1_channels: int, conv2_channels: int, ker
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
-    history = {
-        "epoch": [],
-        "train_loss": [],
-        "train_acc": [],
-        "val_loss": [],
-        "val_acc": [],
-    }
-
+    history = []
+    logs = []
     start_time = time.time()
 
     for epoch in range(1, epochs + 1):
         model.train()
         running_loss = 0.0
-        correct = 0
         total = 0
+        correct = 0
 
         for images, labels in train_loader:
             images, labels = images.to(DEVICE), labels.to(DEVICE)
@@ -228,32 +245,36 @@ def train_model(dataset_name: str, conv1_channels: int, conv2_channels: int, ker
             correct += (preds == labels).sum().item()
             total += labels.size(0)
 
-        train_loss = running_loss / total if total > 0 else 0.0
-        train_acc = correct / total if total > 0 else 0.0
+        train_loss = running_loss / total if total else 0.0
+        train_acc = correct / total if total else 0.0
         val_loss, val_acc = evaluate(model, val_loader, criterion)
 
-        history["epoch"].append(epoch)
-        history["train_loss"].append(train_loss)
-        history["train_acc"].append(train_acc)
-        history["val_loss"].append(val_loss)
-        history["val_acc"].append(val_acc)
-
-        yield {
-            "status": (
-                f"Epoch {epoch}/{epochs} | "
-                f"train_loss={train_loss:.4f}, train_acc={train_acc:.4f}, "
-                f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}"
-            ),
-            "history": history,
-            "finished": False,
-            "models": None,
+        row = {
+            "epoch": epoch,
+            "train_loss": round(train_loss, 4),
+            "train_acc": round(train_acc, 4),
+            "val_loss": round(val_loss, 4),
+            "val_acc": round(val_acc, 4),
         }
+        history.append(row)
+
+        logs.append(
+            f"Epoch {epoch}/{epochs} | "
+            f"train_loss={train_loss:.4f}, train_acc={train_acc:.4f}, "
+            f"val_loss={val_loss:.4f}, val_acc={val_acc:.4f}"
+        )
+
+        yield (
+            "\n".join(logs),
+            history,
+            gr.update(),
+        )
 
     test_loss, test_acc = evaluate(model, test_loader, criterion)
     elapsed = time.time() - start_time
 
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    safe_tag = model_tag.strip().replace(" ", "_") if model_tag else dataset_name.lower()
+    safe_tag = model_tag.strip().replace(" ", "_") if model_tag.strip() else dataset_name.lower()
     model_name = f"{safe_tag}_{timestamp}"
 
     config = {
@@ -267,49 +288,53 @@ def train_model(dataset_name: str, conv1_channels: int, conv2_channels: int, ker
         "batch_size": batch_size,
         "epochs": epochs,
     }
+
     training_summary = {
-        "final_train_loss": history["train_loss"][-1],
-        "final_train_acc": history["train_acc"][-1],
-        "final_val_loss": history["val_loss"][-1],
-        "final_val_acc": history["val_acc"][-1],
-        "test_loss": test_loss,
-        "test_acc": test_acc,
-        "elapsed_seconds": elapsed,
+        "final_train_loss": history[-1]["train_loss"] if history else None,
+        "final_train_acc": history[-1]["train_acc"] if history else None,
+        "final_val_loss": history[-1]["val_loss"] if history else None,
+        "final_val_acc": history[-1]["val_acc"] if history else None,
+        "test_loss": round(test_loss, 4),
+        "test_acc": round(test_acc, 4),
+        "elapsed_seconds": round(elapsed, 2),
         "device": str(DEVICE),
     }
+
     save_model(model, model_name, config, training_summary)
 
-    final_message = (
-        f"Training finished.\n\n"
-        f"Saved model: {model_name}\n"
-        f"Device: {DEVICE}\n"
-        f"Test loss: {test_loss:.4f}\n"
-        f"Test accuracy: {test_acc:.4f}\n"
-        f"Elapsed time: {elapsed:.1f}s"
-    )
+    logs.append("")
+    logs.append("Training finished.")
+    logs.append(f"Saved model: {model_name}")
+    logs.append(f"Device: {DEVICE}")
+    logs.append(f"Test loss: {test_loss:.4f}")
+    logs.append(f"Test accuracy: {test_acc:.4f}")
+    logs.append(f"Elapsed time: {elapsed:.1f}s")
 
-    yield {
-        "status": final_message,
-        "history": history,
-        "finished": True,
-        "models": list_saved_models(),
-    }
+    models = list_saved_models()
+    selected = model_name if model_name in models else (models[0] if models else None)
+
+    yield (
+        "\n".join(logs),
+        history,
+        gr.update(choices=models, value=selected),
+    )
 
 
 # ============================================================
-# Inference helpers
+# Inference
 # ============================================================
 def preprocess_uploaded_image(image: Image.Image):
     if image is None:
         raise ValueError("Please upload an image.")
 
-    transform = transforms.Compose([
-        transforms.Grayscale(num_output_channels=1),
-        transforms.Resize((28, 28)),
-        transforms.ToTensor(),
-        transforms.Normalize((0.5,), (0.5,))
-    ])
-
+    transform = transforms.Compose(
+        [
+            transforms.Grayscale(num_output_channels=1),
+            transforms.Resize((28, 28)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ]
+    )
     tensor = transform(image).unsqueeze(0)
     return tensor
 
@@ -326,16 +351,14 @@ def predict_uploaded_image(model_name: str, image: Image.Image):
         probs = torch.softmax(logits, dim=1).squeeze(0).cpu().tolist()
         pred_idx = int(torch.argmax(logits, dim=1).item())
 
-    conf = max(probs)
     result_text = (
         f"Prediction: {CLASS_NAMES[pred_idx]}\n"
-        f"Confidence: {conf:.4f}\n\n"
+        f"Confidence: {max(probs):.4f}\n\n"
         f"Model: {model_name}\n"
         f"Dataset: {meta['config']['dataset_name']}"
     )
-
-    prob_table = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
-    return result_text, prob_table
+    prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
+    return result_text, prob_dict
 
 
 def test_random_sample(model_name: str):
@@ -344,9 +367,9 @@ def test_random_sample(model_name: str):
 
     model, meta = load_model(model_name)
     dataset_name = meta["config"]["dataset_name"]
-    _, test_dataset = get_datasets(dataset_name)
 
-    idx = torch.randint(low=0, high=len(test_dataset), size=(1,)).item()
+    _, test_dataset = get_datasets(dataset_name)
+    idx = random.randint(0, len(test_dataset) - 1)
     image_tensor, label = test_dataset[idx]
 
     with torch.no_grad():
@@ -354,8 +377,8 @@ def test_random_sample(model_name: str):
         probs = torch.softmax(logits, dim=1).squeeze(0).cpu().tolist()
         pred_idx = int(torch.argmax(logits, dim=1).item())
 
-    display_img = image_tensor.squeeze(0).cpu()
-    prob_table = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
+    display_img = image_tensor.squeeze(0).cpu().numpy()
+
     result_text = (
         f"Random test sample\n"
         f"Ground truth: {label}\n"
@@ -363,18 +386,21 @@ def test_random_sample(model_name: str):
         f"Confidence: {max(probs):.4f}\n"
         f"Model dataset: {dataset_name}"
     )
-    return display_img, result_text, prob_table
+    prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
+    return display_img, result_text, prob_dict
 
 
 def get_model_info(model_name: str):
     if not model_name:
-        return "No model selected."
+        return {"message": "No model selected."}
+
     meta_file = model_meta_path(model_name)
     if not os.path.exists(meta_file):
-        return "Selected model metadata not found."
+        return {"message": "Metadata not found."}
+
     with open(meta_file, "r", encoding="utf-8") as f:
         meta = json.load(f)
-    return json.dumps(meta, indent=2, ensure_ascii=False)
+    return meta
 
 
 def refresh_models_dropdown():
@@ -382,97 +408,55 @@ def refresh_models_dropdown():
     return gr.update(choices=models, value=models[0] if models else None)
 
 
-# ============================================================
-# Gradio callbacks
-# ============================================================
-def training_callback(dataset_name, conv1_channels, conv2_channels, kernel_size,
-                      dropout, fc_dim, learning_rate, batch_size, epochs, model_tag):
-    for step in train_model(
-        dataset_name=dataset_name,
-        conv1_channels=conv1_channels,
-        conv2_channels=conv2_channels,
-        kernel_size=kernel_size,
-        dropout=dropout,
-        fc_dim=fc_dim,
-        learning_rate=learning_rate,
-        batch_size=batch_size,
-        epochs=epochs,
-        model_tag=model_tag,
-    ):
-        line_data = [
-            [e, tl, ta, vl, va]
-            for e, tl, ta, vl, va in zip(
-                step["history"]["epoch"],
-                step["history"]["train_loss"],
-                step["history"]["train_acc"],
-                step["history"]["val_loss"],
-                step["history"]["val_acc"],
-            )
-        ]
-
-        dropdown_update = gr.update()
-        if step["finished"] and step["models"] is not None:
-            models = step["models"]
-            dropdown_update = gr.update(choices=models, value=models[0] if models else None)
-
-        yield step["status"], line_data, dropdown_update, dropdown_update
-
-
 # ============================================================
 # UI
 # ============================================================
 initial_models = list_saved_models()
 
-with gr.Blocks(title="CNN Trainer and Tester") as demo:
-    gr.Markdown("# Simple CNN Trainer and Tester")
+with gr.Blocks(title="Image Classification") as demo:
+    gr.Markdown("# Image Classification")
     gr.Markdown(
-        "This app is designed for lightweight image classification experiments on MNIST or FashionMNIST. "
-        "Tab 1 trains a simple CNN. Tab 2 loads a saved model and tests it on uploaded images or random test samples."
+        "Train a simple CNN on MNIST or FashionMNIST, then test saved models "
+        "with an uploaded image or a random sample."
     )
 
-with gr.Tab("Train"):
-    with gr.Row():
-        with gr.Column(scale=1):
-            dataset_name = gr.Dropdown(
-                choices=["MNIST", "FashionMNIST"],
-                value="MNIST",
-                label="Dataset"
-            )
-            conv1_channels = gr.Slider(8, 64, value=16, step=8, label="Conv1 Channels")
-            conv2_channels = gr.Slider(16, 128, value=32, step=16, label="Conv2 Channels")
-            kernel_size = gr.Dropdown(choices=[3, 5], value=3, label="Kernel Size")
-            dropout = gr.Slider(0.0, 0.7, value=0.2, step=0.05, label="Dropout")
-            fc_dim = gr.Slider(32, 256, value=128, step=32, label="FC Hidden Dimension")
-            learning_rate = gr.Number(value=0.001, label="Learning Rate")
-            batch_size = gr.Dropdown(choices=[32, 64, 128, 256], value=64, label="Batch Size")
-            epochs = gr.Slider(1, 10, value=3, step=1, label="Epochs")
-            model_tag = gr.Textbox(label="Model Tag", placeholder="e.g. mnist_demo")
-            train_btn = gr.Button("Start Training", variant="primary")
-
-        with gr.Column(scale=1):
-            train_status = gr.Textbox(label="Training Status", lines=10)
-            train_plot = gr.LinePlot(
-                x="epoch",
-                y="value",
-                color="metric",
-                title="Training Curves",
-                y_title="Value",
-                x_title="Epoch",
-            )
+    with gr.Tabs():
+        with gr.Tab("Train"):
+            with gr.Row():
+                with gr.Column():
+                    dataset_name = gr.Dropdown(
+                        choices=["MNIST", "FashionMNIST"],
+                        value="MNIST",
+                        label="Dataset"
+                    )
+                    conv1_channels = gr.Slider(8, 64, value=16, step=8, label="Conv1 Channels")
+                    conv2_channels = gr.Slider(16, 128, value=32, step=16, label="Conv2 Channels")
+                    kernel_size = gr.Dropdown(choices=[3, 5], value=3, label="Kernel Size")
+                    dropout = gr.Slider(0.0, 0.7, value=0.2, step=0.05, label="Dropout")
+                    fc_dim = gr.Slider(32, 256, value=128, step=32, label="FC Hidden Dimension")
+                    learning_rate = gr.Number(value=0.001, label="Learning Rate")
+                    batch_size = gr.Dropdown(choices=[32, 64, 128, 256], value=64, label="Batch Size")
+                    epochs = gr.Slider(1, 10, value=3, step=1, label="Epochs")
+                    model_tag = gr.Textbox(label="Model Tag", placeholder="e.g. mnist_demo")
+                    train_btn = gr.Button("Start Training", variant="primary")
+
+                with gr.Column():
+                    train_status = gr.Textbox(label="Training Log", lines=18)
+                    train_history = gr.JSON(label="Training History")
 
         with gr.Tab("Test"):
             with gr.Row():
-                with gr.Column(scale=1):
+                with gr.Column():
                     model_selector = gr.Dropdown(
                         choices=initial_models,
                         value=initial_models[0] if initial_models else None,
                         label="Select Saved Model"
                     )
                     refresh_btn = gr.Button("Refresh Model List")
-                    model_info = gr.Code(label="Model Metadata", language="json")
                     load_info_btn = gr.Button("Show Model Info")
+                    model_info = gr.JSON(label="Model Metadata")
 
-                with gr.Column(scale=1):
+                with gr.Column():
                     upload_image = gr.Image(type="pil", label="Upload Image")
                     predict_btn = gr.Button("Predict Uploaded Image", variant="primary")
                     predict_text = gr.Textbox(label="Prediction Result", lines=6)
@@ -480,34 +464,27 @@ with gr.Tab("Train"):
 
             with gr.Row():
                 random_test_btn = gr.Button("Test Random Sample")
+
             with gr.Row():
                 random_sample_image = gr.Image(type="numpy", label="Random Test Image")
                 random_sample_text = gr.Textbox(label="Random Sample Result", lines=6)
                 random_sample_probs = gr.Label(label="Random Sample Probabilities")
 
-    def format_lineplot_rows(rows):
-        output = []
-        for epoch, train_loss, train_acc, val_loss, val_acc in rows:
-            output.append({"epoch": epoch, "value": train_loss, "metric": "train_loss"})
-            output.append({"epoch": epoch, "value": train_acc, "metric": "train_acc"})
-            output.append({"epoch": epoch, "value": val_loss, "metric": "val_loss"})
-            output.append({"epoch": epoch, "value": val_acc, "metric": "val_acc"})
-        return output
-
-    def wrapped_training_callback(*args):
-        for status, rows, train_dd_update, test_dd_update in training_callback(*args):
-            yield status, format_lineplot_rows(rows), train_dd_update, test_dd_update
-
-    train_model_selector_hidden = gr.Dropdown(visible=False)
-    test_model_selector_hidden = gr.Dropdown(visible=False)
-
     train_btn.click(
-        fn=wrapped_training_callback,
+        fn=train_model,
         inputs=[
-            dataset_name, conv1_channels, conv2_channels, kernel_size,
-            dropout, fc_dim, learning_rate, batch_size, epochs, model_tag
+            dataset_name,
+            conv1_channels,
+            conv2_channels,
+            kernel_size,
+            dropout,
+            fc_dim,
+            learning_rate,
+            batch_size,
+            epochs,
+            model_tag,
         ],
-        outputs=[train_status, train_plot, train_model_selector_hidden, model_selector],
+        outputs=[train_status, train_history, model_selector],
     )
 
     refresh_btn.click(
@@ -534,5 +511,6 @@ with gr.Tab("Train"):
         outputs=[random_sample_image, random_sample_text, random_sample_probs],
     )
 
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()