import gradio as gr
import open_clip
import torch
import requests
import numpy as np
from PIL import Image
from io import BytesIO

# Sidebar content
sidebar_markdown = """

Note, this demo can classify 300 items. If you didn't find what you're looking for, reach out to us on our [Community](https://join.slack.com/t/marqo-community/shared_invite/zt-2iab0260n-QJrZLUSOJYUifVxf964Gdw) and request an item to be added.

## Documentation
📚 [Blog Post](https://www.marqo.ai/blog/search-model-for-fashion)

📝 [Use Case Blog Post](https://www.marqo.ai/blog/ecommerce-image-classification-with-marqo-fashionclip)

## Code
💻 [GitHub Repo](https://github.com/marqo-ai/marqo-FashionCLIP)

🤝 [Google Colab](https://colab.research.google.com/drive/1nq978xFJjJcnyrJ2aE5l82GHAXOvTmfd?usp=sharing)

🤗 [Hugging Face Collection](https://huggingface.co/collections/Marqo/marqo-fashionclip-and-marqo-fashionsiglip-66b43f2d09a06ad2368d4af6)

## Citation
If you use Marqo-FashionSigLIP or Marqo-FashionCLIP, please cite us:
```
@software{Jung_Marqo-FashionCLIP_and_Marqo-FashionSigLIP_2024,
author = {Jung, Myong Chol and Clark, Jesse},
month = aug,
title = {{Marqo-FashionCLIP and Marqo-FashionSigLIP}},
url = {https://github.com/marqo-ai/marqo-FashionCLIP},
version = {1.0.0},
year = {2024}
```
"""

# List of fashion items
items = [
    'abaya', 'anorak', 'apron', 'ascot tie',
    'balaclava', 'ball gown', 'bandanna', 'baseball cap', 'bathing suit', 
    'beanie', 'bedclothes', 'bell-bottoms', 'belt', 'beret', 
    'Bermuda shorts', 'baby clothes', 'bib', 'bikini', 'blazer', 'bloomers', 'blouse', 'boa', 
    'bonnet', 'boot', 'bow', 'bow tie', 'boxer shorts', 'boxers', 'bra', 
    'bracelet', 'brassiere', 'breeches', 'briefs', 'buckle', 'button', 
    'caftan', 'camisole', 'camouflage', 'cap', 
    'cap and gown', 'cape', 'capris', 'cardigan', 'chemise', 'cloak', 'clogs', 
    'coat', 'collar', 'corset', 'costume', 'coveralls', 
    'cowboy boots', 'cowboy hat', 'cravat', 'crown', 'cuff', 'cuff links', 
    'culottes', 'dashiki', 'diaper', 'dinner jacket', 'dirndl', 
    'drawers', 'dress', 'dress shirt', 'duds', 'dungarees', 'earmuffs', 
    'earrings', 'elastic', 'evening gown', 'fashion', 'fedora', 
    'fez', 'flak jacket', 'flannel nightgown', 'flannel shirt', 'flip-flops', 
    'formal wear', 'frock', 'fur', 'fur coat', 'gabardine', 'gaiters', 
    'galoshes', 'garb', 'garters', 'getup', 'gilet', 
    'girdle', 'glasses', 'gloves', 'gown', 'halter top', 'handbag', 
    'handkerchief', 'hat', 'Hawaiian shirt', 'hazmat suit', 'headscarf', 
    'helmet', 'hem', 'high heels', 'hoodie', 'hook and eye', 'hose', 
    'hosiery', 'hospital gown', 'houndstooth', 'housecoat', 'jacket', 
    'jeans', 'jersey', 'jewelry', 'jodhpurs', 'jumper', 'jumpsuit', 'kerchief', 
    'khakis', 'kilt', 'kimono', 'kit', 'knickers', 'lab coat', 'lapel', 
    'leather jacket', 'leg warmers', 'leggings', 'leotard', 'life jacket', 
    'lingerie', 'loafers', 'loincloth', 'long johns', 'long underwear', 
    'miniskirt', 'mittens', 'moccasins', 'muffler', 'muumuu', 'neckerchief', 
    'necklace', 'nightgown', 'nightshirt', 'onesies', 'outerwear', 'outfit', 
    'overalls', 'overcoat', 'overshirt', 'pajamas', 'pants', 
    'pantsuit', 'pantyhose', 'parka', 'pea coat', 'peplum', 'petticoat', 
    'pinafore', 'pleat', 'pocket', 'pocketbook', 'polo shirt', 'poncho', 
    'poodle skirt', 'pullover', 'pumps', 'purse', 'raincoat', 
    'ring', 'robe', 'rugby shirt', 'sandals', 'sari', 'sarong', 'scarf', 
    'school uniform', 'scrubs', 'shawl', 'shirt', 
    'shoes', 'shorts', 'shoulder pads', 'shrug', 'singlet', 'skirt', 
    'slacks', 'slip', 'slippers', 'smock', 'snaps', 'sneakers', 'socks', 
    'sombrero', 'spacesuit', 'stockings', 'stole', 'suit', 
    'sun hat', 'sunbonnet', 'sundress', 'sunglasses', 'suspenders', 
    'sweater', 'sweatpants', 'sweatshirt', 'sweatsuit', 'swimsuit', 
    'T-shirt', 'tam', 'tank top', 'teddy', 'threads', 'tiara', 'tie', 
    'tie clip', 'tights', 'toga', 'tog', 'top', 'top coat', 'top hat', 'train', 
    'trench coat', 'trousers', 'trunks', 'tube top', 'tunic', 'turban', 
    'turtleneck', 'turtleneck shirt', 'tutu', 'tuxedo', 'tweed jacket', 
    'twin set', 'umbrella', 'underclothes', 'undershirt', 
    'underwear', 'uniform', 'veil', 'Velcro', 'vest', 'vestments', 'visor', 
    'waders', 'waistcoat', 'wear', 'wedding gown', 'Wellingtons', 'wetsuit', 
    'white tie', 'wig', 'windbreaker', 'woolens', 'wrap', 'yoke', 'zipper', 
    'zoris', 'jogger', 'palazzo', 'cargo', 'dresspants', 'chinos', 
    'crop top', 'romper', 'insulated jacket', 'fleece', 'rain jacket', 
    'running jacket', 'graphic top', 'pant', 'legging', 'skort', 'brief', 
    'sports bra', 'water shorts', 'cover up', 'goggle', 'glove', 'mitten', 
    'leg gaiter', 'neck gaiter', 'watch', 'bag', 'swim trunk', 
    'pocket watch', 'insoles', "climbing shoes",
]


# Initialize the model and tokenizer
model_name = 'hf-hub:Marqo/marqo-fashionSigLIP'
model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(model_name)
tokenizer = open_clip.get_tokenizer(model_name)

def generate_description(item):
    if "Pants" in item or item in ["Leggings", "Jogger", "Cargo", "Chinos", "Palazzo", "Dresspants", "Sweatpants", "Pant", "Legging", "Skort", "Trouser"]:
        return f"A pair of {item} pants"
    elif item in ["Dress", "Blouse", "T-Shirt", "Tank Top", "Sweater", "Cardigan", "Hoodie", "Coat", "Jacket", "Polo Shirt", "Crop Top", "Romper", "Blazer", "Vest", "Bodysuit", "Maxi Dress", "Graphic Top", "Shirt", "Base Layer Top", "Base Layer Bottom", "Swimsuit", "Rashguard", "Cover Up", "Smock", "Tuxedo"]:
        return f"A {item}"
    elif item in ["Hat", "Sunglasses", "Glasses", "Sun Hat", "Goggle", "Balaclava"]:
        return f"A {item} worn on the head or face"
    elif item in ["Shoes", "Sandals", "Heels", "Trainers", "Boots", "Slippers", "Sneakers", "Insoles", "Socks"]:
        return f"A pair of {item} worn on the feet"
    elif item in ["Jeans", "Skirt", "Shorts", "Dungarees", "Poncho", "Overalls", "Brief", "Boxer", "Swim Trunk", "Ring", "Necklace", "Earing", "Pocket Watch"]:
        return f"A {item} piece of clothing"
    elif item in ["Boxing Gloves", "Glove", "Mitten"]:
        return f"An item of {item} worn on the hands"
    else:
        return f"A fashion item called {item}"
items_desc = [generate_description(item) for item in items]
text = tokenizer(items_desc)

# Encode text features (unchanged)
with torch.no_grad(), torch.amp.autocast('cuda'):
    text_features = model.encode_text(text)
    text_features /= text_features.norm(dim=-1, keepdim=True)

# Prediction function
def predict(image, url):
    if url:
        response = requests.get(url)
        image = Image.open(BytesIO(response.content))
    
    processed_image = preprocess_val(image).unsqueeze(0)

    with torch.no_grad(), torch.amp.autocast('cuda'):
        image_features = model.encode_image(processed_image)
        image_features /= image_features.norm(dim=-1, keepdim=True)

        text_probs = (100 * image_features @ text_features.T).softmax(dim=-1)
        
        sorted_confidences = sorted(
            {items[i]: float(text_probs[0, i]) for i in range(len(items))}.items(), 
            key=lambda x: x[1], 
            reverse=True
        )
        top_10_confidences = dict(sorted_confidences[:10])
        
    return image, top_10_confidences

# Clear function
def clear_fields():
    return None, ""

# Gradio interface
title = "Fashion Item Classifier with Marqo-FashionSigLIP"
description = "Upload an image or provide a URL of a fashion item to classify it using [Marqo-FashionSigLIP](https://huggingface.co/Marqo/marqo-fashionSigLIP)!"

examples = [
    ["images/dress.jpg", "Dress"],
    ["images/sweatpants.jpg", "Sweatpants"],
    ["images/t-shirt.jpg", "T-Shirt"],
    ["images/hat.jpg", "Hat"],
    ["images/blouse.jpg", "Blouse"],
    ["images/cargo.jpg", "Cargos"],
    ["images/sunglasses.jpg", "Sunglasses"],
    ["images/polo-shirt.jpg", "Polo Shirt"],
]

with gr.Blocks(css="""
    .remove-btn {
        font-size: 24px !important; /* Increase the font size of the cross button */
        line-height: 24px !important;
        width: 30px !important; /* Increase the width */
        height: 30px !important; /* Increase the height */
    }
""") as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(f"# {title}")
            gr.Markdown(description)
            gr.Markdown(sidebar_markdown)
            gr.Markdown(" ", elem_id="vertical-line")  # Add an empty Markdown with a custom ID
        with gr.Column(scale=2):
            input_image = gr.Image(type="pil", label="Upload Fashion Item Image", height=312)
            input_url = gr.Textbox(label="Or provide an image URL")
            with gr.Row():
                predict_button = gr.Button("Classify")
                clear_button = gr.Button("Clear")
            gr.Markdown("Or click on one of the images below to classify it:")
            gr.Examples(examples=examples, inputs=input_image)
            output_label = gr.Label(num_top_classes=6)
            predict_button.click(predict, inputs=[input_image, input_url], outputs=[input_image, output_label])
            clear_button.click(clear_fields, outputs=[input_image, input_url])
        
# Launch the interface
demo.launch()