app.py

import os
import json
import time
import random
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


# ============================================================
# 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)

# Force CPU on Hugging Face Spaces for this lightweight demo
DEVICE = torch.device("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),
        )

        # 28x28 -> 14x14 -> 7x7
        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


# ============================================================
# Dataset helpers
# ============================================================
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 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 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):
    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 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)

    cfg = meta["config"]

    model = SimpleCNN(
        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)
    model.to(DEVICE)
    model.eval()
    return model, meta


# ============================================================
# Train / Eval
# ============================================================
def evaluate(model: nn.Module, loader: DataLoader, criterion: nn.Module):
    model.eval()
    total_loss = 0.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)

            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 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,
):
    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 = []
    logs = []
    start_time = time.time()

    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        total = 0
        correct = 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 else 0.0
        train_acc = correct / total if total else 0.0
        val_loss, val_acc = evaluate(model, val_loader, criterion)

        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.strip() 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[-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)

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

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

    result_text = (
        f"Prediction: {CLASS_NAMES[pred_idx]}\n"
        f"Confidence: {max(probs):.4f}\n\n"
        f"Model: {model_name}\n"
        f"Dataset: {meta['config']['dataset_name']}"
    )
    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):
    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 = random.randint(0, len(test_dataset) - 1)
    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().numpy()

    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}"
    )
    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 {"message": "No model selected."}

    meta_file = model_meta_path(model_name)
    if not os.path.exists(meta_file):
        return {"message": "Metadata not found."}

    with open(meta_file, "r", encoding="utf-8") as f:
        meta = json.load(f)
    return meta


def refresh_models_dropdown():
    models = list_saved_models()
    return gr.update(choices=models, value=models[0] if models else None)


# ============================================================
# UI
# ============================================================
initial_models = list_saved_models()

with gr.Blocks(title="Image Classification") as demo:
    gr.Markdown("# Image Classification")
    gr.Markdown(
        "Train a simple CNN on MNIST or FashionMNIST, then test saved models "
        "with an uploaded image or a random sample."
    )

    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():
                    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")
                    load_info_btn = gr.Button("Show Model Info")
                    model_info = gr.JSON(label="Model Metadata")

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

    train_btn.click(
        fn=train_model,
        inputs=[
            dataset_name,
            conv1_channels,
            conv2_channels,
            kernel_size,
            dropout,
            fc_dim,
            learning_rate,
            batch_size,
            epochs,
            model_tag,
        ],
        outputs=[train_status, train_history, 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(ssr_mode=False)