Create app.py

app.py

@@ -0,0 +1,542 @@
+import os
+import json
+import time
+import math
+from datetime import datetime
+from typing import List, Tuple
+
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, random_split
+from torchvision import datasets, transforms
+from PIL import Image
+
+# ============================================================
+# Configuration
+# ============================================================
+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)
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+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):
+        super().__init__()
+        padding = kernel_size // 2
+
+        self.features = nn.Sequential(
+            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
+        flattened_dim = conv2_channels * 7 * 7
+
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(flattened_dim, fc_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(fc_dim, 10),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        return x
+
+
+# ============================================================
+# Data utilities
+# ============================================================
+def get_datasets(dataset_name: str):
+    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)
+        test_dataset = datasets.MNIST(DATA_DIR, train=False, download=True, transform=transform)
+    elif dataset_name == "FashionMNIST":
+        train_dataset = datasets.FashionMNIST(DATA_DIR, train=True, download=True, transform=transform)
+        test_dataset = datasets.FashionMNIST(DATA_DIR, train=False, download=True, transform=transform)
+    else:
+        raise ValueError(f"Unsupported dataset: {dataset_name}")
+
+    return train_dataset, test_dataset
+
+
+def make_loaders(dataset_name: str, batch_size: int, val_ratio: float = 0.1):
+    train_dataset, test_dataset = get_datasets(dataset_name)
+
+    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
+# ============================================================
+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 save_model(model: nn.Module, model_name: str, config: dict, training_summary: dict):
+    torch.save(model.state_dict(), model_weight_path(model_name))
+    payload = {
+        "model_name": model_name,
+        "config": config,
+        "training_summary": training_summary,
+        "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+    }
+    with open(model_meta_path(model_name), "w", encoding="utf-8") as f:
+        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)
+
+    if not os.path.exists(meta_file):
+        raise FileNotFoundError(f"Metadata not found for model: {model_name}")
+    if not os.path.exists(weight_file):
+        raise FileNotFoundError(f"Weights not found for model: {model_name}")
+
+    with open(meta_file, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+
+    config = 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"],
+    )
+    state_dict = torch.load(weight_file, map_location=DEVICE)
+    model.load_state_dict(state_dict)
+    model.to(DEVICE)
+    model.eval()
+    return model, meta
+
+
+# ============================================================
+# Training / evaluation
+# ============================================================
+def evaluate(model: nn.Module, loader: DataLoader, criterion: nn.Module):
+    model.eval()
+    total_loss = 0.0
+    correct = 0
+    total = 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)
+
+            total_loss += loss.item() * images.size(0)
+            preds = outputs.argmax(dim=1)
+            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
+    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):
+    train_loader, val_loader, test_loader = make_loaders(dataset_name, batch_size)
+
+    model = SimpleCNN(
+        conv1_channels=conv1_channels,
+        conv2_channels=conv2_channels,
+        kernel_size=kernel_size,
+        dropout=dropout,
+        fc_dim=fc_dim,
+    ).to(DEVICE)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+    history = {
+        "epoch": [],
+        "train_loss": [],
+        "train_acc": [],
+        "val_loss": [],
+        "val_acc": [],
+    }
+
+    start_time = time.time()
+
+    for epoch in range(1, epochs + 1):
+        model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        for images, labels in train_loader:
+            images, labels = images.to(DEVICE), labels.to(DEVICE)
+
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * images.size(0)
+            preds = outputs.argmax(dim=1)
+            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
+        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,
+        }
+
+    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()
+    model_name = f"{safe_tag}_{timestamp}"
+
+    config = {
+        "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,
+    }
+    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,
+        "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"
+    )
+
+    yield {
+        "status": final_message,
+        "history": history,
+        "finished": True,
+        "models": list_saved_models(),
+    }
+
+
+# ============================================================
+# Inference helpers
+# ============================================================
+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,))
+    ])
+
+    tensor = transform(image).unsqueeze(0)
+    return tensor
+
+
+def predict_uploaded_image(model_name: str, image: Image.Image):
+    if not model_name:
+        return "Please select a model.", None
+
+    model, meta = load_model(model_name)
+    tensor = preprocess_uploaded_image(image).to(DEVICE)
+
+    with torch.no_grad():
+        logits = model(tensor)
+        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"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
+
+
+def test_random_sample(model_name: str):
+    if not model_name:
+        return None, "Please select a model.", None
+
+    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()
+    image_tensor, label = test_dataset[idx]
+
+    with torch.no_grad():
+        logits = model(image_tensor.unsqueeze(0).to(DEVICE))
+        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))}
+    result_text = (
+        f"Random test sample\n"
+        f"Ground truth: {label}\n"
+        f"Prediction: {pred_idx}\n"
+        f"Confidence: {max(probs):.4f}\n"
+        f"Model dataset: {dataset_name}"
+    )
+    return display_img, result_text, prob_table
+
+
+def get_model_info(model_name: str):
+    if not model_name:
+        return "No model selected."
+    meta_file = model_meta_path(model_name)
+    if not os.path.exists(meta_file):
+        return "Selected model 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)
+
+
+def refresh_models_dropdown():
+    models = list_saved_models()
+    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")
+    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."
+    )
+
+    with gr.Tabs():
+        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",
+                        width=700,
+                        height=400,
+                    )
+
+        with gr.Tab("Test"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    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")
+
+                with gr.Column(scale=1):
+                    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)
+                    predict_probs = gr.Label(label="Class Probabilities")
+
+            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,
+        inputs=[
+            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],
+    )
+
+    refresh_btn.click(
+        fn=refresh_models_dropdown,
+        inputs=None,
+        outputs=model_selector,
+    )
+
+    load_info_btn.click(
+        fn=get_model_info,
+        inputs=model_selector,
+        outputs=model_info,
+    )
+
+    predict_btn.click(
+        fn=predict_uploaded_image,
+        inputs=[model_selector, upload_image],
+        outputs=[predict_text, predict_probs],
+    )
+
+    random_test_btn.click(
+        fn=test_random_sample,
+        inputs=[model_selector],
+        outputs=[random_sample_image, random_sample_text, random_sample_probs],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()