Update server.py

server.py

@@ -1,13 +1,14 @@
+import os
+import zipfile
+import shutil
 from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
-import shutil
-import os
-import zipfile
-import uvicorn
+from huggingface_hub import hf_hub_download
 
-app = FastAPI(title="ORBIT Visual Commerce API")
+app = FastAPI(title="ORBIT Visual Search Engine")
 
+# 1. Enable Global CORS (Crucial for Netlify)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,68 +17,96 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- NEW: UNZIP LOGIC ---
-def extract_resources():
-    # Change 'data.zip' to the actual name of your zip file
-    zip_name = "data.zip" 
-    
-    if os.path.exists(zip_name):
-        print(f"Extracting {zip_name}...")
-        with zipfile.ZipFile(zip_name, 'r') as zip_ref:
-            zip_ref.extractall(".")
-        print("Extraction complete.")
-    else:
-        print(f"Note: {zip_name} not found, skipping extraction.")
+# --- CONFIGURATION ---
+DATASET_REPO = "aniketkumar1106/your-private-dataset"  # CHANGE THIS
+ZIP_FILENAME = "product_data.zip"                    # CHANGE THIS
+IMAGE_DIR = "Productimages"
+DB_FILE = "orbiitt.db"
 
-# Run extraction before anything else
-extract_resources()
+def initialize_backend():
+    """Downloads and unzips data from a private dataset on startup."""
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("CRITICAL: HF_TOKEN secret is missing from Space Settings!")
+        return
 
-# --- SETUP DIRECTORIES ---
-IMAGE_DIR = "Productimages"
-if not os.path.exists(IMAGE_DIR):
-    os.makedirs(IMAGE_DIR)
+    try:
+        # Download the zip file into the current directory
+        print(f"Fetching {ZIP_FILENAME} from private dataset...")
+        zip_path = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=ZIP_FILENAME,
+            repo_type="dataset",
+            token=token,
+            local_dir="."
+        )
 
-app.mount(f"/{IMAGE_DIR}", StaticFiles(directory=IMAGE_DIR), name=IMAGE_DIR)
+        # Extract everything
+        print("Extracting archive...")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(".")
+            
+        # Ensure the image folder exists for mounting
+        if not os.path.exists(IMAGE_DIR):
+            os.makedirs(IMAGE_DIR)
+            
+        print(f"Setup complete. Found {len(os.listdir(IMAGE_DIR))} images.")
+    except Exception as e:
+        print(f"Startup Error: {e}")
 
-# --- LOAD ENGINE ---
+# Trigger extraction before app starts
+initialize_backend()
+
+# 2. Mount Static Files (Must happen after directory is created)
+if os.path.exists(IMAGE_DIR):
+    app.mount(f"/{IMAGE_DIR}", StaticFiles(directory=IMAGE_DIR), name=IMAGE_DIR)
+
+# 3. Load the Engine (Handles your 4GB weights)
+engine = None
 try:
     from orbiitt_engine import OrbiittEngine
-    # The engine should now find orbiitt.db in the root folder
+    # The engine should now find orbiitt.db in the root
     engine = OrbiittEngine()
-    print("OrbiittEngine initialized successfully.")
+    print("OrbiittEngine successfully initialized.")
 except Exception as e:
-    print(f"Error loading OrbiittEngine: {e}")
-    engine = None
+    print(f"Engine Load Error: {e}")
 
 @app.get("/")
-def read_root():
-    # Diagnostic: Check if the DB file exists now
-    db_exists = os.path.exists("orbiitt.db")
+def health_check():
+    """Diagnostic endpoint to check file health."""
     return {
-        "message": "Server is running",
-        "db_found": db_exists,
-        "engine_ready": engine is not None,
-        "files_in_root": os.listdir(".")[:10] # Shows first 10 files for debugging
+        "status": "online",
+        "database_ready": os.path.exists(DB_FILE),
+        "images_count": len(os.listdir(IMAGE_DIR)) if os.path.exists(IMAGE_DIR) else 0,
+        "engine_loaded": engine is not None
     }
 
 @app.post("/search")
-async def search_endpoint(text: str = Form(None), weight: float = Form(0.5), file: UploadFile = File(None)):
-    if engine is None:
-        raise HTTPException(status_code=503, detail="Engine not loaded")
+async def search(text: str = Form(None), weight: float = Form(0.5), file: UploadFile = File(None)):
+    if not engine:
+        raise HTTPException(status_code=503, detail="Search engine is still loading or failed.")
 
-    temp_path = None
+    temp_image = None
     try:
+        # Save incoming search image if provided
         if file:
-            temp_path = f"temp_{file.filename}"
-            with open(temp_path, "wb") as buffer:
-                content = await file.read()
-                buffer.write(content)
+            temp_image = f"search_input_{file.filename}"
+            with open(temp_image, "wb") as buffer:
+                buffer.write(await file.read())
 
-        results = engine.search(text_query=text, image_file=temp_path, text_weight=weight)
+        # Perform the actual model search
+        results = engine.search(text_query=text, image_file=temp_image, text_weight=weight)
         return {"results": results}
+
+    except Exception as e:
+        print(f"Runtime Search Error: {e}")
+        raise HTTPException(status_code=500, detail="Search failed internally.")
     finally:
-        if temp_path and os.path.exists(temp_path):
-            os.remove(temp_path)
+        # Clean up temporary search image
+        if temp_image and os.path.exists(temp_image):
+            os.remove(temp_image)
 
 if __name__ == "__main__":
+    import uvicorn
+    # Use standard port for Hugging Face Spaces
     uvicorn.run(app, host="0.0.0.0", port=7860)