fix

Dockerfile

@@ -7,6 +7,7 @@ RUN apt-get update && apt-get install -y \
     curl \
     software-properties-common \
     git \
+    libgl1-mesa-glx \
     && rm -rf /var/lib/apt/lists/*
 
 COPY requirements.txt ./

src/model.py

@@ -1,16 +1,14 @@
 import faiss
 import torch
 import clip
-from openai import OpenAI
 import numpy as np
 from PIL import Image
 from fastapi import FastAPI
 from typing import List
 import segmentation
 
-client = OpenAI()
 device = "cpu"
-model, preprocess = clip.load("ViT-B/32", device=device)
+model, preprocess = clip.load("ViT-B/32", device=device, download_root="./clip_cache")
 
 def get_image_features(image: Image.Image) -> np.ndarray:
     """Extract CLIP features from an image."""
@@ -29,9 +27,9 @@ def save_image_in_index(image_features: np.ndarray, metadata: dict):
     index.add(image_features)
     meta_data_store.append(metadata)
 
-def process_image_embedding(image_url: str, labels=['clothes']) -> np.ndarray:
+def process_image_embedding(image_or_url, labels=['clothes']) -> np.ndarray:
     """Get feature embedding for a query image."""
-    search_image, search_detections = segmentation.grounded_segmentation(image=image_url, labels=labels)
+    search_image, search_detections = segmentation.grounded_segmentation(image=image_or_url, labels=labels)
     cropped_image = segmentation.cut_image(search_image, search_detections[0].mask, search_detections[0].box)
 
     # Convert to valid RGB
@@ -40,8 +38,7 @@ def process_image_embedding(image_url: str, labels=['clothes']) -> np.ndarray:
     if cropped_image.ndim == 2:
         cropped_image = np.stack([cropped_image] * 3, axis=-1)
 
-    pil_image = Image.fromarray(cropped_image)
-    return pil_image
+    return Image.fromarray(cropped_image)
 
 def get_top_k_results(image_url: str, k: int = 10) -> List[dict]:
     """Find top-k similar images from the index."""
@@ -63,3 +60,5 @@ def get_top_k_results(image_url: str, k: int = 10) -> List[dict]:
 
 
 
+
+