app.py

import os

os.environ["HUGGINGFACE_DEMO"] = "1"  # set before import from app

from dotenv import load_dotenv

load_dotenv()
################################################################################################

import gradio as gr
import uuid
import shutil

from app.config import get_settings
from app.schemas.requests import Attribute
from app.request_handler import handle_extract
from app.services.factory import AIServiceFactory


settings = get_settings()
IMAGE_MAX_SIZE = 1536


async def forward_request(
    attributes, product_taxonomy, product_data, ai_model, pil_images
):
    # prepare temp folder
    request_id = str(uuid.uuid4())
    request_temp_folder = os.path.join("gradio_temp", request_id)
    os.makedirs(request_temp_folder, exist_ok=True)

    try:
        # convert attributes to schema
        attributes = "attributes_object = {" + attributes + "}"
        try:
            attributes = exec(attributes, globals())
        except:
            raise gr.Error(
                "Invalid `Attribute Schema`. Please insert valid schema following the example."
            )
        for key, value in attributes_object.items():  # type: ignore
            attributes_object[key] = Attribute(**value)  # type: ignore

        if product_data == "":
            product_data = "{}"
        product_data_code = f"product_data_object = {product_data}"

        try:
            exec(product_data_code, globals())
        except:
            raise gr.Error(
                "Invalid `Product Data`. Please insert valid dictionary or leave it empty."
            )

        if pil_images is None:
            raise gr.Error("Please upload image(s) of the product")
        pil_images = [pil_image[0] for pil_image in pil_images]
        img_paths = []
        for i, pil_image in enumerate(pil_images):
            if max(pil_image.size) > IMAGE_MAX_SIZE:
                ratio = IMAGE_MAX_SIZE / max(pil_image.size)
                pil_image = pil_image.resize(
                    (int(pil_image.width * ratio), int(pil_image.height * ratio))
                )
            img_path = os.path.join(request_temp_folder, f"{i}.jpg")
            if pil_image.mode in ("RGBA", "LA") or (
                pil_image.mode == "P" and "transparency" in pil_image.info
            ):
                pil_image = pil_image.convert("RGBA")
                if pil_image.getchannel("A").getextrema() == (
                    255,
                    255,
                ):  # if fully opaque, save as JPEG
                    pil_image = pil_image.convert("RGB")
                    image_format = "JPEG"
                else:
                    image_format = "PNG"
            else:
                image_format = "JPEG"
            pil_image.save(img_path, image_format, quality=100, subsampling=0)
            img_paths.append(img_path)

        # mapping
        if ai_model in settings.OPENAI_MODELS:
            ai_vendor = "openai"
        elif ai_model in settings.ANTHROPIC_MODELS:
            ai_vendor = "anthropic"
        service = AIServiceFactory.get_service(ai_vendor)

        try:
            json_attributes, reevaluated = await service.extract_attributes_with_validation(
                attributes_object,  # type: ignore
                ai_model,
                None,
                product_taxonomy,
                product_data_object,  # type: ignore
                img_paths=img_paths,
            )
        except Exception as e:
            print(e)
            raise gr.Error("Failed to extract attributes. Something went wrong.")
    finally:
        # remove temp folder anyway
        shutil.rmtree(request_temp_folder)

    gr.Info("Process completed!")
    return json_attributes, reevaluated


def add_attribute_schema(attributes, attr_name, attr_desc, attr_type, allowed_values):
    schema = f"""
"{attr_name}": {{
    "description": "{attr_desc}",
    "data_type": "{attr_type}",
    "allowed_values": [
        {', '.join([f'"{v.strip()}"' for v in allowed_values.split(',')]) if allowed_values != "" else ""}
    ]
}},
"""
    return attributes + schema, "", "", "", ""

import_for_schema = """
from enum import Enum
from pydantic import BaseModel, Field
from typing import List
"""

sample_schema = """from pydantic import BaseModel, Field


class Length(BaseModel):
    maxi: int = Field(..., description="Maxi: Dress extends to the ankles or floor.")
    knee_length: int = Field(..., description="Knee Length: Dress ends around the knees.")
    mini: int = Field(..., description="Mini: Short dress that ends well above the knees.")
    midi: int = Field(..., description="Midi: Dress falls between the knee and ankle.")


class Style(BaseModel):
    a_line: int = Field(..., description="A Line: Fitted at the top and gradually flares toward the hem, forming an 'A' shape.")
    bodycon: int = Field(..., description="Bodycon: Tight-fitting and figure-hugging, usually made with stretchy fabric.")
    shirt_dress: int = Field(..., description="Shirt Dress: Structured like a shirt with buttons, collar, and sleeves; may include a belt.")
    wrap_dress: int = Field(..., description="Wrap Dress: Features a front closure that wraps and ties at the side or back.")
    slip: int = Field(..., description="Slip: Lightweight, spaghetti-strap dress with minimal structure, often bias-cut.")
    smock: int = Field(..., description="Smock: Loose-fitting with gathered or shirred sections, usually on bodice or neckline.")
    corset: int = Field(..., description="Corset: Structured bodice with boning or lacing that shapes the waist.")
    jumper_dress: int = Field(..., description="Jumper Dress: Layered dress style similar to a pinafore, often more casual or thick-strapped.")
    shift: int = Field(..., description="Shift: Simple, straight dress with no defined waist, typically above the knee.")


class SleeveLength(BaseModel):
    sleeveless: int = Field(..., description="Sleeveless: No sleeves.")
    three_quarters_sleeve: int = Field(..., description="Three quarters Sleeve: Sleeves that end between the elbow and wrist.")
    long_sleeve: int = Field(..., description="Long Sleeve: Sleeves that extend to the wrist.")
    short_sleeve: int = Field(..., description="Short Sleeve: Sleeves that end above the elbow.")
    strapless: int = Field(..., description="Strapless: No shoulder straps or sleeves.")


class Neckline(BaseModel):
    v_neck: int = Field(..., description="V Neck: Neckline dips down in the shape of a 'V', varying from shallow to deep.")
    sweetheart: int = Field(..., description="Sweetheart: A heart-shaped neckline, often curving over the bust and dipping in the center.")
    round_neck: int = Field(..., description="Round Neck: Circular neckline sitting around the base of the neck.")
    square_neck: int = Field(..., description="Square Neck: Straight horizontal cut across the chest with vertical sides, forming a square.")
    high_neck: int = Field(..., description="High Neck: Extends up the neck slightly but not folded like a turtle neck.")
    crew_neck: int = Field(..., description="Crew Neck: High, rounded neckline that sits close to the neck.")
    turtle_neck: int = Field(..., description="Turtle Neck: High neckline that folds over and covers the neck completely.")
    off_the_shoulder: int = Field(..., description="Off the Shoulder: Sits below the shoulders, exposing the shoulders and collarbone.")


class Pattern(BaseModel):
    floral: int = Field(..., description="Floral pattern")
    stripe: int = Field(..., description="Stripe pattern")
    leopard_print: int = Field(..., description="Leopard print")
    plain: int = Field(..., description="Plain")
    geometric: int = Field(..., description="Geometric pattern")
    logo: int = Field(..., description="Logo print")
    other: int = Field(..., description="Other pattern")


class Fabric(BaseModel):
    cotton: int = Field(..., description="Cotton")
    denim: int = Field(..., description="Denim")
    linen: int = Field(..., description="Linen")
    satin: int = Field(..., description="Satin")
    silk: int = Field(..., description="Silk")
    leather: int = Field(..., description="Leather")
    velvet: int = Field(..., description="Velvet")
    polyester: int = Field(..., description="Polyester")
    viscose: int = Field(..., description="Viscose")


class Features(BaseModel):
    pockets: int = Field(..., description="Has pockets")
    lined: int = Field(..., description="Lined")
    cut_out: int = Field(..., description="Cut out design")
    backless: int = Field(..., description="Backless")
    none: int = Field(..., description="No special features")


class Closure(BaseModel):
    button: int = Field(..., description="Button closure")
    zip: int = Field(..., description="Zip closure")
    press_stud: int = Field(..., description="Press stud closure")
    clasp: int = Field(..., description="Clasp closure")


class BodyFit(BaseModel):
    petite: int = Field(..., description="Petite fit")
    maternity: int = Field(..., description="Maternity fit")
    regular: int = Field(..., description="Regular fit")
    tall: int = Field(..., description="Tall fit")
    plus_size: int = Field(..., description="Plus size fit")


class Occasion(BaseModel):
    beach: int = Field(..., description="Suitable for beach")
    casual: int = Field(..., description="Casual wear")
    cocktail: int = Field(..., description="Cocktail event")
    day: int = Field(..., description="Day wear")
    evening: int = Field(..., description="Evening wear")
    mother_of_the_bride: int = Field(..., description="Mother of the bride dress")
    party: int = Field(..., description="Party wear")
    prom: int = Field(..., description="Prom dress")


class Season(BaseModel):
    spring: int = Field(..., description="Spring season")
    summer: int = Field(..., description="Summer season")
    autumn: int = Field(..., description="Autumn season")
    winter: int = Field(..., description="Winter season")


class Product(BaseModel):
    length: Length = Field(..., description="Single value ,Length of the dress")
    style: Style = Field(..., description="Can have multiple values, Style of the dress")
    sleeve_length: SleeveLength = Field(..., description="Single value ,Sleeve length of the dress")
    neckline: Neckline = Field(..., description="Single value ,Neckline of the dress")
    pattern: Pattern = Field(..., description="Can have multiple values, Pattern of the dress")
    fabric: Fabric = Field(..., description="Can have multiple values, Fabric of the dress")
    features: Features = Field(..., description="Can have multiple values, Features of the dress")
    closure: Closure = Field(..., description="Can have multiple values ,Closure of the dress")
    body_fit: BodyFit = Field(..., description="Single value ,Body fit of the dress")
    occasion: Occasion = Field(..., description="Can have multiple values ,Occasion of the dress")
    season: Season = Field(..., description="Single value ,Season of the dress")
"""

cf_style_schema = """
"Length": {
    "description": "Length of dress",
    "data_type": "string",
    "allowed_values": [
        "Maxi",
        "Knee Length",
        "Mini",
        "Midi"
    ]
},
"Style": {
    "description": "Select the most appropriate dress style based on the garment's silhouette, fit, structural features, and overall design. Focus on how the dress is constructed and worn: whether it is fitted or loose, whether it has defining elements such as shirring, boning, buttons, collars, tiers, layering, or wrap ties. Ignore color, pattern, or fabric unless they directly influence the structure (e.g., stretch fabric for Bodycon). Use the visual cues of the neckline, sleeves, waistline, hemline, and closure type to guide your choice. Only select one style that best captures the dominant structural or design identity of the dress. Refer to the following definitions when uncertain: - 'A Line': Fitted at the top and gradually flares toward the hem, forming an 'A' shape.- 'Bodycon': Tight-fitting and figure-hugging, usually made with stretchy fabric.- 'Column': Straight silhouette from top to bottom, with minimal shaping or flare.- 'Shirt Dress': Structured like a shirt with buttons, collar, and sleeves; may include a belt.- 'Wrap Dress': Features a front closure that wraps and ties at the side or back.- 'Slip': Lightweight, spaghetti-strap dress with minimal structure, often bias-cut.- 'Kaftan': Very loose, flowing garment with wide sleeves and minimal shaping.- 'Smock': Loose-fitting with gathered or shirred sections (usually bodice or neckline).- 'Corset': Structured bodice with boning or lacing that shapes the waist.- 'Pinafore': Sleeveless over-dress, often worn layered over another top.- 'Jumper Dress': Layered dress style similar to a pinafore, often more casual or thick-strapped.- 'Blazer Dress': Tailored like a blazer or suit jacket, often double-breasted or lapelled.- 'Tunic': Loose and straight-cut, often worn short or over pants/leggings.- 'Gown': Full-length, formal dress with a structured or dramatic silhouette.- 'Asymmetric': Dress with a non-symmetrical hem, neckline, or sleeve design.- 'Shift': Simple, straight dress with no defined waist, typically above the knee.- 'Drop waist': Waistline sits low on the hips, usually with a loose top and flared skirt.- 'Empire': High waistline just below the bust, flowing skirt from there downward.- 'Modest': Covers most of the body, with high neckline, long sleeves, and longer hemline. Use structural cues over stylistic interpretation. Do not infer intent (e.g., party, formal) unless it’s directly tied to the construction.",
    "data_type": "string",
    "allowed_values": [
        "A Line",
        "Bodycon",
        "Column",
        "Shirt Dress",
        "Wrap Dress",
        "Smock",
        "Corset",
        "Tunic",
        "Asymmetric",
        "Shift",
        "Modest"
    ]
},
"Sleeve_length": {
    "description": "Length of sleeves on dress",
    "data_type": "string",
    "allowed_values": [
        "Sleeveless",
        "Three quarters Sleeve",
        "Long Sleeve",
        "Short Sleeve",
        "Strapless"
    ]
},
"Neckline": {
    "description": "Identify the neckline style based on the visible shape and structure of the neckline area. Focus on the cut and contour around the collarbone, shoulders, and upper chest. Only choose the neckline that best represents the dominant design — ignore collars, patterns, or styling details unless they significantly alter the neckline shape. Use the following definitions for clarity: - 'V Neck': Neckline dips down in the shape of a 'V', varying from shallow to deep. - 'Sweetheart': A heart-shaped neckline, often curving over the bust and dipping in the center. - 'Round Neck': Circular neckline sitting around the base of the neck, not as high as a crew neck. - 'Halter Neck': Straps go around the neck, leaving shoulders and upper back exposed. - 'Square Neck': Straight horizontal cut across the chest with vertical sides, forming a square. - 'High Neck': Extends up the neck slightly but not folded like a turtle neck. - 'Crew Neck': High, rounded neckline that sits close to the neck (commonly found in T-shirts). - 'Cowl Neck': Draped or folded neckline that hangs in soft folds. - 'Turtle Neck': High neckline that folds over and covers the neck completely. - 'Off the Shoulder': Sits below the shoulders, exposing the shoulders and collarbone. - 'One Shoulder': Covers one shoulder only, leaving the other bare. - 'Bandeau': Straight, strapless neckline that wraps across the bust. - 'Boat Neck': Wide, shallow neckline that runs almost horizontally from shoulder to shoulder. - 'Scoop Neck': U-shaped neckline, typically deeper than a round neck. - Always prioritize structure over styling — for example, a dress with embellishment or a mesh overlay still counts as 'V Neck' if the main shape is a V. If a neckline is borderline between two types, choose the simpler or more dominant structure.",
    "data_type": "string",
    "allowed_values": [
        "V Neck",
        "Sweetheart",
        "Round Neck",
        "Halter Neck",
        "Square Neck",
        "Cowl Neck",
        "Turtle Neck",
        "Off the shoulder",
        "Boat Neck",
        "Scoop Neck"
    ]
},
"pattern": {
    "description": "Pattern of the garment",
    "data_type": "string",
    "allowed_values": [
        "Floral",
        "Stripe",
        "Leopard Print",
        "Spot",
        "Plain",
        "Geometric",
        "Logo",
        "Graphic print",
        "Check",
        "other"
    ]
},
"fabric": {
    "description": "Material of the garment",
    "data_type": "string",
    "allowed_values": [
        "Cotton",
        "Denim",
        "Jersey",
        "Linen",
        "Satin",
        "Silk",
        "Leather",
        "Velvet",
        "Corduroy",
        "Ponte",
        "Knit",
        "Lace",
        "Polyester",
        "Viscose"
    ]
},
"features": {
    "description": "special features of the garment",
    "data_type": "list[string]",
    "allowed_values": [
        "Pockets",
        "Lined",
        "Cut Out",
        "Backless",
        "none"
    ]
},
"Closure": {
    "description": "Closure of the garment. How it is closed",
    "data_type": "list[string]",
    "allowed_values": [
        "Button",
        "Zip",
        "Press Stud",
        "Clasp"
    ]
},
"Body_Fit": {
    "description": "How the dress fits the body",
    "data_type": "string",
    "allowed_values": [
        "Petite",
        "Maternity",
        "Regular",
        "Tall",
        "Plus Size"
    ]
},
"Occasion": {
    "description": "What occasions do the dress match",
    "data_type": "list[string]",
    "allowed_values": [
        "Beach",
        "Casual",
        "Cocktail",
        "Day",
        "Bridal",
        "Bridesmaid",
        "Evening",
        "Mother of the Bride",
        "Party",
        "Prom"
    ]
},
"Season": {
    "description": "What season do the dress match",
    "data_type": "list[string]",
    "allowed_values": [
        "Spring",
        "Summer",
        "Autumn",
        "Winter"
    ]
}
"""[1:]

description = """
This is a simple demo for Attribution. Follow the steps below:

1. Upload image(s) of a product.
2. Enter the product taxonomy (e.g. 'upper garment', 'lower garment', 'bag'). If only one product is in the image, you can leave this field empty.
3. Select the AI model to use.
4. Enter known attributes (optional).
5. Enter the attribute schema or use the "Add Attributes" section to add attributes.
6. Click "Extract Attributes" to get the extracted attributes.
"""

product_data_placeholder = """Example:
{
    "brand": "Leaf",
    "size": "M",
    "product_name": "Leaf T-shirt",
    "color": "red"
}
"""
product_data_value = """
{
    "data1": "",
    "data2": ""
}
"""

with gr.Blocks(title="Internal Demo for Attribution") as demo:
    with gr.Row():
        with gr.Column(scale=12):
            gr.Markdown(
                """<div style="text-align: center; font-size: 24px;"><strong>Internal Demo for Attribution</strong></div>"""
            )
            gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=12):
            with gr.Row():
                with gr.Column():
                    gallery = gr.Gallery(
                        label="Upload images of your product here", type="pil"
                    )
                    product_taxnomy = gr.Textbox(
                        label="Product Taxonomy",
                        placeholder="Enter product taxonomy here (e.g. 'upper garment', 'lower garment', 'bag')",
                        lines=1,
                        max_lines=1,
                    )
                    ai_model = gr.Dropdown(
                        label="AI Model",
                        choices=settings.SUPPORTED_MODELS,
                        interactive=True,
                    )
                    product_data = gr.TextArea(
                        label="Product Data (Optional)",
                        placeholder=product_data_placeholder,
                        value=product_data_value.strip(),
                        interactive=True,
                        lines=10,
                        max_lines=10,
                    )

                    # track_count = gr.State(1)
                    # @gr.render(inputs=track_count)
                    # def render_tracks(count):
                    #     ka_names = []
                    #     ka_values = []
                    #     with gr.Column():
                    #         for i in range(count):
                    #             with gr.Column(variant="panel"):
                    #                 with gr.Row():
                    #                     ka_name = gr.Textbox(placeholder="key", key=f"key-{i}", show_label=False)
                    #                     ka_value = gr.Textbox(placeholder="data", key=f"data-{i}", show_label=False)
                    #                     ka_names.append(ka_name)
                    #                     ka_values.append(ka_value)

                    # add_track_btn = gr.Button("Add Product Data")
                    # remove_track_btn = gr.Button("Remove Product Data")
                    # add_track_btn.click(lambda count: count + 1, track_count, track_count)
                    # remove_track_btn.click(lambda count: count - 1, track_count, track_count)

                with gr.Column():
                    attributes = gr.TextArea(
                        label="Attribute Schema",
                        value=cf_style_schema,
                        placeholder="Enter schema here or use Add Attributes below",
                        interactive=True,
                        lines=30,
                        max_lines=30,
                    )

                    # with gr.Accordion("Add Attributes", open=False):
                    #     attr_name = gr.Textbox(
                    #         label="Attribute name", placeholder="Enter attribute name"
                    #     )
                    #     attr_desc = gr.Textbox(
                    #         label="Description", placeholder="Enter description"
                    #     )
                    #     attr_type = gr.Dropdown(
                    #         label="Type",
                    #         choices=[
                    #             "string",
                    #             "list[string]",
                    #             "int",
                    #             "list[int]",
                    #             "float",
                    #             "list[float]",
                    #             "bool",
                    #             "list[bool]",
                    #         ],
                    #         interactive=True,
                    #     )
                    #     allowed_values = gr.Textbox(
                    #         label="Allowed values (separated by comma)",
                    #         placeholder="yellow, red, blue",
                    #     )
                    #     add_btn = gr.Button("Add Attribute")

            with gr.Row():
                submit_btn = gr.Button("Extract Attributes")

        with gr.Column(scale=6):
            output_json = gr.Json(
                label="Extracted Attributes", value={}, show_indices=False
            )
            reevaluated_output_json = gr.Json(
                label="Extracted Attributes", value={}, show_indices=False
            )

    # add_btn.click(
    #     add_attribute_schema,
    #     inputs=[attributes, attr_name, attr_desc, attr_type, allowed_values],
    #     outputs=[attributes, attr_name, attr_desc, attr_type, allowed_values],
    # )

    submit_btn.click(
        forward_request,
        inputs=[attributes, product_taxnomy, product_data, ai_model, gallery],
        outputs=[output_json, reevaluated_output_json],
    )


attr_user = os.getenv("ATTR_USER", "1")
attr_pass = os.getenv("ATTR_PASS", "a")
auth = (attr_user, attr_pass)
demo.launch(auth=auth, debug=True, ssr_mode=False)