import gradio as gr
from PIL import Image
import torch
import numpy as np
from rembg import remove
from transformers import CLIPProcessor, CLIPModel

# Lade das Modell und den Prozessor
model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")
processor = CLIPProcessor.from_pretrained("patrickjohncyh/fashion-clip")

# Prompts für jede Merkmalsgruppe
category_prompts = ["a t-shirt", "a long-sleeved shirt", "a hoodie", "a sweatshirt", "a pullover", "a tank top"]
color_prompts = ["a red garment", "a blue garment", "a black garment", "a white garment", "a green garment", "a yellow garment", "a gray garment", "a brown garment", "a pink garment", "a purple garment"]
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"]

# Hilfsfunktion: finde das passendste Prompt für eine Gruppe
def predict_best_prompt(image, prompts):
    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
        probs = logits_per_image.softmax(dim=1).squeeze().tolist()
    best_idx = torch.tensor(probs).argmax().item()
    return prompts[best_idx], probs[best_idx]

# Hauptfunktion für die App
def analyze_image(image):
    if image is None:
        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(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 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=[gr.Image(label="Segmentiertes Bild"), gr.Textbox(label="Vorhersage")],
    title="Fashion Attribute Predictor (Prototype 2 + Segmentation)",
    description="Das Modell entfernt zuerst den Hintergrund (Segmentierung) und erkennt dann Attribute mit FashionCLIP."
)

if __name__ == "__main__":
    iface.launch()