added /list-images and faceswap accepts categoryId

app.py

@@ -1,3 +1,4 @@
+# --------------------- List Images Endpoint ---------------------
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 import shutil
@@ -13,7 +14,7 @@ import insightface
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 
-from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security
+from fastapi import FastAPI, UploadFile, File, HTTPException, Response, Depends, Security, Query
 from fastapi.responses import RedirectResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
@@ -103,7 +104,7 @@ ensure_codeformer()
 # --------------------- MongoDB (for API logs only) ---------------------
 MONGODB_URL = os.getenv("MONGODB_URL")
 
-client: AsyncIOMotorClient = None
+client = None
 database = None
 
 # --------------------- FastAPI ---------------------
@@ -235,26 +236,58 @@ def download_from_spaces(key):
 def root():
     return RedirectResponse("/gradio")
 
-@fastapi_app.get("/health", dependencies=[Depends(verify_token)])
+@fastapi_app.get("/health")
 async def health():
     return {"status": "healthy"}
 
+from fastapi import Form
+import requests
+
 @fastapi_app.post("/face-swap", dependencies=[Depends(verify_token)])
 async def face_swap_api(
     source: UploadFile = File(...),
-    target: UploadFile = File(...),
+    target_category_id: str = Form(...),
     credentials: HTTPAuthorizationCredentials = Security(security)
 ):
+
     try:
-        # Read images
+        # Read source image
         src_bytes = await source.read()
-        tgt_bytes = await target.read()
 
-        # Save to Spaces (source/target)
+        # Save source to Spaces
         src_key = f"faceswap/source/{uuid.uuid4().hex}_{source.filename}"
-        tgt_key = f"faceswap/target/{uuid.uuid4().hex}_{target.filename}"
         upload_to_spaces(src_bytes, src_key, content_type=source.content_type)
-        upload_to_spaces(tgt_bytes, tgt_key, content_type=target.content_type)
+
+        # Find target image URL from categoryId by searching all categories
+        client = get_spaces_client()
+        base_prefix = "faceswap/target/"
+        resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
+        categories = [prefix["Prefix"].split("/")[2] for prefix in resp.get("CommonPrefixes", [])]
+        target_url = None
+        for category in categories:
+            original_prefix = f"faceswap/target/{category}/original/"
+            thumb_prefix = f"faceswap/target/{category}/thumb/"
+            original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
+            thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
+            original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+            thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+            for idx, filename in enumerate(sorted(original_files), start=1):
+                cid = f"{category.lower()}image_{idx}"
+                if filename in thumb_files and cid == target_category_id:
+                    target_url = f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}"
+                    break
+            if target_url:
+                break
+        if not target_url:
+            await log_faceswap_hit(credentials.credentials, status="error")
+            raise HTTPException(status_code=404, detail="Target categoryId not found")
+
+        # Download target image from Spaces
+        resp = requests.get(target_url)
+        if resp.status_code != 200:
+            await log_faceswap_hit(credentials.credentials, status="error")
+            raise HTTPException(status_code=404, detail="Target image not found in Spaces")
+        tgt_bytes = resp.content
 
         # Decode for processing
         src_array = np.frombuffer(src_bytes, np.uint8)
@@ -301,6 +334,45 @@ async def preview_result(result_key: str):
         headers={"Content-Disposition": "inline; filename=result.png"}
     )
 
+@fastapi_app.get("/list-images", dependencies=[Depends(verify_token)])
+def list_images(category: str = Query(None, description="Category name to list images for. If not provided, returns all available categories.")):
+    client = get_spaces_client()
+    base_prefix = "faceswap/target/"
+
+    # If no category is provided, list all categories (folders) under faceswap/target/
+    if not category:
+        # List all prefixes (categories)
+        resp = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=base_prefix, Delimiter="/")
+        categories = []
+        for prefix in resp.get("CommonPrefixes", []):
+            cat = prefix["Prefix"].split("/")[2]  # faceswap/target/<category>/
+            categories.append(cat)
+        return {"categories": categories, "note": "These are available categories. Use /list-images?category=<category> to list images."}
+
+    # If category is provided, list images in that category
+    original_prefix = f"faceswap/target/{category}/original/"
+    thumb_prefix = f"faceswap/target/{category}/thumb/"
+
+    original_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=original_prefix)
+    thumb_objects = client.list_objects_v2(Bucket=DO_SPACES_BUCKET, Prefix=thumb_prefix)
+
+    original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+    thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+
+    items = []
+    for idx, filename in enumerate(sorted(original_files), start=1):
+        if filename in thumb_files:
+            items.append({
+                "categoryId": f"{category.lower()}image_{idx}",
+                "url": f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{original_prefix}{filename}",
+                "thumburl": f"https://{DO_SPACES_BUCKET}.blr1.digitaloceanspaces.com/{thumb_prefix}{filename}"
+            })
+
+    return {
+        "categoryName": category,
+        "items": items,
+        "note": "The URLs returned are public and do not require a bearer token."
+    }
 # --------------------- Mount Gradio ---------------------
 fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 

list.py

@@ -0,0 +1,45 @@
+# from fastapi import FastAPI
+# import boto3
+# import os
+
+# app = FastAPI()
+
+# # Configure DigitalOcean Spaces (S3 Compatible)
+# session = boto3.session.Session()
+# client = session.client(
+#     's3',
+#     region_name='blr1',  # DigitalOcean region
+#     endpoint_url='https://blr1.digitaloceanspaces.com',  # endpoint
+#     aws_access_key_id=os.getenv("DO_SPACES_KEY"),
+#     aws_secret_access_key=os.getenv("DO_SPACES_SECRET")
+# )
+
+# BUCKET_NAME = "milestone"
+# BASE_URL = f"https://{BUCKET_NAME}.blr1.digitaloceanspaces.com"
+
+# @app.get("/list-images")
+# def list_images():
+#     category_name = "BirthdayParty"
+#     original_prefix = f"faceswap/target/{category_name}/original/"
+#     thumb_prefix = f"faceswap/target/{category_name}/thumb/"
+
+#     # List objects in original
+#     original_objects = client.list_objects_v2(Bucket=BUCKET_NAME, Prefix=original_prefix)
+#     thumb_objects = client.list_objects_v2(Bucket=BUCKET_NAME, Prefix=thumb_prefix)
+
+#     original_files = [obj["Key"].split("/")[-1] for obj in original_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+#     thumb_files = [obj["Key"].split("/")[-1] for obj in thumb_objects.get("Contents", []) if obj["Key"].endswith(".png")]
+
+#     items = []
+#     for idx, filename in enumerate(sorted(original_files), start=1):
+#         if filename in thumb_files:  # match only if both exist
+#             items.append({
+#                 "categoryId": f"birthdaypartyimage_{idx}",
+#                 "url": f"{BASE_URL}/{original_prefix}{filename}",
+#                 "thumburl": f"{BASE_URL}/{thumb_prefix}{filename}"
+#             })
+
+#     return {
+#         "categoryName": category_name,
+#         "items": items
+#     }