update app.py

app.py

@@ -17,7 +17,7 @@ from transformers import (
 
 
 # =====================================================================
-# CONFIGURATION GLOBALE
+# CONFIG GLOBALE
 # =====================================================================
 
 DEFAULT_MODEL_NAME = "facebook/convnextv2-tiny-1k-224"
@@ -31,11 +31,9 @@ MODEL_UPLOAD_DIR = "uploaded_model"
 # =====================================================================
 
 def extract_zip(zip_path, dest_dir):
-    """Extrait un fichier ZIP dans dest_dir (remplace si existe)."""
-
+    """Extrait un ZIP et remplace le dossier existant."""
     if os.path.isdir(dest_dir):
         shutil.rmtree(dest_dir)
-
     os.makedirs(dest_dir, exist_ok=True)
 
     try:
@@ -43,59 +41,120 @@ def extract_zip(zip_path, dest_dir):
             zf.extractall(dest_dir)
         return True, None
     except Exception as e:
-        return False, f"Erreur lors de l'extraction : {e}"
+        return False, f"Erreur extraction ZIP : {e}"
+
+
+def find_true_dataset_root(root):
+    """
+    Trouve automatiquement le vrai dossier contenant les classes :
+    - Bonne/
+    - Mauvaise/
+
+    Même si le ZIP contient une couche inutile :
+    dataset.zip
+        images de soudures/
+            bonne/
+            mauvaise/
+
+    Cette fonction retourne le dossier qui contient réellement les classes.
+    """
+    subdirs = [d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))]
+    # Cas idéal : les classes sont directement présentes
+    if any(name.lower() in ["bonne", "mauvaise"] for name in subdirs):
+        return root
+
+    # Sinon, entrer dans le premier sous-dossier
+    if len(subdirs) == 1:
+        sub = os.path.join(root, subdirs[0])
+        deeper = [d for d in os.listdir(sub) if os.path.isdir(os.path.join(sub, d))]
+
+        if any(name.lower() in ["bonne", "mauvaise"] for name in deeper):
+            return sub
+
+    return root  # fallback
 
 
 # =====================================================================
-# PAGE 1 : ENTRAÎNEMENT DU MODÈLE
+# PAGE 1 : ENTRAÎNEMENT
 # =====================================================================
 
 def train_model(zip_dataset_path, model_name, epochs, batch_size, lr):
-    # 1) Vérifier le dataset
+
     if zip_dataset_path is None:
-        return "Erreur : aucun ZIP de dataset fourni."
+        return "Erreur : aucun dataset ZIP fourni."
 
+    # 1) Extraire le ZIP
     success, err = extract_zip(zip_dataset_path, DATASET_EXTRACT_DIR)
     if not success:
         return err
 
-    # Structure attendue :
-    # DATASET_EXTRACT_DIR/
-    #   Bonne/
-    #   Mauvaise/
+    # 2) Trouver le vrai dossier racine du dataset
+    true_root = find_true_dataset_root(DATASET_EXTRACT_DIR)
 
+    # 3) Charger dataset HF
     try:
-        dataset = load_dataset("imagefolder", data_dir=DATASET_EXTRACT_DIR)
+        dataset = load_dataset("imagefolder", data_dir=true_root)
     except Exception as e:
-        return f"Erreur : impossible de charger le dataset au format imagefolder : {e}"
+        return f"Erreur lors du chargement du dataset imagefolder : {e}"
 
+    # Afficher colonnes détectées
+    column_info = f"Colonnes détectées : {dataset['train'].column_names}\n"
+    feature_info = f"Features : {dataset['train'].features}\n"
+    debug_log = column_info + feature_info
+
+    # Vérifier que la colonne label existe
     if "label" not in dataset["train"].column_names:
-        return f"Erreur : le dataset ne contient pas de colonne 'label'. Structure attendue : ZIP → Bonne/ et Mauvaise/."
+        return debug_log + "\nErreur : aucune colonne label détectée."
 
     label_names = dataset["train"].features["label"].names
     num_labels = len(label_names)
 
-    # 2) Préprocessing
+    # 4) Préprocesseur
     processor = AutoImageProcessor.from_pretrained(model_name)
 
+    # Détecter la colonne image réellement présente
+    def detect_image_key(keys):
+        if "image" in keys:
+            return "image"
+        if "file" in keys:
+            return "file"
+        if "path" in keys:
+            return "path"
+        # fallback: première colonne non-label
+        for k in keys:
+            if k != "label":
+                return k
+        raise KeyError(f"Aucune colonne image trouvée dans {keys}")
+
+    image_key = detect_image_key(dataset["train"].column_names)
+
+    # 5) Transformation robuste
     def transform(batch):
-        images = [img.convert("RGB") for img in batch["image"]]
-        inputs = processor(images, return_tensors="pt")
+        raw_imgs = batch[image_key]
+        pil_images = []
+
+        for elem in raw_imgs:
+            if isinstance(elem, Image.Image):
+                pil_images.append(elem.convert("RGB"))
+            else:
+                pil_images.append(Image.open(elem).convert("RGB"))
+
+        inputs = processor(pil_images, return_tensors="pt")
         inputs["labels"] = batch["label"]
         return inputs
 
     dataset = dataset.with_transform(transform)
 
-    # 3) Charger modèle
+    # 6) Charger modèle pré-entraîné
     model = AutoModelForImageClassification.from_pretrained(
         model_name,
         num_labels=num_labels,
         id2label={i: n for i, n in enumerate(label_names)},
         label2id={n: i for i, n in enumerate(label_names)},
-        ignore_mismatched_sizes=True,  # IMPORTANT POUR ADAPTER 1000 → 2 CLASSES
+        ignore_mismatched_sizes=True,  # indispensable pour adapter 1000 → 2 classes
     )
 
-    # 4) Métriques
+    # 7) Métrique
     metric = evaluate.load("accuracy")
 
     def compute_metrics(eval_pred):
@@ -103,7 +162,7 @@ def train_model(zip_dataset_path, model_name, epochs, batch_size, lr):
         preds = np.argmax(logits, axis=-1)
         return metric.compute(predictions=preds, references=labels)
 
-    # 5) TrainingArguments
+    # 8) TrainingArguments
     args = TrainingArguments(
         output_dir=TRAIN_OUTPUT_DIR,
         num_train_epochs=int(epochs),
@@ -117,7 +176,7 @@ def train_model(zip_dataset_path, model_name, epochs, batch_size, lr):
         report_to=[],
     )
 
-    # 6) Trainer
+    # 9) Trainer
     trainer = Trainer(
         model=model,
         args=args,
@@ -128,11 +187,11 @@ def train_model(zip_dataset_path, model_name, epochs, batch_size, lr):
 
     trainer.train()
 
-    # 7) Sauvegarde finale
+    # 10) Sauvegarde finale
     model.save_pretrained(TRAIN_OUTPUT_DIR)
     processor.save_pretrained(TRAIN_OUTPUT_DIR)
 
-    return f"Entraînement terminé. Modèle sauvegardé dans : {TRAIN_OUTPUT_DIR}"
+    return debug_log + f"\nEntraînement terminé. Modèle sauvegardé dans : {TRAIN_OUTPUT_DIR}"
 
 
 # =====================================================================
@@ -140,10 +199,8 @@ def train_model(zip_dataset_path, model_name, epochs, batch_size, lr):
 # =====================================================================
 
 def extract_model(zip_model_path):
-    """Dézippe un modèle téléchargé côté utilisateur."""
     if os.path.isdir(MODEL_UPLOAD_DIR):
         shutil.rmtree(MODEL_UPLOAD_DIR)
-
     os.makedirs(MODEL_UPLOAD_DIR, exist_ok=True)
 
     try:
@@ -156,18 +213,17 @@ def extract_model(zip_model_path):
 
 def predict(model_zip_path, image):
     if model_zip_path is None:
-        return "Erreur : aucun modèle importé."
+        return "Erreur : aucun modèle ZIP fourni."
 
     success, err = extract_model(model_zip_path)
     if not success:
         return err
 
-    # Charger modèle
     try:
         model = AutoModelForImageClassification.from_pretrained(MODEL_UPLOAD_DIR)
         processor = AutoImageProcessor.from_pretrained(MODEL_UPLOAD_DIR)
     except Exception as e:
-        return f"Erreur de chargement du modèle : {e}"
+        return f"Erreur lors du chargement du modèle : {e}"
 
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
@@ -190,20 +246,18 @@ with gr.Blocks() as demo:
 
     gr.Markdown("# Classification de Soudures — Entraînement & Inférence")
 
-    # -----------------------------------------------------------------
+    # -------------------------------------------------------------
     # ONGLET 1 : ENTRAÎNEMENT
-    # -----------------------------------------------------------------
+    # -------------------------------------------------------------
     with gr.Tab("1 • Entraîner un modèle"):
-        gr.Markdown("### Importer un dataset ZIP au format : Bonne/ et Mauvaise/")
-
-        dataset_zip = gr.File(label="Dataset .zip", type="filepath")
+        dataset_zip = gr.File(label="Dataset ZIP (Bonne/ et Mauvaise/)", type="filepath")
         model_name = gr.Textbox(label="Modèle de départ", value=DEFAULT_MODEL_NAME)
-        epochs = gr.Slider(label="Nombre d'époques", minimum=1, maximum=50, value=5)
+        epochs = gr.Slider(label="Époques", minimum=1, maximum=50, value=5)
         batch = gr.Slider(label="Batch size", minimum=2, maximum=64, value=8)
         lr = gr.Number(label="Learning rate", value=5e-5)
 
         train_btn = gr.Button("Lancer l'entraînement")
-        train_log = gr.Textbox(label="Logs", lines=8)
+        train_log = gr.Textbox(label="Logs", lines=10)
 
         train_btn.click(
             train_model,
@@ -211,16 +265,14 @@ with gr.Blocks() as demo:
             outputs=train_log
         )
 
-    # -----------------------------------------------------------------
+    # -------------------------------------------------------------
     # ONGLET 2 : INFÉRENCE
-    # -----------------------------------------------------------------
+    # -------------------------------------------------------------
     with gr.Tab("2 • Tester un modèle"):
-        gr.Markdown("### Importer un modèle ZIP (export HF) + une image de soudure")
-
-        model_zip = gr.File(label="Modèle .zip", type="filepath")
+        model_zip = gr.File(label="Modèle ZIP", type="filepath")
         input_image = gr.Image(label="Image de soudure")
         predict_btn = gr.Button("Prédire")
-        result = gr.Label(label="Probabilités")
+        result = gr.Label(label="Résultat")
 
         predict_btn.click(
             predict,
@@ -229,6 +281,5 @@ with gr.Blocks() as demo:
         )
 
 
-# Lancer l’app dans un Space HF
 if __name__ == "__main__":
     demo.launch()

dataset.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7367892765884ad898bb791e05d4900fd9c76f38dcf56bc69542287853b9414c
+size 488258521