app.py

import io
import os
import ast
import json
import base64
import spaces
import requests
import numpy as np
import gradio as gr
from PIL import Image
from io import BytesIO
import face_recognition
from turtle import title
from openai import OpenAI
from collections import Counter, defaultdict
from typing import List, Optional, Set, Dict, Any
from transformers import pipeline

import urllib.request
from transformers import YolosImageProcessor, YolosForObjectDetection
import torch
import matplotlib.pyplot as plt
from torchvision.transforms import ToTensor, ToPILImage

    
client = OpenAI()

pipe = pipeline("zero-shot-image-classification", model="patrickjohncyh/fashion-clip")

color_file_path = 'color_config.json'
attributes_file_path = 'attributes_config.json'
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")


# Open and read the COLOR JSON file
with open(color_file_path, 'r') as file:
    color_data = json.load(file)

# Open and read the ATTRIBUTES JSON file
with open(attributes_file_path, 'r') as file:
    attributes_data = json.load(file)

COLOURS_DICT = color_data['color_mapping']
ATTRIBUTES_DICT = attributes_data['attribute_mapping']

DETAILS_THRESHOLD = 0.4 # This is how high the total score of an additional detail attribute should be for it to be included (number scales up linearly with more images)
OPENAI_USER_PROMPT = "I have a product called 'Round Neck Knitted Ribbed Crop Top with a Neck Detail'. The description reads: 'Introducing the Round Neck Knitted Ribbed Crop Top with a Neck Detail. This sleeveless top is a must-have addition to your wardrobe. Its round neck and ribbed design give it a modern and stylish look, while the neck detail adds a unique touch. Made from high-quality materials, this crop top is comfortable to wear all day long. The solid white color makes it easy to pair with any outfit, making it a versatile piece for any occasion. Whether you're dressing up for a night out or keeping it casual, this crop top is sure to become one of your favorite go-to staples. Upgrade your style with this trendy and chic top today!'. Arabic naming: 'بلوزة قصيرة محبوكة ومضلعة بياقة دائرية وتفاصيل رقبة'. Arabic description: 'نقدم لك بلوزة قصيرة محبوكة ومضلعة برقبة دائرية مع تفاصيل الرقبة. يعد هذا القميص بلا أكمام إضافة لا غنى عنها لخزانة ملابسك. تمنحه رقبته المستديرة وتصميمه المضلع مظهرًا عصريًا وأنيقًا، بينما تضيف تفاصيل الرقبة لمسة فريدة. مصنوع من مواد عالية الجودة، هذا القميص القصير مريح للارتداء طوال اليوم. اللون الأبيض الصلب يجعل من السهل ارتداؤه مع أي ملابس، مما يجعله قطعة متعددة الاستخدامات لأي مناسبة. سواء كنت ترتدي ملابس ليلية أو تبقيها غير رسمية، فمن المؤكد أن هذه السترة القصيرة ستصبح واحدة من القطع الأساسية المفضلة لديك. قم بترقية أسلوبك مع هذا القميص العصري والأنيق اليوم!'\n    The attributes are:\n\t• Category: 'top > croptop'\n\t• Material: 'polyamide'\n\t• Target Audience: women\n\t• Fabric Style: Knitted\n\t• Pattern: Solid\n\t• Colar Type: Round Neck\n\t• Sleeve Length: Sleeveless\n\t• Details: Ribbed Neck Detail \n\t• Tags: 'neutral, staples, casual, springsummer, bestseller, summer_light_palette, coastal_charm, allnewin' (do not use any seasonal tags or other discount info nothing with '_' or color)\n    \n\nI have a product called 'Crew Neck Crop Top with a V-Shaped Front'. The description reads: 'Step up your casual style with this Crew Neck Crop Top with a V-Shaped Front. Crafted from high-quality materials, this sleeveless top features a flattering v-shaped front design that adds a touch of sophistication to any outfit. The Sacramento Green color adds a pop of freshness and pairs perfectly with neutral tones. Whether you're lounging at home or running errands, this crop top is a staple in your wardrobe. Its versatile design makes it easy to dress up or down for any occasion. Elevate your loungewear game and add this Crew Neck Crop Top to your collection today.'. Arabic naming: 'بلوزة قصيرة بياقة دائرية ومقدمة على شكل حرف V'. Arabic description: 'ارتقي بأسلوبك الكاجوال مع هذا القميص القصير ذو الرقبة الدائرية والمقدمة على شكل حرف V. مصنوع من مواد عالية الجودة، يتميز هذا القميص بدون أكمام بتصميم أمامي جذاب على شكل حرف V يضيف لمسة من الرقي إلى أي ملابس. يضيف لون Sacramento Green لمسة من النضارة ويمتزج بشكل مثالي مع الألوان المحايدة. سواء كنت تسترخي في المنزل أو تقوم بمهمات، فإن هذا القميص القصير هو قطعة أساسية في خزانة ملابسك. تصميمه متعدد الاستخدامات يجعل من السهل ارتداؤه لأعلى أو لأسفل في أي مناسبة. ارفع مستوى ملابسك المريحة وأضف هذا القميص ذو الرقبة الدائرية إلى مجموعتك اليوم.'\n    The attributes are:\n\t• Category: 'top > croptop'\n\t• Material: 'cotton, polyester'\n\t• Target Audience: women\n\t• Fabric Style: Woven\n\t• Pattern: Solid\n\t• Colar Type: Crew Neck\n\t• Sleeve Length: Sleeveless\n\t• Details: V-Shaped Front \n\t• Tags: 'neutral, casual, staples, loungewear, bestseller' (do not use any seasonal tags or other discount info nothing with '_' or color)\n    \n\nI have a product called 'Corduroy Polo Collar Buttoned Pocket Jacket'. The description reads: 'Step up your fashion game with this Corduroy Polo Collar Buttoned Pocket Jacket. Made from high-quality corduroy fabric, this jacket is both stylish and comfortable. The unique polo collar adds a touch of sophistication, while the buttoned pockets provide a practical yet fashionable twist. Whether you're heading to a casual outing or a night on the town, this jacket is sure to turn heads. Its versatile design makes it easy to pair with any outfit, giving you endless options for creating trendy looks. Upgrade your wardrobe with this Hiccup fashion statement and be the envy of everyone around you.'. Arabic naming: 'سترة سروال قصير بياقة بولو وجيب بأزرار'. Arabic description: 'ارتقِ بأسلوبك في عالم الموضة مع سترة الجيب ذات الأزرار بياقة بولو من سروال قصير. مصنوع من قماش كوردروي عالي الجودة، هذا الجاكيت أنيق ومريح. تضيف ياقة البولو الفريدة لمسة من الرقي، بينما توفر الجيوب ذات الأزرار لمسة عملية وعصرية. سواء كنت متجهًا إلى نزهة غير رسمية أو ليلة في المدينة، فمن المؤكد أن هذه السترة ستلفت الأنظار. تصميمه متعدد الاستخدامات يجعل من السهل ارتداؤه مع أي ملابس، مما يمنحك خيارات لا حصر لها للحصول على إطلالات عصرية. قم بترقية خزانة ملابسك مع هذا التصميم الأنيق من Hiccup وكن موضع حسد الجميع من حولك.'\n    The attributes are:\n\t• Category: 'outwear > coatjacket'\n\t• Material: 'polyester'\n\t• Target Audience: women\n\t• Fabric Style: Corduroy\n\t• Pattern: Solid\n\t• Colar Type: Polo Collar\n\t• Sleeve Length: Long Sleeve\n\t• Details: Buttoned, Pocket\n\t• Tags: 'autumn, black, cosy, onsale, discount55_2403' (do not use any seasonal tags or other discount info nothing with '_' or color)\n    \n\nI have a product called 'High Waist Skirt with a Belt'. The description reads: 'Step into style with our High Waist Skirt with a Belt in Brown. This chic and versatile skirt is designed to flatter your figure and elevate your fashion game. The high waistline accentuates your curves while the belt adds a touch of sophistication. Made from a premium blend of fabrics, this skirt is comfortable and durable, ensuring that it will become a staple in your wardrobe for years to come. Whether you're dressing up for a special occasion or keeping it casual for everyday wear, this High Waist Skirt with a Belt in Brown is the perfect choice. Embrace timeless elegance and make a statement with this must-have piece.'. Arabic naming: 'تنورة عالية الخصر مع حزام'. Arabic description: 'نقدم لك تنورة عالية الخصر مع حزام، وهي إضافة لا بد منها لأي خزانة ملابس مواكبة للموضة. ارتقي بأسلوبك مع هذه القطعة الأنيقة والمتعددة الاستخدامات التي تجمع بين الرقي والراحة بسهولة. يبرز تصميم الخصر العالي منحنياتك، بينما يشد الحزام المتضمن خصرك للحصول على مظهر جذاب. مصنوعة من قماش عالي الجودة، هذه التنورة تنسدل بشكل جميل ومثالية للمناسبات غير الرسمية والرسمية. ارتديه مع بلوزة وحذاء بكعب لإطلالة أنيقة، أو ارتديه مع تي شيرت وحذاء رياضي لإطلالة أكثر استرخاءً. تمتعي بالأناقة الخالدة مع تنورتنا ذات الخصر العالي مع الحزام واخرجي بأناقة.'\n    The attributes are:\n\t• Category: 'bottom > skirt'\n\t• Material: 'polyester'\n\t• Target Audience: women\n\t• Fabric Style: Woven\n\t• Waist: High Waist\n\t• Style: Mini\n\t• Details: Belt\n\t• Tags: 'black, springsummer, allnewin, neutral, tailoring' (do not use any seasonal tags or other discount info nothing with '_' or color)\n    \n\nI have a product called 'Knitted Crew Neck Crop Top'. The description reads: 'Introducing our Knitted Crew Neck Crop Top, the perfect addition to your wardrobe. Made from high-quality knitted fabric, this crop top offers both style and comfort. The crew neck design adds a touch of sophistication, while the cropped length adds a modern twist. Whether you pair it with high-waisted jeans for a casual look or dress it up with a skirt for a night out, this crop top is versatile and effortlessly chic. The soft fabric feels luxurious against the skin, making it a pleasure to wear all day long. Upgrade your style game with our Knitted Crew Neck Crop Top and turn heads wherever you go.'. Arabic naming: 'محبوك طاقم الرقبة المحاصيل الأعلى'. \nArabic description: 'نقدم لك القميص القصير ذو الرقبة الدائرية المحبوك، الإضافة المثالية لخزانة ملابسك. مصنوع من قماش محبوك عالي الجودة، هذا القميص القصير يوفر الأناقة والراحة. يضيف تصميم الرقبة الدائرية لمسة من الرقي، بينما يضيف الطول القصير لمسة عصرية. سواء كنت ترتديه مع بنطال جينز عالي الخصر لمظهر غير رسمي أو ترتديه مع تنورة لقضاء ليلة في الخارج، فإن هذا القميص القصير متعدد الاستخدامات وأنيق دون عناء. يعطي النسيج الناعم ملمسًا فاخرًا على البشرة، مما يجعل ارتداؤه ممتعًا طوال اليوم. قم بترقية أسلوبك في الأناقة مع القميص المحبوك ذو الرقبة الدائرية وجذب الأنظار أينما ذهبت.'\n    The attributes are:\n\t• Category: 'top > croptop'\n\t• Material: 'viscose, polyamide'\n\t• Target Audience: women\n\t• Fabric Style: Knitted\n\t• Pattern: Solid\n\t• Colar Type: Crew Neck\n\t• Sleeve Length: Sleeveless\n\t• Tags: 'autumn, neutral' (do not use any seasonal tags or other discount info nothing with '_' or color)\n   "


def shot(input, category, level):
    output_dict = {}
    if level == 'variant':
        subColour, mainColour, score = get_colour(ast.literal_eval(str(input)), category)    
        openai_parsed_response = get_openAI_tags(ast.literal_eval(str(input)))
        face_embeddings = get_face_embeddings(ast.literal_eval(str(input)))
        cropped_images = get_cropped_images(ast.literal_eval(str(input)), category)

        # Ensure all outputs are JSON serializable
        output_dict['colors'] = {
            "main": mainColour,
            "sub": subColour,
            "score": score
        }
        output_dict['image_mapping'] = openai_parsed_response
        output_dict['face_embeddings'] = face_embeddings
        output_dict['cropped_images'] = cropped_images

    if level == 'product':
        common_result = get_predicted_attributes(ast.literal_eval(str(input)), category)    
        output_dict['attributes'] = common_result
        output_dict['subcategory'] = category
        print(common_result)
        output_dict['description'] = get_product_description(category=category, attributes=common_result)
    
    # # Convert the dictionary to a JSON-serializable format
    # try:
    #     serialized_output = json.dumps(output_dict)
    # except TypeError as e:
    #     print(f"Serialization Error: {e}")
    #     return {"error": "Serialization failed"}

    return json.dumps(output_dict)



# @spaces.GPU  
# def get_colour(image_urls, category):
#     colourLabels = list(COLOURS_DICT.keys())
#     for i in range(len(colourLabels)):
#         colourLabels[i] = colourLabels[i] + " clothing: " + category

#     responses = pipe(image_urls, candidate_labels=colourLabels)
#     # Get the most common colour
#     mainColour = responses[0][0]['label'].split(" clothing:")[0]


#     if mainColour not in COLOURS_DICT:
#         return None, None, None

#     # Add category to the end of each label
#     labels = COLOURS_DICT[mainColour]
#     for i in range(len(labels)):
#         labels[i] = labels[i] + " clothing: " + category

#     # Run pipeline in one go
#     responses = pipe(image_urls, candidate_labels=labels)
#     subColour = responses[0][0]['label'].split(" clothing:")[0]

#     return subColour, mainColour, responses[0][0]['score']

@spaces.GPU
def get_colour(image_urls, category):
    # Prepare color labels
    colourLabels = [f"{color} clothing: {category}" for color in COLOURS_DICT.keys()]
    print("Colour Labels:", colourLabels)  # Debug: Print colour labels
    print("Image URLs:", image_urls)       # Debug: Print image URLs

    # Split labels into two batches
    mid_index = len(colourLabels) // 2
    first_batch = colourLabels[:mid_index]
    second_batch = colourLabels[mid_index:]

    # Process the first batch
    responses_first_batch = pipe(image_urls, candidate_labels=first_batch)
    # Get the top 3 from the first batch
    top3_first_batch = sorted(responses_first_batch[0], key=lambda x: x['score'], reverse=True)[:3]

    # Process the second batch
    responses_second_batch = pipe(image_urls, candidate_labels=second_batch)
    # Get the top 3 from the second batch
    top3_second_batch = sorted(responses_second_batch[0], key=lambda x: x['score'], reverse=True)[:3]

    # Combine the top 3 from each batch
    combined_top6 = top3_first_batch + top3_second_batch
    # Get the final top 3 from the combined list
    final_top3 = sorted(combined_top6, key=lambda x: x['score'], reverse=True)[:3]

    mainColour = final_top3[0]['label'].split(" clothing:")[0]

    if mainColour not in COLOURS_DICT:
        return None, None, None

    # Get sub-colors for the main color
    labels = [f"{label} clothing: {category}" for label in COLOURS_DICT[mainColour]]
    print("Labels for pipe:", labels)  # Debug: Confirm labels are correct

    responses = pipe(image_urls, candidate_labels=labels)
    subColour = responses[0][0]['label'].split(" clothing:")[0]

    return subColour, mainColour, responses[0][0]['score']



# Function for get_predicted_attributes
def get_most_common_label(responses):
    feature_scores = defaultdict(float)
    for response in responses:
        label, score = response[0]['label'].split(", clothing:")[0], response[0]['score']
        feature_scores[label] += score
    return max(feature_scores, key=feature_scores.get), feature_scores[max(feature_scores, key=feature_scores.get)]

@spaces.GPU  
def get_predicted_attributes(image_urls, category):
    # Assuming ATTRIBUTES_DICT and pipe are defined outside this function
    attributes = list(ATTRIBUTES_DICT.get(category, {}).keys())

    # Mapping of possible values per attribute
    common_result = []
    for attribute in attributes:
        values = ATTRIBUTES_DICT.get(category, {}).get(attribute, [])

        if len(values) == 0:
            continue

        # Adjust labels for the pipeline
        attribute = attribute.replace("colartype", "collar").replace("sleevelength", "sleeve length").replace("fabricstyle", "fabric")
        values = [f"{attribute}: {value.strip()}, clothing: {category}" for value in values]
    
        # Get the predicted values for the attribute
        responses = pipe(image_urls, candidate_labels=values)
        most_common, score = get_most_common_label(responses)
        common_result.append(most_common)

        if attribute == "details":
            # Process additional details labels if the score is higher than 0.8
            for _ in range(2):
                values = [value for value in values if value != f"{most_common}, clothing: {category}"]
                responses = pipe(image_urls, candidate_labels=values)
                most_common, score = get_most_common_label(responses)
                if score > DETAILS_THRESHOLD * len(image_urls):
                    common_result.append(most_common)

    # Convert common_result into a dictionary
    final = {}
    details_count = 0

    for result in common_result:
        result = result.replace("collar", "colartype").replace("sleeve length", "sleevelength").replace("fabric", "fabricstyle")
        key, value = result.split(": ")
        if key == "details":
            if details_count > 0:
                key += str(details_count)
            details_count += 1
        final[key] = value.lower()

    return final



def generate_prompt(category: Optional[str], tags: Optional[Set[str]], 
                    materials: Optional[List[Dict[str, int]]], attributes: Optional[List[Dict[str, str]]]) -> str:

    print(attributes)
    for attr in attributes:
        print(attr)
    formatted_attributes = [f"{attr['key']}: {attr['value']}" for attr in attributes] if attributes else []

    formatted_string = "\\n".join(formatted_attributes) if formatted_attributes else "No attributes provided."

    processed_category = category.replace("women-", "").replace("-", " > ") if category else None

    material_keys = ", ".join([str(material['key']) for material in materials]) if materials else None

    processed_tags = ", ".join(tags) if tags else None

    return (
        f"I have a new product.\\nThe attributes are: Category: {processed_category}\\n"
        f"Material: {material_keys}\\n{formatted_string}\\n"
        f"Tags: {processed_tags} (do not use any seasonal tags or other discount info nothing with '_' or color)\\n"
        "Print the output as a json object no extra text, keys are name_en, name_ar, summary_en, summary_ar"
    ).strip()

def send_request_to_openai(prompt: str) -> Optional[Dict[str, Any]]:
    messages = [
        {"role": "system", "content": "You are a labelling assistant, you will help create product metadata for my online e-commerce platform. Your tasks will be to return product names and product descriptions. You will be given some product attributes and/or metadata guideline outputs and are expected to follow the sample output to generate new outputs.\n\nYou will be asked to generate product names and descriptions in both english and arabic."},
        {"role": "user", "content": "Hello, I need help creating product metadata for new apparel items. I will provide examples, I want you to understand the current structure and setup. You are expected to match the current tone of voice and overall information."},
        {"role": "assistant", "content": "Sure, I'd be happy to help! Please provide the attributes and any existing product names you'd like me to consider"},
        {"role": "user", "content": OPENAI_USER_PROMPT},
        {"role": "assistant", "content": "This was very helpful, I have taken note of the overall data inputs and expected product name and description outputs. Please provide the attributes for the newly listed product so I can generate the name and description for you."},
        {"role": "user", "content": prompt}
    ]

    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo-16k",
            messages=messages,
            temperature=1.0,
            max_tokens=512,
            top_p=1.0,
            frequency_penalty=0.0,
            presence_penalty=0.0
        )

        json_string = response['choices'][0]['message']['content']
        parsed_response = json.loads(json_string)
        
        print(f"English name: {parsed_response['name_en']}")
        print(f"Arabic name: {parsed_response['name_ar']}")
        print(f"English description: {parsed_response['summary_en']}")
        print(f"Arabic description: {parsed_response['summary_ar']}")
        
        return parsed_response

    except openai.error.OpenAIError as e:
        print(f"Failed to fetch details: {e}")
        return None



def get_product_description(category: Optional[str] = None, tags: Optional[Set[str]] = None,
                            attributes: Optional[List[Dict[str, str]]] = None, materials: Optional[List[Dict[str, int]]] = None) -> Optional[Dict[str, Any]]:
    prompt = generate_prompt(category, tags, materials, attributes)
    return send_request_to_openai(prompt)

# Example usage:
# response = get_product_description(api_key, category="women-shirts", tags={"summer", "casual"}, attributes=[{"key": "color", "value": "blue"}], materials=[{"key": "cotton", "value": 100}])

def get_openAI_tags(image_urls):
    # Create list containing JSONs of each image URL
    imageList = []
    for image in image_urls:
        imageList.append({"type": "image_url", "image_url": {"url": image}})
    try:
        openai_response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                "role": "system",
                "content": [
                    {
                    "type": "text",
                    "text": "You're a tagging assistant, you will help label and tag product pictures for my online e-commerce platform. Your tasks will be to return which angle the product images were taken from. You will have to choose from 'full-body', 'half-body', 'side', 'back', or 'zoomed' angles. You should label each of the images with one of these labels depending on which you think fits best (ideally, every label should be used at least once, but only if there are 5 or more images), and should respond with an unformatted dictionary where the key is a string representation of the url index of the url and the value is the assigned label."
                    }
                ]
                },
                {
                "role": "user",
                "content": imageList
                },
            ],
            temperature=1,
            max_tokens=500,
            top_p=1,
            frequency_penalty=0,
            presence_penalty=0
        )
        response = json.loads(openai_response.choices[0].message.content)
        return response
    except Exception as e:
        print(f"OpenAI API Error: {e}")
        return {}


@spaces.GPU
def get_face_embeddings(image_urls):
    # Initialize a dictionary to store the face encodings or errors
    results = {}

    # Loop through each image URL
    for index, url in enumerate(image_urls):
        try:
            # Try to download the image from the URL
            response = requests.get(url)
            # Raise an exception if the response is not successful
            response.raise_for_status()

            # Load the image using face_recognition
            image = face_recognition.load_image_file(BytesIO(response.content))

            # Get the face encodings for all faces in the image
            face_encodings = face_recognition.face_encodings(image)

            # If no faces are detected, store an empty list
            if not face_encodings:
                results[str(index)] = []
            else:
                # Otherwise, store the first face encoding as a list
                results[str(index)] = face_encodings[0].tolist()
        except Exception as e:
            # If any error occurs during the download or processing, store the error message
            print(f"Error processing image: {str(e)}")

    return results

# new
ACCURACY_THRESHOLD = 0.86

def open_image_from_url(url):
    # Fetch the image from the URL
    response = requests.get(url, stream=True)
    response.raise_for_status()  # Check if the request was successful

    # Open the image using PIL
    image = Image.open(BytesIO(response.content))

    return image

# Add the main data to the session state
main = [['Product Id', 'Sku', 'Color', 'Images', 'Status', 'Category', 'Text']]

# This is the order of the categories list. NO NOT CHANGE. Just for visualization purposes
cats = ['shirt, blouse', 'top, t-shirt, sweatshirt', 'sweater', 'cardigan', 'jacket', 'vest', 'pants', 'shorts', 'skirt', 'coat', 'dress', 'jumpsuit', 'cape', 'glasses', 'hat', 'headband, head covering, hair accessory', 'tie', 'glove', 'watch', 'belt', 'leg warmer', 'tights, stockings', 'sock', 'shoe', 'bag, wallet', 'scarf', 'umbrella', 'hood', 'collar', 'lapel', 'epaulette', 'sleeve', 'pocket', 'neckline', 'buckle', 'zipper', 'applique', 'bead', 'bow', 'flower', 'fringe', 'ribbon', 'rivet', 'ruffle', 'sequin', 'tassel']

filter = ['dress', 'jumpsuit', 'cape', 'glasses', 'hat', 'headband, head covering, hair accessory', 'tie', 'glove', 'watch', 'belt', 'leg warmer', 'tights, stockings', 'sock', 'shoe', 'scarf', 'umbrella', 'hood', 'collar', 'lapel', 'epaulette', 'sleeve', 'pocket', 'neckline', 'buckle', 'zipper', 'applique', 'bead', 'bow', 'flower', 'fringe', 'ribbon', 'rivet', 'ruffle', 'sequin', 'tassel']

# 0 for full body, 1 for upper body, 2 for lower body, 3 for over body (jacket, coat, etc), 4 for accessories
yolo_mapping = {
    'shirt, blouse': 3,
    'top, t-shirt, sweatshirt' : 1,
    'sweater': 1,
    'cardigan': 1,
    'jacket': 3,
    'vest': 1,
    'pants': 2,
    'shorts': 2,
    'skirt': 2,
    'coat': 3,
    'dress': 0,
    'jumpsuit': 0,
    'bag, wallet': 4
}

# First line full body, second line upper body, third line lower body, fourth line over body, fifth line accessories
label_mapping = [
    ['women-dress-mini', 'women-dress-dress', 'women-dress-maxi', 'women-dress-midi', 'women-playsuitsjumpsuits-playsuit', 'women-playsuitsjumpsuits-jumpsuit', 'women-coords-coords', 'women-swimwear-onepieces', 'women-swimwear-bikinisets'],
    ['women-sweatersknits-cardigan', 'women-top-waistcoat', 'women-top-blouse', 'women-sweatersknits-blouse', 'women-sweatersknits-sweater', 'women-top-top', 'women-loungewear-hoodie', 'women-top-camistanks', 'women-top-tshirt', 'women-top-croptop', 'women-loungewear-sweatshirt', 'women-top-body'],
    ['women-loungewear-joggers', 'women-bottom-trousers', 'women-bottom-leggings', 'women-bottom-jeans', 'women-bottom-shorts', 'women-bottom-skirt', 'women-loungewear-activewear', 'women-bottom-joggers'],
    ['women-top-shirt', 'women-outwear-coatjacket', 'women-outwear-blazer', 'women-outwear-coatjacket', 'women-outwear-kimonos'],
    ['women-accessories-bags']
]

MODEL_NAME = "valentinafeve/yolos-fashionpedia"

feature_extractor = YolosImageProcessor.from_pretrained('hustvl/yolos-small')
model = YolosForObjectDetection.from_pretrained(MODEL_NAME)

def get_category_index(category):
    # Find index of label mapping
    for i, labels in enumerate(label_mapping):
        if category in labels:
            break
    return i

def get_yolo_index(category):
    # Find index of yolo mapping
    return yolo_mapping[category]

def fix_channels(t):
    """
    Some images may have 4 channels (transparent images) or just 1 channel (black and white images), in order to let the images have only 3 channels. I am going to remove the fourth channel in transparent images and stack the single channel in back and white images.
    :param t: Tensor-like image
    :return: Tensor-like image with three channels
    """
    if len(t.shape) == 2:
        return ToPILImage()(torch.stack([t for i in (0, 0, 0)]))
    if t.shape[0] == 4:
        return ToPILImage()(t[:3])
    if t.shape[0] == 1:
        return ToPILImage()(torch.stack([t[0] for i in (0, 0, 0)]))
    return ToPILImage()(t)

def idx_to_text(i):
    return cats[i]

# Random colors used for visualization
COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]

# for output bounding box post-processing
def box_cxcywh_to_xyxy(x):
    x_c, y_c, w, h = x.unbind(1)
    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
         (x_c + 0.5 * w), (y_c + 0.5 * h)]
    return torch.stack(b, dim=1)

def rescale_bboxes(out_bbox, size):
    img_w, img_h = size
    b = box_cxcywh_to_xyxy(out_bbox)
    b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)

    return b

def plot_results(pil_img, prob, boxes):
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    i = 0

    crops = []
    crop_classes = []
    for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors):
        cl = p.argmax()

        # Save each box as an image
        box_img = pil_img.crop((xmin, ymin, xmax, ymax))
        crops.append(box_img)
        crop_classes.append(idx_to_text(cl))

        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))

        ax.text(xmin, ymin, idx_to_text(cl), fontsize=10,
                bbox=dict(facecolor=c, alpha=0.8))

        i += 1

    # Remove white padding all around the image
    plt.axis('off')
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    output_img = plt.gcf()
    plt.close()

    return output_img, crops, crop_classes


def visualize_predictions(image, outputs, threshold=0.8):
    # Keep only predictions with confidence >= threshold
    probas = outputs.logits.softmax(-1)[0, :, :-1]
    keep = probas.max(-1).values > threshold

    # Convert predicted boxes from [0; 1] to image scales
    bboxes_scaled = rescale_bboxes(outputs.pred_boxes[0, keep].cpu(), image.size)

    # Get filtered probabilities and boxes based on the filter list
    filter_set = set(filter)
    filtered_probas_boxes = [
        (proba, box) for proba, box in zip(probas[keep], bboxes_scaled)
        if idx_to_text(proba.argmax()) not in filter_set
    ]

    # If there is a jumpsuit or dress detected, remove them if there are other clothes detected
    contains_jumpsuit_or_dress = any(idx_to_text(proba.argmax()) in ["jumpsuit", "dress"] for proba, _ in filtered_probas_boxes)
    if contains_jumpsuit_or_dress and len(filtered_probas_boxes) > 1:
        filtered_probas_boxes = [
            (proba, box) for proba, box in filtered_probas_boxes
            if idx_to_text(proba.argmax()) not in ["jumpsuit", "dress"]
        ]

    # Remove duplicates: Only keep one box per class
    unique_classes = set()
    unique_filtered_probas_boxes = []
    for proba, box in filtered_probas_boxes:
        class_text = idx_to_text(proba.argmax())
        if class_text not in unique_classes:
            unique_classes.add(class_text)
            unique_filtered_probas_boxes.append((proba, box))

    # If there are remaining filtered probabilities, plot results
    output_img = None
    crops = None
    crop_classes = None
    if unique_filtered_probas_boxes:
        final_probas, final_boxes = zip(*unique_filtered_probas_boxes)
        output_img, crops, crop_classes = plot_results(image, list(final_probas), torch.stack(final_boxes))

    # Return the classes of the detected objects
    return [proba.argmax().item() for proba, _ in unique_filtered_probas_boxes], output_img, crops, crop_classes

@spaces.GPU 
def get_objects(image, threshold=0.8):
    class_counts = {}
    image = fix_channels(ToTensor()(image))
    image = image.resize((600, 800))

    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    detected_classes, output_img, crops, crop_classes = visualize_predictions(image, outputs, threshold=threshold)
    for cl in detected_classes:
        class_name = idx_to_text(cl)
        if class_name not in class_counts:
            class_counts[class_name] = 0
        class_counts[class_name] += 1

    if crop_classes is not None:
        crop_classes = [get_yolo_index(c) for c in crop_classes]

    return class_counts, output_img, crops, crop_classes

def encode_images_to_base64(cropped_list):
    base64_images = []
    for image in cropped_list:
        with io.BytesIO() as buffer:
            image.convert('RGB').save(buffer, format='JPEG')
            base64_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
            base64_images.append(base64_image)
    return base64_images


# def get_cropped_images(images,category):
#     cropped_list = []
#     resultsPerCategory = {}
#     for num, image in enumerate(images):
#         image = open_image_from_url(image)
#         class_counts, output_img, cropped_images, cropped_classes = get_objects(image, 0.37)
#         if not class_counts:
#             continue
        
#         # Get the inverse category as any other mapping label except the current one corresponding category
#         inverse_category = [label for i, labels in enumerate(label_mapping) for label in labels if i != get_category_index(category) and i != 0]
        
#         # If category is a cardigan, we don't recommend category indices 1 and 3
#         if category == 'women-sweatersknits-cardigan':
#             inverse_category = [label for i, labels in enumerate(label_mapping) for label in labels if i != get_category_index(category) and i != 1 and i != 3]

#         for i, image in enumerate(cropped_images):
#             cropped_category = cropped_classes[i]
#             print(cropped_category, cropped_classes[i], get_category_index(category))
            
#             specific_category = label_mapping[cropped_category]

#             if cropped_category == get_category_index(category):
#                 continue

#             cropped_list.append(image)


#     base64_images = encode_images_to_base64(cropped_list)

#     return base64_images




def get_cropped_images(images, category):
    cropped_list = []
    resultsPerCategory = {}
    for num, image in enumerate(images):
        try:
            image = open_image_from_url(image)
            class_counts, output_img, cropped_images, cropped_classes = get_objects(image, 0.37)
            
            if not class_counts:
                continue
            
            for i, image in enumerate(cropped_images):
                cropped_list.append(image)
        except Exception as e:
            print(f"Error processing image {num}: {e}")
            return []
    
    # Convert cropped images to base64 strings
    base64_images = encode_images_to_base64(cropped_list)

    return base64_images





# Define the Gradio interface with the updated components
iface = gr.Interface(
    fn=shot, 
    inputs=[
        gr.Textbox(label="Image URLs (starting with http/https) comma seperated "), 
        gr.Textbox(label="Category"),
        gr.Textbox(label="Level; accepted 'variant' or 'product'")
    ], 
    outputs="text",
     examples=[
        [['https://d2q1sfov6ca7my.cloudfront.net/eyJidWNrZXQiOiAiaGljY3VwLWltYWdlLWhvc3RpbmciLCAia2V5IjogIlc4MDAwMDAwMTM0LU9SL1c4MDAwMDAwMTM0LU9SLTEuanBnIiwgImVkaXRzIjogeyJyZXNpemUiOiB7IndpZHRoIjogODAwLCAiaGVpZ2h0IjogMTIwMC4wLCAiZml0IjogIm91dHNpZGUifX19',
 'https://d2q1sfov6ca7my.cloudfront.net/eyJidWNrZXQiOiAiaGljY3VwLWltYWdlLWhvc3RpbmciLCAia2V5IjogIlc4MDAwMDAwMTM0LU9SL1c4MDAwMDAwMTM0LU9SLTIuanBnIiwgImVkaXRzIjogeyJyZXNpemUiOiB7IndpZHRoIjogODAwLCAiaGVpZ2h0IjogMTIwMC4wLCAiZml0IjogIm91dHNpZGUifX19',
 'https://d2q1sfov6ca7my.cloudfront.net/eyJidWNrZXQiOiAiaGljY3VwLWltYWdlLWhvc3RpbmciLCAia2V5IjogIlc4MDAwMDAwMTM0LU9SL1c4MDAwMDAwMTM0LU9SLTMuanBnIiwgImVkaXRzIjogeyJyZXNpemUiOiB7IndpZHRoIjogODAwLCAiaGVpZ2h0IjogMTIwMC4wLCAiZml0IjogIm91dHNpZGUifX19'], "women-top-shirt","variant"]],
    description="Add an image URL (starting with http/https) or upload a picture, and provide a list of labels separated by commas.",
    title="Full product flow"
)

# Launch the interface
iface.launch()