Update app.py

app.py

@@ -1,394 +1,14 @@
-import os
 import json
-import time
-import random
-from datetime import datetime
-from typing import List, Tuple
 
 import spaces
 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 / basic config
-# ============================================================
-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)
-
-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),
-        )
-
-        flattened_dim = conv2_channels * 7 * 7  # 28x28 -> 14x14 -> 7x7
-
-        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):
-    cpu_state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
-    torch.save(cpu_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, device: torch.device) -> 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="cpu")
-    model.load_state_dict(state_dict)
-    model.to(device)
-    model.eval()
-    return model, meta
-
-
-# ============================================================
-# ZeroGPU helpers
-# ============================================================
-def get_runtime_device() -> torch.device:
-    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+from train_utils import train_model, list_saved_models, model_meta_path
+from predict_utils import predict_uploaded_image, test_random_sample
 
 
 @spaces.GPU(duration=120)
-def _train_on_gpu(
-    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,
-):
-    device = get_runtime_device()
-
-    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()
-
-    def evaluate(loader):
-        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
-
-    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(val_loader)
-
-        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}"
-        )
-
-    test_loss, test_acc = evaluate(test_loader)
-    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")
-
-    return "\n".join(logs), history, training_summary, model_name
-
-
-@spaces.GPU(duration=60)
-def _predict_uploaded_image_gpu(model_name: str, image: Image.Image):
-    if not model_name:
-        return "Please select a model.", None
-
-    if image is None:
-        return "Please upload an image.", None
-
-    device = get_runtime_device()
-    model, meta = load_model(model_name, device)
-
-    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).to(device)
-
-    with torch.no_grad():
-        logits = model(tensor)
-        probs = torch.softmax(logits, dim=1).squeeze(0).detach().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']}\n"
-        f"Runtime device: {device}"
-    )
-    prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
-    return result_text, prob_dict
-
-
-@spaces.GPU(duration=60)
-def _test_random_sample_gpu(model_name: str):
-    if not model_name:
-        return None, "Please select a model.", None
-
-    device = get_runtime_device()
-    model, meta = load_model(model_name, device)
-    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).detach().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}\n"
-        f"Runtime device: {device}"
-    )
-    prob_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(10)}
-    return display_img, result_text, prob_dict
-
-
-# ============================================================
-# UI callbacks
-# ============================================================
 def train_callback(
-    dataset_name,
     conv1_channels,
     conv2_channels,
     kernel_size,
@@ -400,8 +20,7 @@ def train_callback(
     model_tag,
 ):
     try:
-        logs, history, summary, model_name = _train_on_gpu(
-            dataset_name,
+        logs, history, summary, model_name = train_model(
             int(conv1_channels),
             int(conv2_channels),
             int(kernel_size),
@@ -412,112 +31,136 @@ def train_callback(
             int(epochs),
             model_tag,
         )
+
         models = list_saved_models()
         selected = model_name if model_name in models else (models[0] if models else None)
+
         return logs, history, summary, gr.update(choices=models, value=selected)
+
     except Exception as e:
-        return f"Training failed:\n{str(e)}", None, None, gr.update()
+        return f"Échec de l’entraînement :\n{str(e)}", None, None, gr.update()
 
 
+@spaces.GPU(duration=60)
 def predict_uploaded_image_callback(model_name, image):
     try:
-        return _predict_uploaded_image_gpu(model_name, image)
+        return predict_uploaded_image(model_name, image)
     except Exception as e:
-        return f"Prediction failed:\n{str(e)}", None
+        return f"Échec de la prédiction :\n{str(e)}", None
 
 
+@spaces.GPU(duration=60)
 def test_random_sample_callback(model_name):
     try:
-        return _test_random_sample_gpu(model_name)
+        return test_random_sample(model_name)
     except Exception as e:
-        return None, f"Random test failed:\n{str(e)}", None
+        return None, f"Échec du test aléatoire :\n{str(e)}", None
+
+
+def refresh_models_dropdown():
+    models = list_saved_models()
+    return gr.update(choices=models, value=models[0] if models else None)
 
 
 def get_model_info(model_name: str):
     if not model_name:
-        return {"message": "No model selected."}
+        return {"message": "Aucun modèle sélectionné."}
 
     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)
+    try:
+        with open(meta_file, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except FileNotFoundError:
+        return {"message": "Métadonnées introuvables."}
 
 
-# ============================================================
-# UI
-# ============================================================
 initial_models = list_saved_models()
 
-with gr.Blocks(title="Image Classification") as demo:
-    gr.Markdown("# Image Classification")
+
+with gr.Blocks(title="Classification d’images microscopiques") as demo:
+    gr.Markdown("# Classification d’images microscopiques de charbons de bois")
     gr.Markdown(
-        "Train a simple CNN on MNIST or FashionMNIST, then test saved models "
-        "with an uploaded image or a random sample."
+        "Cette application permet d’entraîner un réseau de neurones convolutif simple "
+        "sur un jeu de données privé Hugging Face, puis de tester les modèles sauvegardés "
+        "sur une image importée ou sur un échantillon aléatoire."
     )
 
     with gr.Tabs():
-        with gr.Tab("Train"):
+        with gr.Tab("Entraîner"):
             with gr.Row():
                 with gr.Column():
-                    dataset_name = gr.Dropdown(
-                        choices=["MNIST", "FashionMNIST"],
-                        value="MNIST",
-                        label="Dataset",
+                    gr.Markdown("### Paramètres d’entraînement")
+
+                    conv1_channels = gr.Slider(
+                        8, 64, value=16, step=8, label="Nombre de canaux - couche convolutionnelle 1"
+                    )
+                    conv2_channels = gr.Slider(
+                        16, 128, value=32, step=16, label="Nombre de canaux - couche convolutionnelle 2"
                     )
-                    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")
+                    kernel_size = gr.Dropdown(
+                        choices=[3, 5], value=3, label="Taille du noyau"
+                    )
+                    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="Dimension de la couche cachée fully-connected"
+                    )
+                    learning_rate = gr.Number(
+                        value=0.001, label="Taux d’apprentissage"
+                    )
+                    batch_size = gr.Dropdown(
+                        choices=[16, 32, 64, 128], value=32, label="Taille du batch"
+                    )
+                    epochs = gr.Slider(
+                        1, 20, value=5, step=1, label="Nombre d’époques"
+                    )
+                    model_tag = gr.Textbox(
+                        label="Nom court du modèle",
+                        placeholder="ex. charbon_cnn_test"
+                    )
+
+                    train_btn = gr.Button("Lancer l’entraînement", variant="primary")
 
                 with gr.Column():
-                    train_status = gr.Textbox(label="Training Log", lines=18)
-                    train_history = gr.JSON(label="Training History")
-                    train_summary = gr.JSON(label="Training Summary")
+                    train_status = gr.Textbox(label="Journal d’entraînement", lines=18)
+                    train_history = gr.JSON(label="Historique d’entraînement")
+                    train_summary = gr.JSON(label="Résumé d’entraînement")
 
-        with gr.Tab("Test"):
+        with gr.Tab("Tester"):
             with gr.Row():
                 with gr.Column():
+                    gr.Markdown("### Modèle sauvegardé")
+
                     model_selector = gr.Dropdown(
                         choices=initial_models,
                         value=initial_models[0] if initial_models else None,
-                        label="Select Saved Model",
+                        label="Sélectionner un modèle",
                     )
-                    refresh_btn = gr.Button("Refresh Model List")
-                    load_info_btn = gr.Button("Show Model Info")
-                    model_info = gr.JSON(label="Model Metadata")
+                    refresh_btn = gr.Button("Actualiser la liste des modèles")
+                    load_info_btn = gr.Button("Afficher les informations du modèle")
+                    model_info = gr.JSON(label="Métadonnées du modèle")
 
                 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=7)
-                    predict_probs = gr.Label(label="Class Probabilities")
+                    gr.Markdown("### Prédiction sur une image importée")
+
+                    upload_image = gr.Image(type="pil", label="Importer une image")
+                    predict_btn = gr.Button("Prédire la classe", variant="primary")
+                    predict_text = gr.Textbox(label="Résultat de la prédiction", lines=7)
+                    predict_probs = gr.Label(label="Probabilités par classe")
 
             with gr.Row():
-                random_test_btn = gr.Button("Test Random Sample")
+                random_test_btn = gr.Button("Tester un échantillon aléatoire")
 
             with gr.Row():
-                random_sample_image = gr.Image(type="numpy", label="Random Test Image")
-                random_sample_text = gr.Textbox(label="Random Sample Result", lines=7)
-                random_sample_probs = gr.Label(label="Random Sample Probabilities")
+                random_sample_image = gr.Image(type="pil", label="Image test aléatoire")
+                random_sample_text = gr.Textbox(label="Résultat sur l’échantillon", lines=7)
+                random_sample_probs = gr.Label(label="Probabilités par classe")
 
     train_btn.click(
         fn=train_callback,
         inputs=[
-            dataset_name,
             conv1_channels,
             conv2_channels,
             kernel_size,