Update app.py

app.py

@@ -1,4 +1,3 @@
-# fashionclip_segmented_app.py
 import gradio as gr
 from PIL import Image
 import torch
@@ -16,17 +15,9 @@ color_prompts = ["a red garment", "a blue garment", "a black garment", "a white
 pattern_prompts = ["a plain shirt", "a striped shirt", "a floral shirt", "a checked shirt", "a dotted shirt", "an abstract patterned shirt"]
 fit_prompts = ["a slim fit shirt", "an oversized top", "a regular fit shirt", "a cropped shirt", "a shirt with a crew neck", "a shirt with a v-neck", "a shirt with a round neckline"]
 
-# Funktion zur Hintergrundentfernung
-def remove_background(image: Image.Image) -> Image.Image:
-    image_np = np.array(image)
-    image_no_bg_np = remove(image_np)
-    return Image.fromarray(image_no_bg_np)
-
 # Hilfsfunktion: finde das passendste Prompt für eine Gruppe
 def predict_best_prompt(image, prompts):
-    print(f"[DEBUG] Image type: {type(image)}, Prompt count: {len(prompts)}")
-    image_clean = remove_background(image)
-    inputs = processor(text=prompts, images=[image_clean], return_tensors="pt", padding=True)
+    inputs = processor(text=prompts, images=[image], return_tensors="pt", padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
         logits_per_image = outputs.logits_per_image
@@ -37,30 +28,36 @@ def predict_best_prompt(image, prompts):
 # Hauptfunktion für die App
 def analyze_image(image):
     if image is None:
-        return "⚠️ Please upload or take a picture first."
+        return None, "⚠️ Please upload or take a picture first."
+    # Hintergrund einmal entfernen!
+    image_np = np.array(image)
+    image_no_bg_np = remove(image_np)
+    segmented_image = Image.fromarray(image_no_bg_np)
 
+    # Vorhersage mit CLIP auf dem segmentierten Bild
     results = {}
-    results["Category"], cat_score = predict_best_prompt(image, category_prompts)
-    results["Color"], color_score = predict_best_prompt(image, color_prompts)
-    results["Pattern"], pattern_score = predict_best_prompt(image, pattern_prompts)
-    results["Fit"], fit_score = predict_best_prompt(image, fit_prompts)
+    results["Category"], cat_score = predict_best_prompt(segmented_image, category_prompts)
+    results["Color"], color_score = predict_best_prompt(segmented_image, color_prompts)
+    results["Pattern"], pattern_score = predict_best_prompt(segmented_image, pattern_prompts)
+    results["Fit"], fit_score = predict_best_prompt(segmented_image, fit_prompts)
+
+    text_result = (
+        f"Category: {results['Category']} ({cat_score:.2f})\n"
+        f"Color: {results['Color']} ({color_score:.2f})\n"
+        f"Pattern: {results['Pattern']} ({pattern_score:.2f})\n"
+        f"Fit: {results['Fit']} ({fit_score:.2f})"
+    )
 
-    return f"""
-    Category: {results['Category']} ({cat_score:.2f})\n
-    Color: {results['Color']} ({color_score:.2f})\n
-    Pattern: {results['Pattern']} ({pattern_score:.2f})\n
-    Fit: {results['Fit']} ({fit_score:.2f})
-    """
+    return segmented_image, text_result
 
 # Gradio UI erstellen
 iface = gr.Interface(
     fn=analyze_image,
     inputs=gr.Image(type="pil", label="Upload or take a picture", sources=["upload", "webcam"]),
-    outputs="text",
+    outputs=[gr.Image(label="Segmentiertes Bild"), gr.Textbox(label="Vorhersage")],
     title="Fashion Attribute Predictor (Prototype 2 + Segmentation)",
-    description="Upload or capture an image of a t-shirt or pullover. The model first removes the background and then predicts category, color, pattern, and fit using FashionCLIP."
+    description="Das Modell entfernt zuerst den Hintergrund (Segmentierung) und erkennt dann Attribute mit FashionCLIP."
 )
 
-# App starten
 if __name__ == "__main__":
     iface.launch()