# fashionclip_app.py
import gradio as gr
from PIL import Image
import torch
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):
    print(f"[DEBUG] Image type: {type(image)}, Prompt count: {len(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 "⚠️ Please upload or take a picture first."

    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)

    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})
    """

# 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",
    title="Fashion Attribute Predictor (Prototype 2)",
    description="Upload or capture an image of a t-shirt or pullover. The model predicts category, color, pattern, and fit using FashionCLIP."
)

# App starten
if __name__ == "__main__":
    iface.launch()