app.py

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",

                    )

        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()