Update main.py

main.py

@@ -1,13 +1,12 @@
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from typing import List
-from PIL import Image
 import os
 import shutil
 import uuid
 import re
-import inflect  # <-- NEW: Import inflect
-
+import inflect 
+import cloudinary.api
 from src.models import AIModelManager
 from src.cloud_db import CloudDB
 
@@ -15,7 +14,7 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], 
+    allow_origins=["*"],  
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -23,82 +22,88 @@ app.add_middleware(
 
 print("Loading AI Models and Cloud DB...")
 ai = AIModelManager()
-db = CloudDB()
-# Initialize the inflect engine
+db = CloudDB() 
 p = inflect.engine() 
 print("Ready!")
 
 os.makedirs("temp_uploads", exist_ok=True)
 
-# --- NEW: Standardization Function ---
 def standardize_category_name(name: str) -> str:
-    """Converts ' Cows ', 'COWS', or 'cow' all into 'cow'."""
-    # 1. Lowercase and strip accidental edge spaces
     clean_name = name.strip().lower()
-    
-    # 2. Replace inner spaces with underscores (e.g., 'sports cars' -> 'sports_cars')
     clean_name = re.sub(r'\s+', '_', clean_name)
-    
-    # 3. Remove weird special characters just in case (keep only letters, numbers, underscores)
     clean_name = re.sub(r'[^\w\s]', '', clean_name)
-    
-    # 4. Convert plural to singular (if it's already singular, it returns False, so we keep the clean_name)
     singular_name = p.singular_noun(clean_name)
-    if singular_name:
-        return singular_name
-        
+    return singular_name if singular_name else clean_name
+
+def sanitize_filename(filename: str) -> str:
+    clean_name = re.sub(r'\s+', '_', filename)
+    clean_name = re.sub(r'[^\w\.\-]', '', clean_name)
     return clean_name
-# -------------------------------------
 
 @app.post("/api/upload")
 async def upload_new_images(files: List[UploadFile] = File(...), folder_name: str = Form(...)):
-    """Handles bulk uploading of multiple images at once."""
     uploaded_urls = []
-    
-    # Clean the folder name before doing anything else!
     standardized_folder = standardize_category_name(folder_name)
     
     try:
         for file in files:
-            temp_path = f"temp_uploads/{file.filename}"
+            safe_filename = sanitize_filename(file.filename)
+            temp_path = f"temp_uploads/{safe_filename}"
             with open(temp_path, "wb") as buffer:
                 shutil.copyfileobj(file.file, buffer)
                 
-            # Upload to Cloudinary using the perfectly clean folder name
             image_url = db.upload_image(temp_path, standardized_folder)
             
-            img = Image.open(temp_path).convert('RGB')
-            vector = ai.encode_image(img)
+            vectors_to_save = ai.process_image(temp_path, is_query=False)
             
-            image_id = str(uuid.uuid4())
-            db.add_vector(vector, image_url, image_id)
+            for vec_dict in vectors_to_save:
+                image_id = str(uuid.uuid4())
+                db.add_vector(vec_dict, image_url, image_id)
             
             os.remove(temp_path)
             uploaded_urls.append(image_url)
             
-        # Return the standardized name so the frontend knows what was actually saved
-        return {
-            "message": f"Successfully added {len(files)} images to category '{standardized_folder}'!", 
-            "urls": uploaded_urls
-        }
-        
+        return {"message": "Success!", "urls": uploaded_urls}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# ... (Your /api/search endpoint stays exactly the same) ...
 @app.post("/api/search")
 async def search_database(file: UploadFile = File(...)):
     try:
-        temp_path = f"temp_uploads/query_{file.filename}"
+        safe_filename = sanitize_filename(file.filename)
+        temp_path = f"temp_uploads/query_{safe_filename}"
+        
         with open(temp_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
             
-        img = Image.open(temp_path).convert('RGB')
-        vector = ai.encode_image(img)
+        vectors_to_search = ai.process_image(temp_path, is_query=True)
         
-        results = db.search(vector, top_k=10)
+        all_results = []
+        for vec_dict in vectors_to_search:
+            results = db.search(vec_dict, top_k=10)
+            all_results.extend(results)
+            
         os.remove(temp_path)
         
-        return {"results": results}
+        unique_results = {}
+        for r in all_results:
+            url = r["url"]
+            if url not in unique_results or r["score"] > unique_results[url]["score"]:
+                unique_results[url] = r
+                
+        final_results = sorted(unique_results.values(), key=lambda x: x["score"], reverse=True)
+        
+        return {"results": final_results[:10]}
+    except Exception as e:
+        print(f"Production Search Error: {str(e)}") 
+        raise HTTPException(status_code=500, detail=str(e))
+    
+@app.get("/api/categories")
+async def get_categories():
+    try:
+        result = cloudinary.api.root_folders()
+        folders = [folder["name"] for folder in result.get("folders", [])]
+        return {"categories": folders}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        print(f"Error fetching categories from Cloudinary: {e}")
+        return {"categories": []}