Update app.py

app.py

@@ -14,9 +14,8 @@ import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 
-from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Query
+from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Form
 from fastapi.responses import RedirectResponse
-from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient
 
@@ -28,7 +27,9 @@ from gradio import mount_gradio_app
 import boto3
 from botocore.client import Config
 from io import BytesIO
-from typing import Optional 
+from typing import Optional
+import requests
+
 # --------------------- Logging ---------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -42,13 +43,40 @@ os.makedirs(MODELS_DIR, exist_ok=True)
 
 # --------------------- Secrets ---------------------
 HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face private repo token
-API_SECRET_TOKEN = os.getenv("API_SECRET_TOKEN")  # Bearer token for API
+# Firebase credentials JSON
+FIREBASE_CREDENTIALS_PATH = os.getenv("FIREBASE_CREDENTIALS_PATH")
+
 # --------------------- DigitalOcean Spaces Credentials ---------------------
-DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")  # Default region = Bangalore
+DO_SPACES_REGION = os.getenv("DO_SPACES_REGION", "blr1")
 DO_SPACES_ENDPOINT = os.getenv("DO_SPACES_ENDPOINT", f"https://{DO_SPACES_REGION}.digitaloceanspaces.com")
-DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")  # Your Access Key
-DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")  # Your Secret Key
-DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")  # Bucket Name
+DO_SPACES_KEY = os.getenv("DO_SPACES_KEY")
+DO_SPACES_SECRET = os.getenv("DO_SPACES_SECRET")
+DO_SPACES_BUCKET = os.getenv("DO_SPACES_BUCKET")
+
+# --------------------- Firebase Auth ---------------------
+import firebase_admin
+from firebase_admin import credentials, auth
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+if not firebase_admin._apps:
+    cred = credentials.Certificate(FIREBASE_CREDENTIALS_PATH)
+    firebase_admin.initialize_app(cred)
+    logger.info("✅ Firebase initialized successfully")
+
+security = HTTPBearer()
+
+def verify_firebase_token(credentials: HTTPAuthorizationCredentials = Security(security)):
+    """Verify Firebase ID token from Authorization header."""
+    try:
+        id_token = credentials.credentials
+        decoded_token = auth.verify_id_token(id_token)
+        user_email = decoded_token.get("email")
+        if not user_email:
+            raise HTTPException(status_code=401, detail="Firebase token invalid or missing email")
+        return user_email
+    except Exception as e:
+        logger.error(f"Firebase auth failed: {e}")
+        raise HTTPException(status_code=401, detail="Unauthorized: Invalid Firebase token")
 
 # --------------------- Download Models ---------------------
 def download_models():
@@ -100,7 +128,7 @@ def ensure_codeformer():
 
 ensure_codeformer()
 
-# --------------------- MongoDB (for API logs only) ---------------------
+# --------------------- MongoDB ---------------------
 MONGODB_URL = os.getenv("MONGODB_URL")
 
 client = None
@@ -124,22 +152,14 @@ async def shutdown_db():
         client.close()
         logger.info("MongoDB connection closed")
 
-# --------------------- Auth ---------------------
-security = HTTPBearer()
-
-def verify_token(credentials: HTTPAuthorizationCredentials = Security(security)):
-    if credentials.credentials != API_SECRET_TOKEN:
-        raise HTTPException(status_code=401, detail="Invalid or missing token")
-    return credentials.credentials
-
 # --------------------- Logging API Hits ---------------------
-async def log_faceswap_hit(token: str, status: str = "success"):
+async def log_faceswap_hit(user_email: str, status: str = "success"):
     global database
     if database is None:
         return
     await database.api_logs.insert_one({
-        "token": token,
-        "endpoint": "/faceswap",
+        "user": user_email,
+        "endpoint": "/face-swap",
         "status": status,
         "timestamp": datetime.utcnow()
     })
@@ -150,7 +170,6 @@ swap_lock = threading.Lock()
 def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     try:
         with swap_lock:
-            # Prepare temporary directory
             if os.path.exists(temp_dir):
                 shutil.rmtree(temp_dir)
             os.makedirs(temp_dir, exist_ok=True)
@@ -164,71 +183,16 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
             if not src_faces or not tgt_faces:
                 return None, None, "❌ Face not detected in source or target image"
 
-            def expand_bbox(bbox, img_shape, scale=1.6):
-                ih, iw = img_shape[:2]
-                x1, y1, x2, y2 = map(int, bbox)
-                w, h = x2 - x1, y2 - y1
-                cx, cy = x1 + w // 2, y1 + h // 2
-                new_w, new_h = int(w * scale), int(h * scale)
-                nx1 = max(0, cx - new_w // 2)
-                ny1 = max(0, cy - new_h // 2)
-                nx2 = min(iw, cx + new_w // 2)
-                ny2 = min(ih, cy + new_h // 2)
-                return nx1, ny1, nx2, ny2
-
             src_face0 = src_faces[0]
             tgt_face0 = tgt_faces[0]
 
-            # More accurate source face crop with slight expansion
-            s_x1, s_y1, s_x2, s_y2 = expand_bbox(src_face0.bbox, src_bgr.shape, scale=1.4)
-            src_crop = src_bgr[s_y1:s_y2, s_x1:s_x2]
-            src_crop_faces = face_analysis_app.get(src_crop)
-            if src_crop_faces:
-                src_for_swap = src_crop
-                src_face_for_swap = src_crop_faces[0]
-            else:
-                src_for_swap = src_bgr
-                src_face_for_swap = src_face0
-
-            # More aggressive target crop for precise landmark detection
-            t_x1, t_y1, t_x2, t_y2 = expand_bbox(tgt_face0.bbox, tgt_bgr_full.shape, scale=1.6)
-            tgt_crop = tgt_bgr_full[t_y1:t_y2, t_x1:t_x2]
-            tgt_crop_faces = face_analysis_app.get(tgt_crop)
-
-            if tgt_crop_faces:
-                tgt_for_swap = tgt_crop
-                tgt_face_for_swap = tgt_crop_faces[0]
-
-                swapped_crop = swapper.get(tgt_for_swap, tgt_face_for_swap, src_face_for_swap)
-                if swapped_crop is None:
-                    return None, None, "❌ Face swap failed on crop"
-
-                # Create mask with threshold for seamlessClone
-                mask = cv2.cvtColor(swapped_crop, cv2.COLOR_BGR2GRAY)
-                _, mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
-
-                center = ((t_x1 + t_x2) // 2, (t_y1 + t_y2) // 2)
-
-                try:
-                    blended = cv2.seamlessClone(swapped_crop, tgt_bgr_full, mask, center, cv2.NORMAL_CLONE)
-                except Exception:
-                    # Fallback to direct paste if seamlessClone fails
-                    blended = tgt_bgr_full.copy()
-                    h, w = swapped_crop.shape[:2]
-                    blended[t_y1:t_y1+h, t_x1:t_x1+w] = swapped_crop
-
-                swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
-                cv2.imwrite(swapped_path, blended)
-
-            else:
-                # Fallback: swap on full image if crop detection fails
-                swapped_bgr_full = swapper.get(tgt_bgr_full, tgt_face0, src_face0)
-                if swapped_bgr_full is None:
-                    return None, None, "❌ Face swap failed on full image"
-                swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
-                cv2.imwrite(swapped_path, swapped_bgr_full)
-
-            # Run CodeFormer enhancement on the swapped image
+            swapped_bgr_full = swapper.get(tgt_bgr_full, tgt_face0, src_face0)
+            if swapped_bgr_full is None:
+                return None, None, "❌ Face swap failed"
+
+            swapped_path = os.path.join(temp_dir, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
+            cv2.imwrite(swapped_path, swapped_bgr_full)
+
             cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {temp_dir} --bg_upsampler realesrgan --face_upsample"
             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
             if result.returncode != 0:
@@ -247,8 +211,6 @@ def face_swap_and_enhance(src_img, tgt_img, temp_dir="/tmp/faceswap_work"):
     except Exception as e:
         return None, None, f"❌ Error: {str(e)}"
 
-
-
 # --------------------- Gradio ---------------------
 with gr.Blocks() as demo:
     gr.Markdown("Face Swap")
@@ -300,72 +262,55 @@ def root():
 async def health():
     return {"status": "healthy"}
 
-from fastapi import Form
-import requests
-
-@fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
+@fastapi_app.post("/face-swap")
 async def face_swap_api(
     source: UploadFile = File(...),
     target_category_id: str = Form(...),
-    credentials: HTTPAuthorizationCredentials = Security(security)
+    user_email: str = Depends(verify_firebase_token)
 ):
     try:
-        # Read source image
         src_bytes = await source.read()
-
-        # Save source image to Spaces under bikini-theme/source/
         src_key = f"bikini-theme/source/{uuid.uuid4().hex}_{source.filename}"
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
-        # Build target image URL directly (no category scanning)
         target_filename = f"{target_category_id}.jpg"
         target_url = f"https://{DO_SPACES_BUCKET}.{DO_SPACES_REGION}.digitaloceanspaces.com/bikini-theme/target/{target_filename}"
 
-        # Try to fetch the target image
         resp = requests.get(target_url)
         if resp.status_code != 200:
-            await log_faceswap_hit(credentials.credentials, status="error")
-            raise HTTPException(
-                status_code=404, 
-                detail=f"Target image not found at {target_url}"
-            )
+            await log_faceswap_hit(user_email, status="error")
+            raise HTTPException(status_code=404, detail=f"Target image not found at {target_url}")
 
         tgt_bytes = resp.content
 
-        # Decode source and target images
         src_array = np.frombuffer(src_bytes, np.uint8)
         tgt_array = np.frombuffer(tgt_bytes, np.uint8)
         src_bgr = cv2.imdecode(src_array, cv2.IMREAD_COLOR)
         tgt_bgr = cv2.imdecode(tgt_array, cv2.IMREAD_COLOR)
 
         if src_bgr is None or tgt_bgr is None:
-            await log_faceswap_hit(credentials.credentials, status="error")
+            await log_faceswap_hit(user_email, status="error")
             raise HTTPException(status_code=400, detail="Invalid image data")
 
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
-        # Run face swap and enhancement
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
-            await log_faceswap_hit(credentials.credentials, status="error")
+            await log_faceswap_hit(user_email, status="error")
             raise HTTPException(status_code=500, detail=err)
 
-        # Upload the result back to bikini-theme/result/
         with open(final_path, "rb") as f:
             result_bytes = f.read()
         result_key = f"bikini-theme/result/{uuid.uuid4().hex}_enhanced.png"
         result_url = upload_to_spaces(result_bytes, result_key, content_type="image/png")
 
-        # Log success in MongoDB
-        await log_faceswap_hit(credentials.credentials, status="success")
+        await log_faceswap_hit(user_email, status="success")
 
-        return {
-            "result_url": result_url
-        }
+        return {"result_url": result_url}
 
     except Exception as e:
-        await log_faceswap_hit(credentials.credentials, status="error")
+        await log_faceswap_hit(user_email, status="error")
         raise HTTPException(status_code=500, detail=f"Face swap failed: {str(e)}")
 
 @fastapi_app.get("/preview/{result_key:path}")
@@ -379,8 +324,9 @@ async def preview_result(result_key: str):
         media_type="image/png",
         headers={"Content-Disposition": "inline; filename=result.png"}
     )
+
 # --------------------- Mount Gradio ---------------------
 fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 
 if __name__ == "__main__":
-    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)
+    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)