app.py

"""
Visual Search API - HuggingFace Space
"""

import os
import gradio as gr
import torch
import numpy as np
from PIL import Image

# Pinecone config from HF Secrets
PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
PINECONE_HOST = os.environ.get('PINECONE_HOST')

# Model (loaded on first use)
model = None


def load_model():
    """Load Jina CLIP v2 model."""
    global model
    if model is None:
        print("Loading Jina CLIP v2...")
        from transformers import AutoModel
        model = AutoModel.from_pretrained(
            "jinaai/jina-clip-v2",
            trust_remote_code=True
        )
        model.eval()
        print("Model loaded!")
    return model


def get_embedding(image: Image.Image) -> list:
    """Generate 512-dim embedding for an image."""
    m = load_model()

    with torch.no_grad():
        emb = m.encode_image(image)
        if hasattr(emb, 'cpu'):
            emb = emb.cpu().numpy()
        emb = emb.flatten()
        emb = emb / np.linalg.norm(emb)
        if len(emb) > 512:
            emb = emb[:512]
        return emb.tolist()


def query_pinecone(embedding: list, top_k: int = 12) -> list:
    """Query Pinecone for similar products."""
    if not PINECONE_API_KEY or not PINECONE_HOST:
        return []

    import requests

    resp = requests.post(
        f"https://{PINECONE_HOST}/query",
        headers={
            "Api-Key": PINECONE_API_KEY,
            "Content-Type": "application/json"
        },
        json={
            "vector": embedding,
            "topK": top_k,
            "includeMetadata": True
        },
        timeout=15
    )

    if resp.status_code != 200:
        return []

    matches = resp.json().get('matches', [])
    return [
        {
            'handle': m.get('metadata', {}).get('handle', m.get('id')),
            'title': m.get('metadata', {}).get('title', ''),
            'score': m.get('score', 0),
        }
        for m in matches
    ]


def search(image):
    """Main search function."""
    if image is None:
        return "No image provided"

    try:
        embedding = get_embedding(image)
        products = query_pinecone(embedding)

        if not products:
            return "No similar products found"

        result = "\n".join([
            f"{i+1}. {p['title']} ({p['handle']}) - score: {p['score']:.3f}"
            for i, p in enumerate(products)
        ])
        return result
    except Exception as e:
        return f"Error: {str(e)}"


# Simple Gradio interface
demo = gr.Interface(
    fn=search,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Textbox(label="Similar Products", lines=15),
    title="Visual Product Search",
    description="Upload an image to find similar products."
)

if __name__ == "__main__":
    demo.launch()