Update app.py

app.py

@@ -1,111 +1,125 @@
 import gradio as gr
 from transformers import pipeline
-from PIL import Image
+from PIL import Image, ImageOps, ImageFilter
 import numpy as np
+import cv2
 
-# Liste des catégories de vêtements
+# Catégories précises et distinctes
 FASHION_CATEGORIES = [
-    "t-shirt", "long sleeve shirt", "short sleeve shirt", 
-    "sleeveless shirt", "polo shirt", "sweatshirt", 
-    "hoodie", "sweater", "cardigan", "jacket", "coat", 
-    "blazer", "dress", "long dress", "short dress", 
-    "skirt", "long skirt", "short skirt", "jeans", 
-    "pants", "trousers", "shorts", "leggings", 
-    "sports shoes", "sneakers", "boots", "heels", "sandals"
+    "t-shirt", "button-down shirt", "polo shirt", 
+    "sweatshirt", "hoodie", "sweater",
+    "jacket", "coat", "blazer",
+    "dress", "long dress", "short dress",
+    "skirt", "long skirt", "short skirt",
+    "jeans", "pants", "shorts",
+    "sneakers", "boots", "heels", "sandals"
 ]
 
-# Charger le modèle de classification
-print("Loading classification model...")
-class_pipe = pipeline(
-    "zero-shot-image-classification", 
-    model="openai/clip-vit-base-patch32"
-)
-print("Model loaded successfully!")
+# Chargement des modèles
+print("🔧 Loading models...")
+seg_pipe = pipeline("image-segmentation", model="mattmdjaga/segformer_b2_clothes")
+class_pipe = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
+print("✅ Models loaded!")
 
-def classify_image(input_image):
-    """Fonction simple pour classifier les images"""
+def preprocess_image(image):
+    """Prétraitement avancé de l'image"""
+    # Conversion en numpy array si nécessaire
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
+    # Réduction de taille pour meilleures performances
+    image = image.resize((512, 512))
+    
+    return image
+
+def remove_background(image):
+    """Suppression professionnelle de l'arrière-plan"""
+    # Segmentation
+    segments = seg_pipe(image)
+    if not segments:
+        return image
+    
+    # Trouver le vêtement principal
+    largest_segment = max(segments, key=lambda x: np.sum(x['mask']))
+    mask = np.array(largest_segment['mask'])
+    
+    # Application du masque
+    image_np = np.array(image)
+    masked_image = np.zeros_like(image_np)
+    masked_image[mask > 0] = image_np[mask > 0]
+    
+    return Image.fromarray(masked_image)
+
+def classify_fashion(image):
+    """Classification précise"""
     try:
-        if input_image is None:
-            return "Please upload an image first"
+        # Préprocessing
+        processed_image = preprocess_image(image)
         
-        # Convertir en format PIL si nécessaire
-        if isinstance(input_image, np.ndarray):
-            input_image = Image.fromarray(input_image)
+        # Suppression de l'arrière-plan
+        isolated_image = remove_background(processed_image)
         
-        # Redimensionner pour de meilleures performances
-        input_image = input_image.resize((224, 224))
-        
-        # Classification
+        # Classification avec paramètres optimisés
         predictions = class_pipe(
-            input_image, 
+            isolated_image, 
             candidate_labels=FASHION_CATEGORIES,
-            hypothesis_template="This is a photo of {}"
+            hypothesis_template="a clear photo of {}",
+            multi_label=False
         )
         
-        # Formater les résultats
-        result_text = "👗 Classification Results:\n\n"
-        for i, pred in enumerate(predictions[:5]):
-            result_text += f"{i+1}. {pred['label']}: {pred['score']*100:.1f}%\n"
+        # Filtrage des résultats (seulement > 10% de confiance)
+        filtered_predictions = [p for p in predictions if p['score'] > 0.1]
+        
+        if not filtered_predictions:
+            return "❌ No confident prediction. Try a clearer image.", isolated_image
+        
+        # Formatage des résultats
+        result_text = "🎯 **Top Predictions:**\n\n"
+        for i, pred in enumerate(filtered_predictions[:3]):
+            result_text += f"{i+1}. **{pred['label']}**: {pred['score']*100:.1f}%\n"
         
-        return result_text
+        return result_text, isolated_image
         
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"❌ Error: {str(e)}", None
 
-# Interface Gradio simple
-with gr.Blocks(title="Fashion Classifier", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 👗 Fashion Category Classifier
-    Upload a picture of clothing to classify it.
-    """)
+# Interface améliorée
+with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 800px;}") as demo:
+    gr.Markdown("# 👗 **Fashion AI - Professional Classifier**")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(
-                label="📤 Upload Clothing Image",
-                type="pil",
-                height=300
-            )
+            gr.Markdown("### 📤 Upload Image")
+            image_input = gr.Image(type="pil", label="Clothing Image")
+            process_btn = gr.Button("🚀 Analyze Image", variant="primary")
             
-            classify_btn = gr.Button(
-                "🔍 Classify Image", 
-                variant="primary",
-                size="lg"
-            )
-        
         with gr.Column():
-            output_text = gr.Textbox(
-                label="📊 Results",
-                lines=8,
-                interactive=False
-            )
+            gr.Markdown("### 📊 Results")
+            output_text = gr.Markdown(label="Analysis")
+            output_image = gr.Image(label="Processed Image", interactive=False)
     
     # Instructions
     gr.Markdown("""
-    ### 📝 How to use:
-    1. Upload an image of a clothing item
-    2. Click the 'Classify Image' button
-    3. See the classification results
+    ### 💡 **Pro Tips for Best Results:**
+    - ✅ Use clear, well-lit photos
+    - ✅ Center the clothing item
+    - ✅ Use plain backgrounds when possible
+    - ✅ Avoid multiple items in one photo
+    - ❌ Don't use blurry or dark images
     """)
     
-    # Lier le bouton à la fonction
-    classify_btn.click(
-        fn=classify_image,
+    # Events
+    process_btn.click(
+        fn=classify_fashion,
         inputs=image_input,
-        outputs=output_text
+        outputs=[output_text, output_image]
     )
     
-    # Ajouter aussi le changement sur l'upload
     image_input.upload(
-        fn=classify_image,
+        fn=classify_fashion,
         inputs=image_input,
-        outputs=output_text
+        outputs=[output_text, output_image]
     )
 
-# Lancer l'application
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    demo.launch()