Create app.py

app.py

@@ -0,0 +1,176 @@
+# app.py
+
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F # 🚨 CORRECTION : Nécessaire pour F.pad
+from torchvision import transforms
+from PIL import Image
+from huggingface_hub import hf_hub_download
+import sys
+import os
+
+# --- 1. DÉFINITION DE L'ARCHITECTURE DRANINA (U-NET) ---
+# Copie exacte de l'architecture entraînée pour garantir la compatibilité
+
+# Bloc de Convolution de Base
+def double_conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(inplace=True),
+        nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(inplace=True)
+    )
+
+# Bloc d'Augmentation (Upsampling)
+class Up(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        # 🚨 CORRECTION : Assurez-vous que l'Upsampling correspond exactement à l'entraînement
+        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+        self.conv = double_conv(in_channels, out_channels) 
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        
+        # Gestion des bords (padding) via torch.nn.functional
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
+                        diffY // 2, diffY - diffY // 2])
+        
+        # Concaténation le long de la dimension des canaux (dim=1)
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+# Le Modèle Dranina U-Net Complet
+class DraninaUnet(nn.Module):
+    def __init__(self, n_channels=3, n_classes=3):
+        super(DraninaUnet, self).__init__()
+        
+        # Encodeur
+        self.inc = double_conv(n_channels, 64)
+        self.down1 = nn.MaxPool2d(2); self.conv1 = double_conv(64, 128)
+        self.down2 = nn.MaxPool2d(2); self.conv2 = double_conv(128, 256)
+        self.down3 = nn.MaxPool2d(2); self.conv3 = double_conv(256, 512)
+        self.down4 = nn.MaxPool2d(2); 
+        self.conv4 = double_conv(512, 1024) 
+        
+        # Décodeur
+        self.up1 = Up(1024, 512)
+        self.up2 = Up(512, 256)
+        self.up3 = Up(256, 128)
+        self.up4 = Up(128, 64)
+        
+        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)
+
+    def forward(self, x):
+        # Chemin avant (skip connections)
+        x1 = self.inc(x)
+        x2 = self.down1(x1); x2 = self.conv1(x2)
+        x3 = self.down2(x2); x3 = self.conv2(x3)
+        x4 = self.down3(x3); x4 = self.conv3(x4)
+        x5 = self.down4(x4)
+        
+        x5 = self.conv4(x5) # Bas du U
+
+        # Chemin retour (avec skip connections)
+        x = self.up1(x5, x4) 
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        
+        logits = self.outc(x)
+        return logits
+
+# --- 2. INITIALISATION ET CHARGEMENT DU MODÈLE ---
+
+MODEL_ID = "Clemylia/Dranina-Mandala-Colorizer" # L'ID de ton dépôt de modèle publié
+MODEL_FILENAME = "pytorch_model.bin"
+IMAGE_SIZE = 256
+DEVICE = torch.device("cpu") 
+
+model = None
+try:
+    # Télécharger le fichier de poids du modèle depuis le Hub
+    model_path = hf_hub_download(repo_id=MODEL_ID, filename=MODEL_FILENAME)
+    print(f"✅ Modèle téléchargé depuis le Hub : {model_path}")
+
+    # Instancier le modèle et charger les poids
+    model = DraninaUnet(n_channels=3, n_classes=3)
+    # L'argument map_location est CRUCIAL pour que ça tourne sur CPU (le standard du Space Gratuit)
+    state_dict = torch.load(model_path, map_location=DEVICE) 
+    model.load_state_dict(state_dict)
+    model.to(DEVICE)
+    model.eval()
+    print("✅ Modèle Dranina prêt pour la prédiction.")
+
+except Exception as e:
+    print(f"❌ Erreur lors du chargement du modèle Dranina : {e}", file=sys.stderr)
+
+# --- 3. FONCTION DE PRÉDICTION ---
+
+def colorize_mandala(input_image_pil):
+    """
+    Prend une image PIL (non coloriée), la passe dans le modèle et retourne 
+    l'image PIL coloriée.
+    """
+    if model is None:
+        return Image.new('RGB', (IMAGE_SIZE, IMAGE_SIZE), color = 'red')
+
+    # Enregistre la taille originale pour la redimensionner à la fin (meilleur rendu)
+    original_size = input_image_pil.size 
+
+    # 1. Préparation de l'image (Transformation)
+    transform = transforms.Compose([
+        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        transforms.ToTensor(), 
+    ])
+    
+    # Assurez-vous que l'entrée est en RGB (3 canaux)
+    input_tensor = transform(input_image_pil.convert('RGB')).unsqueeze(0).to(DEVICE) 
+
+    # 2. Inférence
+    with torch.no_grad():
+        output_tensor = model(input_tensor)
+
+    # 3. Post-traitement (Conversion en image PIL)
+    # Clamp pour s'assurer que les valeurs sont entre 0 et 1
+    output_tensor = torch.clamp(output_tensor.squeeze(0), 0, 1) 
+    
+    # Convertir le tensor (C, H, W) en PIL Image
+    output_image = transforms.ToPILImage()(output_tensor.cpu()) 
+    
+    # Redimensionner l'image de sortie à la taille de l'entrée originale
+    output_image = output_image.resize(original_size)
+
+    return output_image
+
+# --- 4. INTERFACE GRADIO ---
+
+if model is not None:
+    # Description pour le Space
+    title = "🎨 Dranina : Mandala Colorizer"
+    description = (
+        "Ceci est une démonstration du modèle **Dranina U-Net**, entraîné sur votre dataset pour colorier automatiquement des mandalas. "
+        "Téléversez une image de mandala en noir et blanc pour voir la prédiction du modèle."
+    )
+    
+    # Création de l'interface
+    iface = gr.Interface(
+        fn=colorize_mandala,
+        inputs=gr.Image(type="pil", label="Mandala Non Colorié (Entrée)"),
+        outputs=gr.Image(type="pil", label="Mandala Colorié (Prédiction)"),
+        title=title,
+        description=description,
+        # Ajoute des exemples ici si tu en as dans le Space
+        # examples=["votre_dossier_dans_le_space/exemple1.jpg"],
+        allow_flagging="auto",
+    )
+
+    # Lancement de l'application
+    iface.launch()
+else:
+    print("Application non lancée car le modèle n'a pas pu être chargé. Vérifiez l'ID du modèle et les dépendances.")