Spaces:
Sleeping
Sleeping
| """ | |
| app.py — Fashion Recommendation API (CLIP + FAISS + Re-ranking) | |
| """ | |
| import os | |
| import io | |
| import pickle | |
| import tempfile | |
| import requests | |
| import numpy as np | |
| import pandas as pd | |
| from PIL import Image | |
| from contextlib import asynccontextmanager | |
| from fastapi import FastAPI, File, UploadFile | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| import torch | |
| import open_clip | |
| import faiss | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| # ───────────────────────────────────────────── | |
| # ENV CONFIG | |
| # ───────────────────────────────────────────── | |
| HF_OWNER = os.getenv("HF_DATASET_OWNER") | |
| HF_REPO = os.getenv("HF_DATASET_REPO") | |
| CLOUDINARY_CLOUD = os.getenv("CLOUDINARY_CLOUD_NAME") | |
| if not all([HF_OWNER, HF_REPO, CLOUDINARY_CLOUD]): | |
| raise ValueError("❌ Missing env vars: HF_DATASET_OWNER, HF_DATASET_REPO, CLOUDINARY_CLOUD_NAME") | |
| BASE_HF = f"https://huggingface.co/datasets/{HF_OWNER}/{HF_REPO}/resolve/main" | |
| EMBEDDINGS_URL = f"{BASE_HF}/clip_embeddings.pkl" | |
| FILENAMES_URL = f"{BASE_HF}/clip_filenames.pkl" | |
| FAISS_URL = f"{BASE_HF}/faiss_index.bin" | |
| CSV_URL = f"{BASE_HF}/styles.csv" | |
| CLOUDINARY_URL = f"https://res.cloudinary.com/{CLOUDINARY_CLOUD}/image/upload/" | |
| CLIP_MODEL_NAME = "ViT-B-32" | |
| CLIP_PRETRAINED = "openai" | |
| FAISS_NPROBE = 10 | |
| ANN_CANDIDATES = 50 | |
| FINAL_TOP_K = 5 | |
| ml_models = {} | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ───────────────────────────────────────────── | |
| # HELPERS: Download | |
| # ───────────────────────────────────────────── | |
| def download_bytes(url: str) -> bytes: | |
| print(f" ↓ {url}") | |
| r = requests.get(url, stream=True) | |
| r.raise_for_status() | |
| return r.content | |
| def load_pickle(url: str): | |
| return pickle.loads(download_bytes(url)) | |
| def load_csv(url: str) -> pd.DataFrame: | |
| content = download_bytes(url) | |
| return pd.read_csv(io.BytesIO(content), on_bad_lines="skip") | |
| def load_faiss_index(url: str) -> faiss.Index: | |
| data = download_bytes(url) | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as f: | |
| f.write(data) | |
| tmp_path = f.name | |
| index = faiss.read_index(tmp_path) | |
| os.unlink(tmp_path) | |
| return index | |
| # ───────────────────────────────────────────── | |
| # STARTUP / SHUTDOWN | |
| # ───────────────────────────────────────────── | |
| async def lifespan(app: FastAPI): | |
| print(f"🚀 Starting up (device={DEVICE})…") | |
| try: | |
| print("Loading CLIP model…") | |
| clip_model, _, preprocess = open_clip.create_model_and_transforms( | |
| CLIP_MODEL_NAME, pretrained=CLIP_PRETRAINED, device=DEVICE | |
| ) | |
| clip_model.eval() | |
| print("Downloading embeddings, filenames, FAISS index, styles…") | |
| embeddings = np.array(load_pickle(EMBEDDINGS_URL), dtype="float32") | |
| filenames = load_pickle(FILENAMES_URL) | |
| styles_df = load_csv(CSV_URL) | |
| faiss_index = load_faiss_index(FAISS_URL) | |
| if hasattr(faiss_index, "nprobe"): | |
| faiss_index.nprobe = FAISS_NPROBE | |
| ml_models.update({ | |
| "clip_model": clip_model, | |
| "preprocess": preprocess, | |
| "embeddings": embeddings, | |
| "filenames": filenames, | |
| "styles_df": styles_df, | |
| "faiss_index": faiss_index, | |
| }) | |
| print(f"✅ Ready! Index has {faiss_index.ntotal} vectors.") | |
| except Exception as e: | |
| print(f"🛑 Startup error: {e}") | |
| raise | |
| yield | |
| ml_models.clear() | |
| print("🛑 Shutdown.") | |
| # ───────────────────────────────────────────── | |
| # APP | |
| # ───────────────────────────────────────────── | |
| app = FastAPI(title="Klaro — Fashion Recommendation API (CLIP + FAISS)", lifespan=lifespan) | |
| app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, | |
| allow_methods=["*"], allow_headers=["*"]) | |
| # ───────────────────────────────────────────── | |
| # CORE: Feature extraction | |
| # ───────────────────────────────────────────── | |
| def encode_image(img: Image.Image) -> np.ndarray: | |
| preprocess = ml_models["preprocess"] | |
| clip_model = ml_models["clip_model"] | |
| tensor = preprocess(img).unsqueeze(0).to(DEVICE) | |
| features = clip_model.encode_image(tensor) | |
| features = features / features.norm(dim=-1, keepdim=True) | |
| return features.cpu().numpy().astype("float32") # (1, 512) | |
| # ───────────────────────────────────────────── | |
| # CORE: Two-stage retrieval | |
| # ───────────────────────────────────────────── | |
| def faiss_retrieve(query_vec: np.ndarray, k: int): | |
| index = ml_models["faiss_index"] | |
| scores, indices = index.search(query_vec, k) | |
| return scores[0], indices[0] | |
| def rerank(query_vec: np.ndarray, candidate_indices: np.ndarray, top_k: int): | |
| embeddings = ml_models["embeddings"] | |
| candidate_embs = embeddings[candidate_indices] | |
| scores = (candidate_embs @ query_vec.T).flatten() | |
| order = np.argsort(scores)[::-1][:top_k] | |
| return scores[order], candidate_indices[order] | |
| # ───────────────────────────────────────────── | |
| # HELPERS: Metadata ← FIXED (ml_models + new CSV columns) | |
| # ───────────────────────────────────────────── | |
| def get_metadata(filename: str) -> dict: | |
| df = ml_models["styles_df"] # ← was state["styles_df"], now fixed | |
| try: | |
| image_id = int(os.path.splitext(filename)[0]) | |
| row = df[df["Product ID"] == image_id] | |
| if not row.empty: | |
| return { | |
| "gender": str(row["Age Group"].values[0]), | |
| "articleType": str(row["Sub Category"].values[0]), | |
| "baseColour": str(row["Base Colour"].values[0]), | |
| "usage": str(row["Season"].values[0]), | |
| "name": str(row["Product Name"].values[0]), | |
| "brand": str(row["Brand"].values[0]), | |
| "price": float(row["Price"].values[0]), | |
| "discountedPrice": float(row["Discounted Price"].values[0]), | |
| } | |
| except Exception as e: | |
| print(f"Metadata error for {filename}: {e}") | |
| return {} | |
| # ───────────────────────────────────────────── | |
| # ENDPOINT | |
| # ───────────────────────────────────────────── | |
| async def recommend_fashion(file: UploadFile = File(...)): | |
| try: | |
| contents = await file.read() | |
| img = Image.open(io.BytesIO(contents)).convert("RGB") | |
| # Stage 1: CLIP encode + FAISS ANN | |
| query_vec = encode_image(img) # (1, 512) | |
| _, ann_indices = faiss_retrieve(query_vec, k=ANN_CANDIDATES) | |
| ann_indices = ann_indices[ann_indices >= 0] # remove -1s | |
| # Stage 2: Exact cosine re-rank → top-5 | |
| scores, top_indices = rerank(query_vec, ann_indices, top_k=FINAL_TOP_K) | |
| filenames = ml_models["filenames"] | |
| results = [] | |
| for idx, score in zip(top_indices, scores): | |
| filename = os.path.basename(filenames[idx]) | |
| metadata = get_metadata(filename) | |
| results.append({ | |
| "filename": filename, | |
| "image_url": CLOUDINARY_URL + filename, | |
| "metadata": metadata, | |
| "score": float(score), | |
| }) | |
| return {"recommendations": results} | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| return JSONResponse(status_code=500, content={"error": str(e)}) | |
| def health(): | |
| idx = ml_models.get("faiss_index") | |
| return { | |
| "status": "ok", | |
| "index_size": idx.ntotal if idx else 0, | |
| "device": DEVICE, | |
| } | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |