Update server.py

server.py

@@ -1,14 +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
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 
-app = FastAPI(title="ORBIT Visual Search Engine")
+app = FastAPI(title="ORBIT Visual Engine")
 
-# 1. Enable Global CORS for Netlify communication
+# --- 1. GLOBALLY ENABLE CORS ---
+# Required for your Netlify frontend to talk to this API
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -17,96 +17,68 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- CONFIGURATION (UPDATE THESE) ---
+# --- 2. DATASET CONFIGURATION ---
 DATASET_REPO = "aniketkumar1106/orbit-data" 
-ZIP_FILENAME = "orbiitt_db" 
 IMAGE_DIR = "Productimages"
+DB_FILE = "orbiitt.db"
 
-def setup_resources():
-    """Downloads, Deep-Scans, and Extracts data from private HF Dataset."""
+def sync_dataset():
+    """Syncs the database and images from the HF private dataset folder."""
     token = os.environ.get("HF_TOKEN")
     if not token:
-        print("CRITICAL: HF_TOKEN secret is missing!")
+        print("CRITICAL ERROR: HF_TOKEN not found in Space Secrets.")
         return
 
-    # Clean start: Create image directory if missing
-    if not os.path.exists(IMAGE_DIR):
-        os.makedirs(IMAGE_DIR)
-
     try:
-        # 1. Download zip from private dataset
-        print(f"Fetching {ZIP_FILENAME} from {DATASET_REPO}...")
-        zip_path = hf_hub_download(
+        print(f"Syncing folder-based dataset from {DATASET_REPO}...")
+        
+        # snapshot_download downloads the folder structure as-is
+        snapshot_download(
             repo_id=DATASET_REPO,
-            filename=ZIP_FILENAME,
             repo_type="dataset",
             token=token,
-            local_dir="."
+            local_dir=".", 
+            allow_patterns=[DB_FILE, f"{IMAGE_DIR}/*"]
         )
-
-        # 2. Extract to a temporary folder
-        temp_extract = "temp_data"
-        if os.path.exists(temp_extract):
-            shutil.rmtree(temp_extract)
-            
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            zip_ref.extractall(temp_extract)
-
-        # 3. DEEP SCAN: Recursively find orbiitt.db and images
-        print("Scanning extracted files...")
-        for root, dirs, files in os.walk(temp_extract):
-            for file in files:
-                file_lower = file.lower()
-                source_path = os.path.join(root, file)
-                
-                # If we find the database, move it to the root
-                if file == "orbiitt.db":
-                    shutil.move(source_path, "./orbiitt.db")
-                    print(f"Moved {file} to root.")
-                
-                # If we find images, move them to Productimages
-                elif file_lower.endswith(('.png', '.jpg', '.jpeg', '.webp')):
-                    target_path = os.path.join(IMAGE_DIR, file)
-                    # Move only if it doesn't exist yet
-                    if not os.path.exists(target_path):
-                        shutil.move(source_path, target_path)
-
-        # 4. Cleanup temp folder
-        shutil.rmtree(temp_extract)
-        print("Setup complete. Environment is organized.")
-
+        
+        print("Sync Successful.")
+        print(f"Files found in root: {os.listdir('.')}")
     except Exception as e:
-        print(f"Initialization Failed: {e}")
+        print(f"Sync Failed: {e}")
 
-# Run setup before mounting or engine loading
-setup_resources()
+# Run sync before anything else starts
+sync_dataset()
 
-# 2. Mount Static Files (Frontend access to images)
+# --- 3. MOUNT STATIC IMAGES ---
+# Allows the frontend to access images at: https://your-space.hf.space/Productimages/1.jpg
 if os.path.exists(IMAGE_DIR):
     app.mount(f"/{IMAGE_DIR}", StaticFiles(directory=IMAGE_DIR), name=IMAGE_DIR)
 
-# 3. Load the OrbiittEngine
+# --- 4. INITIALIZE SEARCH ENGINE ---
 engine = None
 try:
     from orbiitt_engine import OrbiittEngine
-    # Ensure database is present before starting engine
-    if os.path.exists("orbiitt.db"):
+    if os.path.exists(DB_FILE):
         engine = OrbiittEngine()
-        print("Engine successfully connected to orbiitt.db")
+        print("OrbiittEngine: Ready and Connected.")
     else:
-        print("Engine NOT loaded: orbiitt.db is missing.")
+        print(f"OrbiittEngine Error: {DB_FILE} not found after sync.")
 except Exception as e:
-    print(f"Error initializing OrbiittEngine: {e}")
+    print(f"Engine initialization failed: {e}")
+
+
+
+# --- 5. API ENDPOINTS ---
 
 @app.get("/")
-def health_check():
-    """Diagnostic check for frontend and debugging."""
+def health_status():
+    """Dashboard for checking if the server has the files it needs."""
+    img_count = len(os.listdir(IMAGE_DIR)) if os.path.exists(IMAGE_DIR) else 0
     return {
         "status": "online",
-        "database_ready": os.path.exists("orbiitt.db"),
-        "images_count": len(os.listdir(IMAGE_DIR)) if os.path.exists(IMAGE_DIR) else 0,
-        "engine_ready": engine is not None,
-        "root_files": os.listdir(".")[:10]
+        "database_ready": os.path.exists(DB_FILE),
+        "images_count": img_count,
+        "engine_ready": engine is not None
     }
 
 @app.post("/search")
@@ -116,30 +88,34 @@ async def search_endpoint(
     file: UploadFile = File(None)
 ):
     if not engine:
-        raise HTTPException(status_code=503, detail="Search engine is not initialized.")
+        raise HTTPException(status_code=503, detail="Search engine is not available.")
 
-    temp_path = None
+    temp_image_path = None
     try:
-        # Save search image if provided
+        # If an image was uploaded for search, save it temporarily
         if file:
-            temp_path = f"search_input_{file.filename}"
-            with open(temp_path, "wb") as buffer:
+            temp_image_path = f"tmp_search_{file.filename}"
+            with open(temp_image_path, "wb") as buffer:
                 buffer.write(await file.read())
 
-        # Perform the actual search
-        results = engine.search(text_query=text, image_file=temp_path, text_weight=weight)
-        
-        # Ensure result URLs point correctly to the mounted folder
+        # Perform the search
+        results = engine.search(
+            text_query=text, 
+            image_file=temp_image_path, 
+            text_weight=weight
+        )
         return {"results": results}
 
     except Exception as e:
-        print(f"Search Error: {e}")
-        raise HTTPException(status_code=500, detail="An internal search error occurred.")
+        print(f"Runtime Search Error: {e}")
+        raise HTTPException(status_code=500, detail="An internal error occurred during search.")
+    
     finally:
-        # Cleanup temp search file
-        if temp_path and os.path.exists(temp_path):
-            os.remove(temp_path)
+        # Always clean up the uploaded search image
+        if temp_image_path and os.path.exists(temp_image_path):
+            os.remove(temp_image_path)
 
 if __name__ == "__main__":
     import uvicorn
+    # Use 7860 as it is the default port for Hugging Face Spaces
     uvicorn.run(app, host="0.0.0.0", port=7860)