Update server.py

server.py

@@ -9,21 +9,25 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from huggingface_hub import snapshot_download
 
-# 1. SETUP LOGGING
+# 1. LOGGING & PATHS
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# 2. DEFINE PATHS
 DATASET_REPO = "aniketkumar1106/orbit-data" 
 IMAGE_DIR = "Productimages"
 DB_TARGET = "orbiitt.db" 
 
-# 3. CRITICAL FIX: Create directory BEFORE FastAPI initialization
+# 2. ACCURACY CONFIGURATION
+# Scores below this are considered "random noise" and discarded.
+# 0.22 is the standard for SigLIP 2 to separate categories effectively.
+MIN_CONFIDENCE_THRESHOLD = 0.22 
+
+# 3. DIRECTORY BOOTSTRAP
+# Fixes "RuntimeError: Directory 'Productimages' does not exist"
 if not os.path.exists(IMAGE_DIR):
     os.makedirs(IMAGE_DIR, exist_ok=True)
-    logger.info(f"Created initial directory: {IMAGE_DIR}")
 
-# 4. INITIALIZE APP
+# 4. APP INITIALIZATION
 app = FastAPI()
 
 app.add_middleware(
@@ -34,16 +38,20 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 5. MOUNT STORAGE (Will no longer crash)
+# Mount static files (guaranteed to find folder now)
 app.mount("/Productimages", StaticFiles(directory=IMAGE_DIR), name="Productimages")
 
-# 6. GLOBAL ENGINE STATE
+# 5. GLOBAL ENGINE STATE
 engine = None
-loading_status = "System Booting..."
+loading_status = "System Initializing..."
 
 def normalize_filename(name):
-    """Handles smart quotes and complex symbols to prevent 404s."""
+    """Deep normalization to ensure AI path strings match local disk files."""
+    if not name: return ""
+    name = urllib.parse.unquote(name)
+    # Convert smart quotes/dashes to standard ASCII
     name = name.replace('’', "'").replace('‘', "'").replace('“', '"').replace('”', '"').replace('–', '-')
+    # Strip non-ASCII/Unicode decomposition
     name = unicodedata.normalize('NFKD', name).encode('ascii', 'ignore').decode('ascii')
     return name.strip()
 
@@ -52,13 +60,13 @@ def setup_environment():
     global engine, loading_status
     token = os.environ.get("HF_TOKEN")
     
-    # We clean and re-create inside startup to ensure a fresh sync
+    # Clean and re-sync on boot
     if os.path.exists(IMAGE_DIR): 
         shutil.rmtree(IMAGE_DIR)
     os.makedirs(IMAGE_DIR, exist_ok=True)
 
     try:
-        loading_status = "Syncing Orbit Assets..."
+        loading_status = "Downloading Orbit Assets..."
         logger.info(loading_status)
         snapshot_download(repo_id=DATASET_REPO, repo_type="dataset", token=token, local_dir=".")
 
@@ -72,13 +80,12 @@ def setup_environment():
                     if f.lower() in ["orbiitt_db", "orbiitt.db"] or (f.endswith(".db") and not f.startswith(".")):
                         shutil.copy(src, f"./{DB_TARGET}")
                     elif f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-                        # Normalize filename on disk
-                        clean_name = normalize_filename(urllib.parse.unquote(f))
+                        # Normalize BEFORE saving to disk
+                        clean_name = normalize_filename(f)
                         shutil.copy(src, os.path.join(IMAGE_DIR, clean_name))
             shutil.rmtree("temp_extract")
 
-        loading_status = "Waking up AI Engine..."
-        logger.info(loading_status)
+        loading_status = "Loading SigLIP 2 Engine..."
         try:
             from orbiitt_engine import OrbiittEngine
             engine = OrbiittEngine()
@@ -108,30 +115,45 @@ async def search(text: str = Form(None), weight: float = Form(0.5), file: Upload
             with open(t_path, "wb") as b: 
                 b.write(await file.read())
         
-        results = engine.search(text_query=text, image_file=t_path, text_weight=weight)
+        # Increase top_k to 40 so we have a larger pool to filter from
+        results = engine.search(text_query=text, image_file=t_path, text_weight=weight, top_k=40)
         
         all_files = os.listdir(IMAGE_DIR)
         valid_results = []
 
         for r in results:
+            score = r.get('score', 0)
+            
+            # ACCURACY FIX: The Confidence Gate
+            # Prevents random cross-category matching
+            if score < MIN_CONFIDENCE_THRESHOLD:
+                continue
+
             raw_path = r.get('url') or r.get('path') or ""
-            fname = normalize_filename(os.path.basename(raw_path))
+            fname_from_ai = normalize_filename(os.path.basename(raw_path))
             
-            # Fuzzy match check
+            # Precise Matching Logic
             match = None
-            if fname in all_files:
-                match = fname
+            if fname_from_ai in all_files:
+                match = fname_from_ai
             else:
+                # Substring fallback for robustness
                 for disk_file in all_files:
-                    if fname[:20] in disk_file: 
+                    if fname_from_ai[:20].lower() in disk_file.lower():
                         match = disk_file
                         break
             
             if match:
-                r['url'] = f"Productimages/{urllib.parse.quote(match)}"
-                valid_results.append(r)
+                valid_results.append({
+                    "id": r.get('id', 'Product'),
+                    "url": f"Productimages/{urllib.parse.quote(match)}",
+                    "score": score
+                })
+        
+        # Keep results strictly sorted by AI confidence
+        valid_results.sort(key=lambda x: x['score'], reverse=True)
                 
-        return {"results": valid_results}
+        return {"results": valid_results[:20]} # Return top 20 verified matches
     except Exception as e:
         logger.error(f"Search error: {e}")
         return {"results": [], "error": str(e)}