Update app.py and training logic

app.py

@@ -1,32 +1,81 @@
-from datasets import load_dataset, load_metric
+from datasets import load_dataset
 from transformers import AutoImageProcessor, AutoModelForImageClassification, TrainingArguments, Trainer
 from PIL import Image
 import numpy as np
 import gradio as gr
 import torch
 import os
+import shutil
+import zipfile
+import evaluate  # pour les métriques
 
 DEFAULT_MODEL_NAME = "facebook/convnextv2-tiny-1k-224"  # ou "google/vit-base-patch16-224"
 DEFAULT_OUTPUT_DIR = "./weld_cls_model_best"
+EXTRACT_DIR = "./uploaded_dataset"  # dossier où l'on extrait l'archive
 
 
-def train_model(data_dir, model_name=DEFAULT_MODEL_NAME, num_epochs=10, batch_size=16, lr=5e-5):
+def extract_archive(archive_path, extract_to=EXTRACT_DIR):
     """
-    Lance l'entraînement sur un dataset d'images de type imagefolder :
-    data_dir/
-        bonne/
-            img1.jpg
-            ...
-        mauvaise/
-            img2.jpg
-            ...
+    Extrait une archive .zip ou .rar dans extract_to.
+    La structure attendue après extraction est de type imagefolder :
+        extract_to/
+            bonne/
+                img1.jpg
+                ...
+            mauvaise/
+                img2.jpg
+                ...
+    """
+    if archive_path is None or not os.path.isfile(archive_path):
+        return None, f"Erreur : aucune archive de dataset fournie."
+
+    # Nettoyer l'ancien dossier, s'il existe
+    if os.path.isdir(extract_to):
+        shutil.rmtree(extract_to)
+    os.makedirs(extract_to, exist_ok=True)
+
+    archive_lower = archive_path.lower()
+
+    try:
+        if archive_lower.endswith(".zip"):
+            with zipfile.ZipFile(archive_path, "r") as zf:
+                zf.extractall(extract_to)
+        elif archive_lower.endswith(".rar"):
+            try:
+                import rarfile
+            except ImportError:
+                return None, (
+                    "Erreur : format .rar demandé mais le module 'rarfile' n'est pas installé.\n"
+                    "Ajoute 'rarfile' dans requirements.txt, ou utilise une archive .zip."
+                )
+            with rarfile.RarFile(archive_path) as rf:
+                rf.extractall(extract_to)
+        else:
+            return None, "Erreur : format d'archive non supporté. Utilise .zip ou .rar."
+    except Exception as e:
+        return None, f"Erreur lors de l'extraction de l'archive : {e}"
+
+    return extract_to, None
+
+
+def train_model(dataset_archive_path, model_name=DEFAULT_MODEL_NAME, num_epochs=10, batch_size=16, lr=5e-5):
+    """
+    Lance l'entraînement à partir d'une archive uploadée (zip/rar) contenant
+    un dataset de type imagefolder.
     """
 
-    if not os.path.isdir(data_dir):
-        return f"Erreur : le dossier {data_dir} n'existe pas."
+    # 0) Extraction de l'archive
+    data_dir, err = extract_archive(dataset_archive_path)
+    if err is not None:
+        return err
+    if data_dir is None or not os.path.isdir(data_dir):
+        return f"Erreur : le dossier de données '{data_dir}' est introuvable après extraction."
 
     # 1) Charger le dataset
-    dataset = load_dataset("imagefolder", data_dir=data_dir)
+    try:
+        dataset = load_dataset("imagefolder", data_dir=data_dir)
+    except Exception as e:
+        return f"Erreur lors du chargement du dataset avec 'imagefolder' : {e}"
 
     label_names = dataset["train"].features["label"].names
     num_labels = len(label_names)
@@ -54,7 +103,7 @@ def train_model(data_dir, model_name=DEFAULT_MODEL_NAME, num_epochs=10, batch_si
     )
 
     # 4) Définir les métriques
-    metric = load_metric("accuracy")
+    metric = evaluate.load("accuracy")
 
     def compute_metrics(eval_pred):
         logits, labels = eval_pred
@@ -128,15 +177,14 @@ def predict(image):
 # -----------------------
 with gr.Blocks() as demo:
     gr.Markdown("# Classification de soudures – Entraînement + Inférence\n"
-                "Interface Gradio pour entraîner un modèle Hugging Face et tester des images.")
+                "Upload d'un dataset (.zip ou .rar), entraînement du modèle, puis test sur des images.")
 
     with gr.Tab("Entraînement"):
         gr.Markdown("## Lancer l'entraînement")
 
-        data_dir_input = gr.Textbox(
-            label="Dossier du dataset (format imagefolder)",
-            value="path/vers/tes_images",
-            placeholder="Ex : /data/soudures"
+        dataset_file_input = gr.File(
+            label="Archive du dataset (.zip ou .rar)",
+            type="filepath"
         )
         model_name_input = gr.Textbox(
             label="Nom du modèle Hugging Face",
@@ -167,7 +215,7 @@ with gr.Blocks() as demo:
 
         train_button.click(
             fn=train_model,
-            inputs=[data_dir_input, model_name_input, epochs_input, batch_input, lr_input],
+            inputs=[dataset_file_input, model_name_input, epochs_input, batch_input, lr_input],
             outputs=train_output
         )