Update app.py

app.py

@@ -2,34 +2,159 @@ import gradio as gr
 from transformers import pipeline
 from PIL import Image
 import numpy as np
+import cv2
+import time
 
-# Liste des catégories
-CATEGORIES = ["t-shirt", "dress", "jeans", "jacket", "skirt", "shoes"]
+# Import des catégories
+from categories import FASHION_CATEGORIES
 
-# Charger le modèle (version simplifiée)
-model = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32")
+# Initialisation des modèles
+print("🔧 Loading models...")
+try:
+    # Modèle de segmentation
+    seg_pipe = pipeline(
+        "image-segmentation", 
+        model="mattmdjaga/segformer_b2_clothes",
+        device=-1  # Force l'utilisation du CPU pour plus de stabilité
+    )
+    
+    # Modèle de classification
+    class_pipe = pipeline(
+        "zero-shot-image-classification", 
+        model="openai/clip-vit-base-patch32",
+        device=-1  # Force l'utilisation du CPU
+    )
+    print("✅ Models loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading models: {e}")
+    raise e
+
+def process_image(input_image):
+    """Traite l'image et retourne les résultats"""
+    try:
+        if input_image is None:
+            return "⚠️ Please upload an image first", None, None
+        
+        # Conversion en PIL Image
+        if isinstance(input_image, np.ndarray):
+            pil_image = Image.fromarray(input_image)
+        else:
+            pil_image = input_image
+        
+        # Redimensionnement pour éviter les problèmes de mémoire
+        pil_image = pil_image.resize((224, 224))
+        
+        # Étape 1: Segmentation
+        print("🔍 Segmenting image...")
+        segments = seg_pipe(pil_image)
+        
+        if not segments:
+            return "❌ No clothing detected", None, None
+        
+        # Trouver le plus grand segment
+        largest_segment = max(segments, key=lambda x: np.sum(x['mask']))
+        mask = largest_segment['mask']
+        
+        # Étape 2: Extraction du vêtement
+        mask_np = np.array(mask).astype(np.uint8) * 255
+        masked_image = cv2.bitwise_and(np.array(pil_image), np.array(pil_image), mask=mask_np)
+        masked_pil = Image.fromarray(masked_image)
+        
+        # Étape 3: Classification
+        print("📊 Classifying...")
+        predictions = class_pipe(
+            masked_pil, 
+            candidate_labels=FASHION_CATEGORIES,
+            hypothesis_template="This is a photo of {}"
+        )
+        
+        # Formatage des résultats
+        result_text = "🎯 Classification Results:\n\n"
+        for i, pred in enumerate(predictions[:3]):
+            result_text += f"{i+1}. {pred['label']}: {pred['score']*100:.1f}%\n"
+        
+        return result_text, masked_pil, pil_image
+        
+    except Exception as e:
+        return f"❌ Error: {str(e)}", None, None
 
-def classify_simple(image):
-    """Version simplifiée sans segmentation"""
-    if image is None:
-        return "Please upload an image"
+# Interface Gradio améliorée
+with gr.Blocks(
+    title="Fashion Classifier", 
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {max-width: 900px !important;}
+    """
+) as demo:
     
-    results = model(image, candidate_labels=CATEGORIES)
+    gr.Markdown("""
+    # 👗 Fashion Category Classifier
+    Upload a picture of clothing. The AI will detect and classify it.
+    """)
     
-    output = "Results:\n"
-    for result in results:
-        output += f"{result['label']}: {result['score']*100:.1f}%\n"
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(
+                label="📤 Upload Image",
+                type="pil",
+                height=200
+            )
+            
+            process_btn = gr.Button(
+                "🚀 Process Image", 
+                variant="primary",
+                size="lg"
+            )
+        
+        with gr.Column():
+            output_text = gr.Textbox(
+                label="📊 Results",
+                lines=5,
+                interactive=False
+            )
+            
+            with gr.Row():
+                original_output = gr.Image(
+                    label="Original",
+                    type="pil",
+                    height=200,
+                    interactive=False
+                )
+                masked_output = gr.Image(
+                    label="Detected Item",
+                    type="pil",
+                    height=200,
+                    interactive=False
+                )
     
-    return output
-
-# Interface minimaliste
-iface = gr.Interface(
-    fn=classify_simple,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(),
-    title="Simple Fashion Classifier",
-    description="Upload a clothing item to classify it"
-)
+    # Instructions
+    gr.Markdown("""
+    ### 📝 Instructions:
+    1. Upload an image of clothing
+    2. Click 'Process Image' 
+    3. See the classification results
+    """)
+    
+    # Lier les événements
+    process_btn.click(
+        fn=process_image,
+        inputs=image_input,
+        outputs=[output_text, masked_output, original_output]
+    )
+    
+    # Exemple de texte
+    gr.Markdown("""
+    ### 💡 Tips:
+    - Use clear, well-lit photos
+    - Focus on one clothing item at a time
+    - Avoid busy backgrounds for better results
+    """)
 
+# Lancement de l'application
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    )