app.py

# app.py
import os
import time
import torch
import numpy as np
import gradio as gr
from PIL import Image, ImageOps
from tqdm.auto import tqdm
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from pinecone import Pinecone, ServerlessSpec
from pinecone_text.sparse import BM25Encoder
from transformers import CLIPProcessor, CLIPModel
import openai

# ------------------- Keys & Setup -------------------
openai.api_key = os.getenv("OPENAI_API_KEY")
pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
spec = ServerlessSpec(cloud=os.getenv("PINECONE_CLOUD") or "aws", region=os.getenv("PINECONE_REGION") or "us-east-1")
index_name = "hybrid-image-search"

if index_name not in pc.list_indexes().names():
    pc.create_index(index_name, dimension=512, metric='dotproduct', spec=spec)
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)
index = pc.Index(index_name)

# ------------------- Models & Dataset -------------------
fashion = load_dataset("ashraq/fashion-product-images-small", split="train")
images = fashion["image"]
metadata = fashion.remove_columns("image").to_pandas()
bm25 = BM25Encoder()
bm25.fit(metadata["productDisplayName"])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SentenceTransformer('sentence-transformers/clip-ViT-B-32', device=device)
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# ------------------- Helper Functions -------------------
def hybrid_scale(dense, sparse, alpha: float):
    if alpha < 0 or alpha > 1:
        raise ValueError("Alpha must be between 0 and 1")
    hsparse = {
        'indices': sparse['indices'],
        'values':  [v * (1 - alpha) for v in sparse['values']]
    }
    hdense = [v * alpha for v in dense]
    return hdense, hsparse

def extract_intent_from_openai(query: str):
    prompt = f"""
You are an assistant for a fashion search engine. Extract the user's intent from the following query.
Return a Python dictionary with keys: category, gender, subcategory, color.
If something is missing, use null.

Query: "{query}"
Only return the dictionary.
"""
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0
        )
        raw = response.choices[0].message['content']
        structured = eval(raw)
        return structured
    except Exception as e:
        print(f"⚠️ OpenAI intent extraction failed: {e}")
        return {"include": {}, "exclude": {}}
#-----------------below changed------------------------------#

import imagehash
from PIL import Image

def is_duplicate(img, existing_hashes, hash_size=16, tolerance=0):
    """
    Checks if the image is a near-duplicate based on perceptual hash.
    :param img: PIL Image
    :param existing_hashes: set of previously seen hashes
    :param hash_size: size of the hash (default=16 for more precision)
    :param tolerance: allowable Hamming distance for near-duplicates
    :return: (bool) whether image is duplicate
    """
    img_hash = imagehash.phash(img, hash_size=hash_size)
    for h in existing_hashes:
        if abs(img_hash - h) <= tolerance:
            return True
    existing_hashes.add(img_hash)
    return False

def extract_metadata_filters(query: str):
    query_lower = query.lower()
    gender = None
    category = None
    subcategory = None
    color = None

    # --- Gender Mapping ---
    gender_map = {
        "men": "Men", "man": "Men", "mens": "Men", "mans": "Men", "male": "Men",
        "women": "Women", "woman": "Women", "womens": "Women", "female": "Women",
        "boys": "Boys", "boy": "Boys",
        "girls": "Girls", "girl": "Girls",
        "kids": "Kids", "kid": "Kids",
        "unisex": "Unisex"
    }
    for term, mapped_value in gender_map.items():
        if term in query_lower:
            gender = mapped_value
            break

    # --- Category Mapping ---
    category_map = {
        "shirt": "Shirts",
        "tshirt": "Tshirts",
        "t-shirt": "Tshirts",
        "jeans": "Jeans",
        "watch": "Watches",
        "kurta": "Kurtas",
        "dress": "Dresses",
        "trousers": "Trousers", "pants": "Trousers",
        "shorts": "Shorts",
        "footwear": "Footwear",
        "shoes": "Shoes",
        "fashion": "Apparel"
    }
    for term, mapped_value in category_map.items():
        if term in query_lower:
            category = mapped_value
            break

    # --- SubCategory Mapping ---
    subCategory_list = [
        "Accessories", "Apparel Set", "Bags", "Bath and Body", "Beauty Accessories",
        "Belts", "Bottomwear", "Cufflinks", "Dress", "Eyes", "Eyewear", "Flip Flops",
        "Fragrance", "Free Gifts", "Gloves", "Hair", "Headwear", "Home Furnishing",
        "Innerwear", "Jewellery", "Lips", "Loungewear and Nightwear", "Makeup",
        "Mufflers", "Nails", "Perfumes", "Sandal", "Saree", "Scarves", "Shoe Accessories",
        "Shoes", "Skin", "Skin Care", "Socks", "Sports Accessories", "Sports Equipment",
        "Stoles", "Ties", "Topwear", "Umbrellas", "Vouchers", "Wallets", "Watches",
        "Water Bottle", "Wristbands"
    ]
    if "topwear" in query_lower or "top" in query_lower:
        subcategory = "Topwear"
    else:
        query_words = query_lower.split()
        for subcat in subCategory_list:
            if subcat.lower() in query_words:
                subcategory = subcat
                break

    # --- Color Extraction ---
    color_list = [
        "red", "blue", "green", "yellow", "black", "white",
        "orange", "pink", "purple", "brown", "grey", "beige"
    ]
    for c in color_list:
        if c in query_lower:
            color = c.capitalize()
            break

    # --- Invalid pairs ---
    invalid_pairs = {
        ("Men", "Dresses"), ("Men", "Sarees"), ("Men", "Skirts"),
        ("Boys", "Dresses"), ("Boys", "Sarees"),
        ("Girls", "Boxers"), ("Men", "Heels")
    }
    if (gender, category) in invalid_pairs:
        print(f"⚠️ Invalid pair: {gender} + {category}, dropping gender")
        gender = None

    # --- Fallback for missing category ---
    if gender and not category:
        category = "Apparel"

    # --- Refine subcategory for party/wedding-related queries ---
    if "party" in query_lower or "wedding" in query_lower or "cocktail" in query_lower:
        if subcategory in ["Loungewear and Nightwear", "Nightdress", "Innerwear"]:
            subcategory = None  # reset it to avoid filtering into wrong items


    return gender, category, subcategory, color

# ------------------- Search Functions -------------------
def search_fashion(query: str, alpha: float, start: int = 0, end: int = 12, gender_override: str = None):
    intent = extract_intent_from_openai(query)

    include = intent.get("include", {})
    exclude = intent.get("exclude", {})

    gender = include.get("gender")
    category = include.get("category")
    subcategory = include.get("subcategory")
    color = include.get("color")

    # Apply override from dropdown
    if gender_override:
        gender = gender_override

    # Build Pinecone filter
    filter = {}

    # Inclusion filters
    if gender:
        filter["gender"] = gender
    if category:
        if category in ["Footwear", "Shoes"]:
            filter["articleType"] = {"$regex": ".*(Shoe|Footwear).*"}
        else:
            filter["articleType"] = category
    if subcategory:
        filter["subCategory"] = subcategory
    
    # Step 4: Exclude irrelevant items for party-like queries
    query_lower = query.lower()
    if any(word in query_lower for word in ["party", "wedding", "cocktail", "traditional", "reception"]):
        filter.setdefault("subCategory", {})
        if isinstance(filter["subCategory"], dict):
            filter["subCategory"]["$nin"] = [
                "Loungewear and Nightwear", "Nightdress", "Innerwear", "Sleepwear", "Vests", "Boxers"
            ]


    if color:
        filter["baseColour"] = color

    # Exclusion filters
    exclude_filter = {}
    if exclude.get("color"):
        exclude_filter["baseColour"] = {"$ne": exclude["color"]}
    if exclude.get("subcategory"):
        exclude_filter["subCategory"] = {"$ne": exclude["subcategory"]}
    if exclude.get("category"):
        exclude_filter["articleType"] = {"$ne": exclude["category"]}

    # Combine all filters
    if filter and exclude_filter:
        final_filter = {"$and": [filter, exclude_filter]}
    elif filter:
        final_filter = filter
    elif exclude_filter:
        final_filter = exclude_filter
    else:
        final_filter = None

    print(f"🔍 Using filter: {final_filter} (showing {start} to {end})")

    # Hybrid encoding
    sparse = bm25.encode_queries(query)
    dense = model.encode(query).tolist()
    hdense, hsparse = hybrid_scale(dense, sparse, alpha=alpha)

    result = index.query(
        top_k=100,
        vector=hdense,
        sparse_vector=hsparse,
        include_metadata=True,
        filter=final_filter
    )

    # Retry fallback
    if len(result["matches"]) == 0:
        print("⚠️ No results, retrying with alpha=0 sparse only")
        hdense, hsparse = hybrid_scale(dense, sparse, alpha=0)
        result = index.query(
            top_k=100,
            vector=hdense,
            sparse_vector=hsparse,
            include_metadata=True,
            filter=final_filter
        )

    # Format results
    imgs_with_captions = []
    seen_hashes = set()

    for r in result["matches"]:
        idx = int(r["id"])
        img = images[idx]
        meta = r.get("metadata", {})
        if not isinstance(img, Image.Image):
            img = Image.fromarray(np.array(img))
        padded = ImageOps.pad(img, (256, 256), color="white")
        caption = str(meta.get("productDisplayName", "Unknown Product"))

        if not is_duplicate(padded, seen_hashes):
            imgs_with_captions.append((padded, caption))

        if len(imgs_with_captions) >= end:
            break

    return imgs_with_captions

def search_by_image(uploaded_image, alpha=0.5, start=0, end=12):
    # Step 1: Preprocess image for CLIP model
    processed = clip_processor(images=uploaded_image, return_tensors="pt").to(device)

    with torch.no_grad():
        image_vec = clip_model.get_image_features(**processed)
        image_vec = image_vec.cpu().numpy().flatten().tolist()

    # Step 2: Query Pinecone index for similar images
    result = index.query(
        top_k=100,  # fetch more to allow deduplication
        vector=image_vec,
        include_metadata=True
    )

    matches = result["matches"]
    imgs_with_captions = []
    seen_hashes = set()

    # Step 3: Deduplicate based on image hash
    for r in matches:
        idx = int(r["id"])
        img = images[idx]
        meta = r.get("metadata", {})
        caption = str(meta.get("productDisplayName", "Unknown Product"))

        if not isinstance(img, Image.Image):
            img = Image.fromarray(np.array(img))

        padded = ImageOps.pad(img, (256, 256), color="white")

        if not is_duplicate(padded, seen_hashes):
            imgs_with_captions.append((padded, caption))

        if len(imgs_with_captions) >= end:
            break

    return imgs_with_captions

import gradio as gr
import whisper

asr_model = whisper.load_model("base")

def handle_voice_search(vf_path, a, offset, gender_ui):
    try:
        transcription = asr_model.transcribe(vf_path)["text"].strip()
    except:
        transcription = ""
    filters = extract_intent_from_openai(transcription) if transcription else {}
    gender_override = gender_ui if gender_ui else filters.get("gender")
    results = search_fashion(transcription, a, 0, 12, gender_override)
    seen_ids = {r[1] for r in results}
    return results, 12, transcription, None, gender_override, results, seen_ids

custom_css = """
/* === Background Styling === */
# html, body {
#     margin: 0;
#     padding: 0;
#     height: 100%;
#     overflow: auto;
# }
html, body {
    height: auto;
    min-height: 100%;
    overflow-x: hidden;
}



# #app-bg {
#     min-height: 100vh;
#     display: flex;
#     justify-content: center;
#     align-items: flex-start;
#     background: radial-gradient(circle at center, #C2C5EF 0%, #E0E2F5 100%) !important;
#     background-attachment: fixed;
#     position: relative;
#     overflow-y: auto;
#     padding: 24px;
# }
#app-bg {
    background: radial-gradient(circle at center, #C2C5EF 0%, #E0E2F5 100%) !important;
    background-attachment: fixed;
    padding: 24px;
    width: 100%;
}




/* === Main Content Container === */
# #main-container {
#     width: 95%;
#     max-width: 1100px;
#     margin: 20px auto;
#     padding: 24px;
#     background: #ffffff;
#     border-radius: 18px;
#     box-shadow: 0 10px 30px rgba(0, 0, 0, 0.08);
#     # border: 2px solid #C2C5EF;
#     border: 2px solid black;
#     position: relative;
#     z-index: 1;
#     overflow: visible;
# }

#main-container {
    width: 95%;
    max-width: 1100px;
    margin: 20px auto;
    padding: 24px;
    background: #ffffff;
    border-radius: 18px;
    box-shadow: 0 10px 30px rgba(0, 0, 0, 0.08);
    # border: 2px solid #C2C5EF;
    border: 2px solid black;
    
}



/* === Card Containers === */
.gr-box, .gr-block, .gr-column, .gr-row, .gr-tab {
    background-color: #C2C5EF !important;
    color: #22284F !important;
    border-radius: 12px;
    padding: 16px !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
}

/* === Headings === */
h1, .gr-markdown h1 {
    font-size: 2.2rem !important;
    font-weight: bold;
    color: #22284F;
    text-align: center;
    margin-bottom: 1rem;
}

/* === Inputs === */
input[type="text"],
.gr-textbox textarea,
.gr-dropdown,
.gr-slider {
    background-color: #C2C5EF !important;
    color: #22284F !important;
    border-radius: 8px;
    border: 1px solid #999 !important;
    padding: 10px !important;
    font-size: 16px;
    box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.05);
}

/* === Gallery Grid === */
.gr-gallery {
    padding-top: 12px;
    overflow-y: auto;
}
.gr-gallery-item {
    width: 128px !important;
    height: 128px !important;
    border-radius: 8px;
    overflow: hidden;
    background-color: #C2C5EF;
    color: #22284F;
    transition: transform 0.3s ease-in-out;
}
.gr-gallery-item:hover {
    transform: scale(1.06);
    box-shadow: 0 3px 12px rgba(0, 0, 0, 0.15);
}
.gr-gallery-item img {
    object-fit: cover;
    width: 100%;
    height: 100%;
    border-radius: 8px;
}

/* === Audio & Image === */
.gr-audio, .gr-image {
    width: 100% !important;
    max-width: 500px !important;
    margin: 0 auto;
    border-radius: 12px;
    background-color: #C2C5EF !important;
    color: #22284F !important;
    box-shadow: 0 1px 5px rgba(0, 0, 0, 0.1);
}

.gr-image {
    height: 256px !important;
}

/* === Buttons === */
.gr-button {
    background-image: linear-gradient(92.88deg, #455EB5 9.16%, #5643CC 43.89%, #673FD7 64.72%);
    color: #ffffff !important;
    border-radius: 8px;
    font-size: 16px;
    font-weight: 500;
    height: 3.5rem;
    padding: 0 1.5rem;
    border: none;
    box-shadow: rgba(80, 63, 205, 0.5) 0 1px 30px;
    transition: all 0.3s;
}
.gr-button:hover {
    transform: translateY(-2px);
    box-shadow: rgba(80, 63, 205, 0.8) 0 2px 20px;
}

/* === Tab Labels === */
button[role="tab"] {
    color: #22284F !important;
    font-weight: 500;
    font-size: 16px;
}
button[role="tab"][aria-selected="true"] {
    color: #f57c00 !important;
    font-weight: bold;
}
button[role="tab"]:hover {
    color: #f57c00 !important;
    font-weight: 600;
    cursor: pointer;
}

/* === Footer === */
.gr-markdown:last-child {
    text-align: center;
    font-size: 14px;
    color: #666;
    padding-top: 1rem;
}

/* === Responsive === */
@media (max-width: 768px) {
    #main-container {
        width: 100%;
        margin: 8px;
        padding: 12px;
    }

    .gr-button {
        font-size: 14px;
        height: 3.2rem;
    }

    input[type="text"], .gr-textbox textarea, .gr-dropdown, .gr-slider {
        font-size: 14px;
        padding: 8px !important;
    }

    h1, .gr-markdown h1 {
        font-size: 1.6rem !important;
    }

    .gr-gallery-item {
        width: 100px !important;
        height: 100px !important;
    }

    .gr-image {
        height: auto !important;
    }
}

"""



with gr.Blocks(css=custom_css) as demo:
  with gr.Column(elem_id="app-bg"):
    with gr.Column(elem_id="main-container"):
      gr.Markdown("# 🛍️ Fashion Product Hybrid Search")

      alpha = gr.Slider(0, 1, value=0.5, label="Hybrid Weight (alpha: 0=sparse, 1=dense)")

      with gr.Tabs():
          with gr.Tab("Text Search"):
              query = gr.Textbox(
                  label="Text Query",
                  placeholder="e.g., floral summer dress for women"
              )
              gender_dropdown = gr.Dropdown(
                  ["", "Men", "Women", "Boys", "Girls", "Kids", "Unisex"],
                  label="Gender Filter (optional)"
              )
              text_search_btn = gr.Button("Search by Text", elem_classes="search-btn")
          with gr.Tab("🎙️ Voice Search"):
            voice_input = gr.Audio(label="Speak Your Query", type="filepath")
            voice_gender_dropdown = gr.Dropdown(["", "Men", "Women", "Boys", "Girls", "Kids", "Unisex"], label="Gender")
            voice_search_btn = gr.Button("Search by Voice")


          with gr.Tab("Image Search"):
              # image_input = gr.Image(
              #     type="pil",
              #     label="Upload an image",
              #     sources=["upload", "clipboard"],
              #     height=256,
              #     width=356
              # )
              image_input = gr.Image(
                type="pil",
                label="Upload an image",
                sources=["upload", "clipboard"],
                # tool=None,
                height=400
              )

              image_gender_dropdown = gr.Dropdown(
                  ["", "Men", "Women", "Boys", "Girls", "Kids", "Unisex"],
                  label="Gender Filter (optional)"
              )
              image_search_btn = gr.Button("Search by Image", elem_classes="search-btn")

      gallery = gr.Gallery(label="Search Results", columns=6, height=None)
      load_more_btn = gr.Button("Load More")

    # --- UI State Holders ---
      search_offset = gr.State(0)
      current_query = gr.State("")
      current_image = gr.State(None)
      current_gender = gr.State("")
      shown_results = gr.State([])
      shown_ids = gr.State(set())

    # --- Unified Search Function ---
    def unified_search(q, uploaded_image, a, offset, gender_ui):
        start = 0
        end = 12

        filters = extract_intent_from_openai(q) if q.strip() else {}
        gender_override = gender_ui if gender_ui else filters.get("gender")

        if uploaded_image is not None:
            results = search_by_image(uploaded_image, a, start, end)
        elif q.strip():
            results = search_fashion(q, a, start, end, gender_override)
        else:
            results = []

        seen_ids = {r[1] for r in results}
        return results, end, q, uploaded_image, gender_override, results, seen_ids

    # Text Search
    # Text Search
    text_search_btn.click(
      unified_search,
      inputs=[query, gr.State(None), alpha, search_offset, gender_dropdown],
      outputs=[gallery, search_offset, current_query, current_image, current_gender, shown_results, shown_ids]
    )

    voice_search_btn.click(
    handle_voice_search,
    inputs=[voice_input, alpha, search_offset, voice_gender_dropdown],
    outputs=[gallery, search_offset, current_query, current_image, current_gender, shown_results, shown_ids]
)

    # Image Search
    image_search_btn.click(
        unified_search,
        inputs=[gr.State(""), image_input, alpha, search_offset, image_gender_dropdown],
        outputs=[gallery, search_offset, current_query, current_image, current_gender, shown_results, shown_ids]
    )

    # --- Load More Button ---
    def load_more_fn(a, offset, q, img, gender_ui, prev_results, prev_ids):
        start = offset
        end = offset + 12
        gender_override = gender_ui

        if img is not None:
            new_results = search_by_image(img, a, start, end)
        elif q.strip():
            new_results = search_fashion(q, a, start, end, gender_override)
        else:
            new_results = []

        filtered_new = []
        new_ids = set()
        for item in new_results:
            img_obj, caption = item
            if caption not in prev_ids:
                filtered_new.append(item)
                new_ids.add(caption)

        combined = prev_results + filtered_new
        updated_ids = prev_ids.union(new_ids)

        return combined, end, combined, updated_ids

    load_more_btn.click(
        load_more_fn,
        inputs=[alpha, search_offset, current_query, current_image, current_gender, shown_results, shown_ids],
        outputs=[gallery, search_offset, shown_results, shown_ids]
    )

    # gr.Markdown("🧠 Powered by OpenAI + Hybrid AI Fashion Search")

demo.launch()