Upload 5 files

Dockerfile

@@ -0,0 +1,26 @@
+# 1. Base Image
+FROM python:3.11-slim
+
+# 2. Install Git LFS to handle the 10GB download
+RUN apt-get update && apt-get install -y git git-lfs && git lfs install
+
+# 3. Setup Hugging Face User
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+WORKDIR /app
+
+# 4. Install Dependencies
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# 5. DOWNLOAD YOUR DATA (The 10GB Part)
+# We use the HF_TOKEN secret to clone your private dataset into the /app folder.
+RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true \
+    git clone https://user:$(cat /run/secrets/hf_token)@huggingface.co/datasets/aniketkumar1106/orbit-data .
+
+# 6. EXPOSE PORT & RUN
+# Hugging Face Spaces mandates port 7860. 
+# We override the port here so you don't have to change your server.py.
+EXPOSE 7860
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860"]

orbiitt_engine.py

@@ -0,0 +1,123 @@
+import os
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import AutoModel, AutoProcessor
+import chromadb
+from tqdm import tqdm
+
+class OrbiittEngine:
+    def __init__(self, db_path="./orbiitt_db"):
+        self.image_folder = "Productimages"
+        self.db_path = db_path
+        
+        # 1. Device detection (Native Apple Silicon Support)
+        self.device = "mps" if torch.backends.mps.is_available() else "cpu"
+        
+        # 2. Load SigLIP 2 (The Modern Champ)
+        print(f"🧠 Loading SigLIP 2 (google/siglip2-base-patch16-256) on {self.device}...")
+        self.model_name = "google/siglip2-base-patch16-256"
+        self.model = AutoModel.from_pretrained(self.model_name).to(self.device).eval()
+        self.processor = AutoProcessor.from_pretrained(self.model_name)
+        
+        # 3. Get Expected Dimension (768 for Base)
+        self.expected_dim = self.model.config.vision_config.hidden_size
+        
+        # 4. Connect to Database with Safety Logic
+        self.client = chromadb.PersistentClient(path=self.db_path)
+        self._check_db_compatibility()
+        
+        # CRITICAL: Set space to 'cosine' for AI search
+        self.collection = self.client.get_or_create_collection(
+            name="product_catalog",
+            metadata={"hnsw:space": "cosine"}
+        )
+
+    def _check_db_compatibility(self):
+        """Ensures the stored vectors match SigLIP 2's 768 dimensions."""
+        try:
+            col = self.client.get_collection(name="product_catalog")
+            sample = col.peek(limit=1)
+            if sample and sample['embeddings']:
+                existing_dim = len(sample['embeddings'][0])
+                if existing_dim != self.expected_dim:
+                    print(f"⚠️  Dimension Mismatch: DB is {existing_dim}, Model is {self.expected_dim}")
+                    if input("Wipe DB and restart? (y/n): ").lower() == 'y':
+                        self.client.delete_collection(name="product_catalog")
+                    else:
+                        exit()
+        except: pass
+
+    def get_image_embedding(self, image_path):
+        """Processes image and returns a normalized 768D vector."""
+        image = Image.open(image_path).convert("RGB")
+        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            features = self.model.get_image_features(**inputs)
+            # Normalize to unit length (Unit Sphere)
+            features = F.normalize(features, p=2, dim=-1)
+        return features.squeeze().cpu().numpy().tolist()
+
+    def get_text_embedding(self, text):
+        """Processes text and returns a normalized 768D vector."""
+        # Use the SigLIP 2 standard prompt template
+        prompt = f"this is a photo of {text}"
+        inputs = self.processor(text=[prompt], padding="max_length", return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            features = self.model.get_text_features(**inputs)
+            features = F.normalize(features, p=2, dim=-1)
+        return features.squeeze().cpu().numpy().tolist()
+
+    def index_images(self):
+        """Scans the Productimages folder and indexes them."""
+        if not os.path.exists(self.image_folder):
+            print(f"❌ Error: {self.image_folder} not found."); return
+
+        files = [f for f in os.listdir(self.image_folder) if f.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))]
+        print(f"🏗️ Indexing {len(files)} products...")
+
+        for fname in tqdm(files, desc="SigLIP 2 Processing"):
+            path = os.path.join(self.image_folder, fname)
+            if len(self.collection.get(ids=[fname])['ids']) > 0: continue
+            try:
+                emb = self.get_image_embedding(path)
+                self.collection.add(ids=[fname], embeddings=[emb], metadatas=[{"path": path}])
+            except Exception as e:
+                tqdm.write(f"⚠️ Skipped {fname}: {e}")
+
+    def search(self, text_query=None, image_file=None, text_weight=0.5):
+        """Hybrid search blending visual and text embeddings."""
+        img_vec = None
+        txt_vec = None
+
+        if image_file:
+            img_vec = torch.tensor(self.get_image_embedding(image_file))
+        if text_query:
+            txt_vec = torch.tensor(self.get_text_embedding(text_query))
+
+        # BLENDING LOGIC
+        if img_vec is not None and txt_vec is not None:
+            # Combined and then re-normalized to maintain 1.0 length
+            combined = (img_vec * (1.0 - text_weight)) + (txt_vec * text_weight)
+            query_emb = F.normalize(combined, p=2, dim=0).tolist()
+        elif img_vec is not None:
+            query_emb = img_vec.tolist()
+        elif txt_vec is not None:
+            query_emb = txt_vec.tolist()
+        else:
+            return []
+
+        results = self.collection.query(query_embeddings=[query_emb], n_results=10)
+
+        output = []
+        for i in range(len(results['ids'][0])):
+            fname = results['ids'][0][i]
+            # distance for 'cosine' is 1 - similarity. 
+            # 0 distance = perfect match.
+            score = round((1.0 - results['distances'][0][i]) * 100)
+            output.append({
+                "id": fname,
+                "url": f"http://localhost:8000/Productimages/{fname}",
+                "score": score
+            })
+        return output

requirements.txt

@@ -0,0 +1,10 @@
+fastapi
+uvicorn
+python-multipart
+chromadb
+sentence-transformers
+torch
+pillow
+transformers
+numpy
+huggingface_hub

server.py

@@ -0,0 +1,32 @@
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from orbiitt_engine import OrbiittEngine
+import shutil
+import os
+
+app = FastAPI()
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+# Mount your specific folder
+app.mount("/Productimages", StaticFiles(directory="Productimages"), name="Productimages")
+
+engine = OrbiittEngine()
+
+@app.post("/search")
+async def search_endpoint(text: str = Form(None), weight: float = Form(0.5), file: UploadFile = File(None)):
+    temp_path = None
+    if file:
+        temp_path = f"temp_{file.filename}"
+        with open(temp_path, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+
+    results = engine.search(text_query=text, image_file=temp_path, text_weight=weight)
+    
+    if temp_path and os.path.exists(temp_path):
+        os.remove(temp_path)
+    return {"results": results}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)