import gradio as gr
from PIL import Image
import torch
from transformers import ViTFeatureExtractor, ViTForImageClassification
from accessories import recommend_accessories

# Load ViT Model for style classification
def load_model():
    feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
    model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
    return feature_extractor, model

extractor, model = load_model()

def analyze_style(image):
    if image is None:
        return "Please upload an image.", None, None

    inputs = extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    predicted_class = outputs.logits.argmax(-1).item()
    style_name = model.config.id2label[predicted_class]

    style_label = style_name.lower()
    rec = recommend_accessories(style_label)

    # Return predicted style string, recommendation string, and the image
    return f"**Predicted Style Class:** {style_name}", rec, image

title = "👗 StyleGuru: AI-Enhanced Fashion Designer"
description = """
**StyleGuru** helps fashion enthusiasts and designers analyze garment styles and get accessory & fabric recommendations. Upload a photo or sketch, and let AI do the magic!

**How to use:**
1. Upload a clear image or sketch of a garment.
2. View the predicted style.
3. See recommended accessories and fabrics to enhance your design.
"""

with gr.Blocks() as demo:
    # CSS to hide webcam and paste buttons
    gr.HTML("""
    <style>
    button[data-testid="paste-button"], 
    button[data-testid="webcam-button"] {
        display: none !important;
    }
    </style>
    """)

    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    with gr.Row():
        image_input = gr.Image(label="Upload a garment image or sketch", type="pil")
        with gr.Column():
            style_output = gr.Markdown(label="Style Analysis")
            rec_output = gr.Markdown(label="💍 Accessory & Fabric Recommendation")

    analyze_button = gr.Button("Analyze Style")
    analyze_button.click(
        fn=analyze_style,
        inputs=image_input,
        outputs=[style_output, rec_output, image_input],
    )

demo.launch()