Update app.py

app.py

@@ -254,55 +254,76 @@ async def face_swap_api(
     source: UploadFile = File(...),
     target_category_id: str = Form(None),
     new_category_id: str = Form(None),
-    user_id: Optional[str] = Form(None)
+    user_id: Optional[str] = Form(None),
+    credentials: HTTPAuthorizationCredentials = Security(security)
 ):
     start_time = datetime.utcnow()
 
     try:
-        # ============ VALIDATION ============
+        # ------------------------------------------------------------------
+        # VALIDATION
+        # ------------------------------------------------------------------
         if target_category_id and new_category_id:
             raise HTTPException(400, "Provide only one of new_category_id or target_category_id.")
 
         if not target_category_id and not new_category_id:
             raise HTTPException(400, "Either new_category_id or target_category_id is required.")
 
-        # Read source image
+        # ------------------------------------------------------------------
+        # READ SOURCE IMAGE
+        # ------------------------------------------------------------------
         src_bytes = await source.read()
         src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
-        upload_to_spaces(src_bytes, src_key)
+        upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
 
-        # ============ CASE 1: MongoDB new_category_id ============
+        # ------------------------------------------------------------------
+        # CASE 1 : new_category_id → MongoDB lookup
+        # ------------------------------------------------------------------
         if new_category_id:
-            doc = await subcategories_col.find_one({"asset_images._id": ObjectId(new_category_id)})
+
+            doc = await subcategories_col.find_one({
+                "asset_images._id": ObjectId(new_category_id)
+            })
 
             if not doc:
-                raise HTTPException(404, "Asset image not found.")
+                raise HTTPException(404, "Asset image not found in database")
 
-            # find correct asset image
-            asset = next((img for img in doc["asset_images"] if str(img["_id"]) == new_category_id), None)
+            # extract correct asset
+            asset = next(
+                (img for img in doc["asset_images"] if str(img["_id"]) == new_category_id),
+                None
+            )
 
             if not asset:
-                raise HTTPException(404, "Asset image URL not found.")
+                raise HTTPException(404, "Asset image URL not found")
 
+            # correct URL
             target_url = asset["url"]
-            category_oid = doc["categoryId"]  # REAL ObjectId
 
-            # ---- MEDIA_CLICKS ONLY IF user_id EXISTS ----
+            # correct categoryId (ObjectId)
+            category_oid = doc["categoryId"]   # <-- DO NOT CONVERT TO STRING
+
+            # ------------------------------------------------------------------
+            # MEDIA_CLICKS (ONLY IF user_id PRESENT)
+            # ------------------------------------------------------------------
             if user_id:
                 try:
-                    user_id = user_id.strip()
-                    user_oid = ObjectId(user_id)
+                    user_oid = ObjectId(user_id.strip())
                     now = datetime.utcnow()
 
                     update_result = await media_clicks_col.update_one(
                         {"userId": user_oid, "categories.categoryId": category_oid},
                         {
-                            "$set": {"updatedAt": now, "categories.$.lastClickedAt": now},
+                            "$set": {
+                                "updatedAt": now,
+                                "categories.$.lastClickedAt": now
+                            },
                             "$inc": {"categories.$.click_count": 1}
                         }
                     )
 
                     if update_result.matched_count == 0:
+                        # category does not exist → push new entry
                         await media_clicks_col.update_one(
                             {"userId": user_oid},
                             {
@@ -319,62 +340,81 @@ async def face_swap_api(
                             upsert=True
                         )
 
-                except Exception as err:
-                    logger.error(f"MEDIA_CLICK ERROR: {err}")
-                    # Do NOT break face-swap
+                except Exception as media_err:
+                    logger.error(f"MEDIA_CLICK ERROR: {media_err}")
+                    # do NOT interrupt faceswap flow
 
-        # ============ CASE 2: DigitalOcean target_category_id ============
+        # ------------------------------------------------------------------
+        # CASE 2 : target_category_id → DigitalOcean path (unchanged logic)
+        # ------------------------------------------------------------------
         if target_category_id:
             client = get_spaces_client()
             base_prefix = "faceswap/target/"
-            resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
-
+            resp = client.list_objects_v2(
+                Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/"
+            )
             categories = [p["Prefix"].split("/")[2] for p in resp.get("CommonPrefixes", [])]
+
             target_url = None
 
             for category in categories:
                 original_prefix = f"faceswap/target/{category}/original/"
-                original_files = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
+                original_objs = client.list_objects_v2(
+                    Bucket=DO_SPACES_BUCKET, Prefix=original_prefix
+                ).get("Contents", [])
 
-                for idx, obj in enumerate(sorted(original_files.get("Contents", [])), start=1):
+                original_files = sorted([obj["Key"] for obj in original_objs])
+
+                for idx, file_key in enumerate(original_files, start=1):
                     cid = f"{category.lower()}image_{idx}"
                     if cid == target_category_id:
-                        filename = obj["Key"].split("/")[-1]
-                        target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{original_prefix}{filename}"
+                        target_url = f"{DO_SPACES_ENDPOINT}/{DO_SPACES_BUCKET}/{file_key}"
                         break
 
                 if target_url:
                     break
 
             if not target_url:
-                raise HTTPException(404, "Target categoryId not found.")
+                raise HTTPException(404, "Target categoryId not found")
+
+        # ------------------------------------------------------------------
+        # DOWNLOAD TARGET IMAGE
+        # ------------------------------------------------------------------
+        tgt_bytes = requests.get(target_url).content
 
-        # ============ DOWNLOAD TARGET IMAGE ============
-        target_bytes = requests.get(target_url).content
-        tgt_bgr = cv2.imdecode(np.frombuffer(target_bytes, np.uint8), cv2.IMREAD_COLOR)
         src_bgr = cv2.imdecode(np.frombuffer(src_bytes, np.uint8), cv2.IMREAD_COLOR)
+        tgt_bgr = cv2.imdecode(np.frombuffer(tgt_bytes, np.uint8), cv2.IMREAD_COLOR)
 
         if src_bgr is None or tgt_bgr is None:
             raise HTTPException(400, "Invalid image data")
 
-        # ============ FACE SWAP ============
         src_rgb = cv2.cvtColor(src_bgr, cv2.COLOR_BGR2RGB)
         tgt_rgb = cv2.cvtColor(tgt_bgr, cv2.COLOR_BGR2RGB)
 
+        # ------------------------------------------------------------------
+        # FACE SWAP EXECUTION
+        # ------------------------------------------------------------------
         final_img, final_path, err = face_swap_and_enhance(src_rgb, tgt_rgb)
         if err:
             raise HTTPException(500, err)
 
-        result_bytes = open(final_path, "rb").read()
-        result_key = f"faceswap/result/{uuid.uuid4().hex}.png"
+        with open(final_path, "rb") as f:
+            result_bytes = f.read()
+
+        result_key = f"faceswap/result/{uuid.uuid4().hex}_enhanced.png"
         result_url = upload_to_spaces(result_bytes, result_key)
 
-        return {"result_key": result_key, "result_url": result_url}
+        return {
+            "result_key": result_key,
+            "result_url": result_url
+        }
 
     except Exception as e:
         raise HTTPException(500, f"Face swap failed: {str(e)}")
 
 
+
+
 # @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
 # async def face_swap_api(
 #     source: UploadFile = File(...),