usig gridfs for storage

app.py

@@ -15,11 +15,15 @@ from fastapi.responses import RedirectResponse
 from pydantic import BaseModel
 from motor.motor_asyncio import AsyncIOMotorClient
 from bson.objectid import ObjectId
+from gridfs import AsyncIOMotorGridFSBucket
 from gradio import mount_gradio_app
 import uvicorn
 import logging
+import io
 
-# Set up logging
+# -------------------------------------------------
+# Logging
+# -------------------------------------------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -37,7 +41,7 @@ os.makedirs(RESULT_DIR, exist_ok=True)
 os.makedirs(MODELS_DIR, exist_ok=True)
 
 # -------------------------------------------------
-# Download models once
+# Download models
 # -------------------------------------------------
 def download_models():
     logger.info("Downloading models...")
@@ -67,7 +71,7 @@ def download_models():
 inswapper_path = download_models()
 
 # -------------------------------------------------
-# Initialize face analysis and swapper
+# Face Analysis + Swapper
 # -------------------------------------------------
 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 logger.info(f"Initializing FaceAnalysis with providers: {providers}")
@@ -94,7 +98,7 @@ def ensure_codeformer():
 ensure_codeformer()
 
 # -------------------------------------------------
-# MongoDB setup
+# MongoDB + GridFS
 # -------------------------------------------------
 MONGODB_URL = os.getenv(
     "MONGODB_URL",
@@ -102,10 +106,8 @@ MONGODB_URL = os.getenv(
 )
 client = AsyncIOMotorClient(MONGODB_URL)
 database = client.FaceSwap
-target_images_collection = database.get_collection("Target_Images")
-source_images_collection = database.get_collection("Source_Images")
-results_collection = database.get_collection("Results")
-logger.info("MongoDB client initialized")
+fs_bucket = AsyncIOMotorGridFSBucket(database)
+logger.info("MongoDB + GridFS initialized")
 
 # -------------------------------------------------
 # Lock for face swap
@@ -113,7 +115,7 @@ logger.info("MongoDB client initialized")
 swap_lock = threading.Lock()
 
 # -------------------------------------------------
-# Pipeline Function
+# Face Swap Pipeline
 # -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
     logger.info("Starting face swap and enhancement")
@@ -125,55 +127,39 @@ def face_swap_and_enhance(src_img, tgt_img):
             os.makedirs(RESULT_DIR, exist_ok=True)
 
             if not isinstance(src_img, np.ndarray) or not isinstance(tgt_img, np.ndarray):
-                logger.error("Invalid input images: not numpy arrays")
-                return None, None, "❌ Invalid input images: not numpy arrays"
+                return None, None, "❌ Invalid input images"
 
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
-            logger.info("Detecting faces...")
             src_faces = face_analysis_app.get(src_bgr)
             tgt_faces = face_analysis_app.get(tgt_bgr)
             if not src_faces or not tgt_faces:
-                logger.error("Face not detected in one of the images")
-                return None, None, "❌ Face not detected in one of the images"
+                return None, None, "❌ Face not detected"
 
-            unique_name = f"swapped_{uuid.uuid4().hex[:8]}.jpg"
-            swapped_path = os.path.join(UPLOAD_DIR, unique_name)
-            logger.info("Performing face swap...")
+            swapped_path = os.path.join(UPLOAD_DIR, f"swapped_{uuid.uuid4().hex[:8]}.jpg")
             swapped_bgr = swapper.get(tgt_bgr, tgt_faces[0], src_faces[0])
             if swapped_bgr is None:
-                logger.error("Face swap failed: swapper returned None")
                 return None, None, "❌ Face swap failed"
 
             cv2.imwrite(swapped_path, swapped_bgr)
-            logger.info(f"Swapped image saved to {swapped_path}")
 
             cmd = f"python {CODEFORMER_PATH} -w 0.7 --input_path {swapped_path} --output_path {RESULT_DIR} --bg_upsampler realesrgan --face_upsample"
-            logger.info(f"Running CodeFormer: {cmd}")
             result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
             if result.returncode != 0:
-                logger.error(f"CodeFormer failed: {result.stderr}")
                 return None, None, f"❌ CodeFormer failed:\n{result.stderr}"
 
             final_results_dir = os.path.join(RESULT_DIR, "final_results")
-            if not os.path.exists(final_results_dir):
-                logger.error("CodeFormer did not produce final results")
-                return None, None, "❌ CodeFormer did not produce final results"
-
             final_files = [f for f in os.listdir(final_results_dir) if f.endswith(".png")]
             if not final_files:
-                logger.error("No enhanced image found in final results")
-                return None, None, "❌ No enhanced image found in final results"
+                return None, None, "❌ No enhanced image found"
 
             final_path = os.path.join(final_results_dir, final_files[0])
             final_img = cv2.cvtColor(cv2.imread(final_path), cv2.COLOR_BGR2RGB)
-            logger.info(f"Enhanced image ready at {final_path}")
 
             return final_img, final_path, ""
 
     except Exception as e:
-        logger.error(f"Face swap error: {str(e)}")
         return None, None, f"❌ Error: {str(e)}"
 
 # -------------------------------------------------
@@ -187,7 +173,6 @@ with gr.Blocks() as demo:
         tgt_input = gr.Image(type="numpy", label="Upload Target Image")
 
     btn = gr.Button("Swap Face")
-
     output_img = gr.Image(type="numpy", label="Enhanced Output")
     download = gr.File(label="⬇️ Download Enhanced Image")
     error_box = gr.Textbox(label="Logs / Errors", interactive=False)
@@ -211,98 +196,72 @@ def root():
 async def health():
     return {"status": "healthy"}
 
+# -------- Upload Endpoints with GridFS --------
 @fastapi_app.post("/source")
 async def upload_source(image: UploadFile = File(...)):
-    logger.info(f"Uploading source image: {image.filename}")
     contents = await image.read()
-    doc = {
-        "filename": image.filename,
-        "content_type": image.content_type,
-        "data": contents
-    }
-    result = await source_images_collection.insert_one(doc)
-    logger.info(f"Source image uploaded with ID: {str(result.inserted_id)}")
-    return {"source_id": str(result.inserted_id)}
+    file_id = await fs_bucket.upload_from_stream(image.filename, contents)
+    return {"source_id": str(file_id)}
 
 @fastapi_app.post("/target")
 async def upload_target(image: UploadFile = File(...)):
-    logger.info(f"Uploading target image: {image.filename}")
     contents = await image.read()
-    doc = {
-        "filename": image.filename,
-        "content_type": image.content_type,
-        "data": contents
-    }
-    result = await target_images_collection.insert_one(doc)
-    logger.info(f"Target image uploaded with ID: {str(result.inserted_id)}")
-    return {"target_id": str(result.inserted_id)}
+    file_id = await fs_bucket.upload_from_stream(image.filename, contents)
+    return {"target_id": str(file_id)}
 
+# -------- Faceswap Endpoint --------
 class FaceSwapRequest(BaseModel):
     source_id: str
     target_id: str
 
 @fastapi_app.post("/faceswap")
 async def perform_faceswap(request: FaceSwapRequest):
-    logger.info(f"Starting face swap for source_id: {request.source_id}, target_id: {request.target_id}")
-    source_doc = await source_images_collection.find_one({"_id": ObjectId(request.source_id)})
-    if not source_doc:
-        logger.error(f"Source image not found: {request.source_id}")
-        raise HTTPException(status_code=404, detail="Source image not found")
-
-    target_doc = await target_images_collection.find_one({"_id": ObjectId(request.target_id)})
-    if not target_doc:
-        logger.error(f"Target image not found: {request.target_id}")
-        raise HTTPException(status_code=404, detail="Target image not found")
-
-    source_array = np.frombuffer(source_doc["data"], np.uint8)
-    source_bgr = cv2.imdecode(source_array, cv2.IMREAD_COLOR)
-    if source_bgr is None:
-        logger.error("Failed to decode source image")
-        raise HTTPException(status_code=500, detail="Failed to decode source image")
-    source_rgb = cv2.cvtColor(source_bgr, cv2.COLOR_BGR2RGB)
-
-    target_array = np.frombuffer(target_doc["data"], np.uint8)
-    target_bgr = cv2.imdecode(target_array, cv2.IMREAD_COLOR)
-    if target_bgr is None:
-        logger.error("Failed to decode target image")
-        raise HTTPException(status_code=500, detail="Failed to decode target image")
-    target_rgb = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2RGB)
-
-    final_img, final_path, err = face_swap_and_enhance(source_rgb, target_rgb)
-    if err:
-        logger.error(f"Face swap failed: {err}")
-        raise HTTPException(status_code=500, detail=err)
-
-    with open(final_path, "rb") as f:
-        final_bytes = f.read()
-
-    result_doc = {
-        "source_id": request.source_id,
-        "target_id": request.target_id,
-        "filename": "enhanced.png",
-        "content_type": "image/png",
-        "data": final_bytes
-    }
-    result = await results_collection.insert_one(result_doc)
-    logger.info(f"Face swap result stored with ID: {str(result.inserted_id)}")
-    return {"result_id": str(result.inserted_id)}
+    try:
+        # Read source
+        source_stream = await fs_bucket.open_download_stream(ObjectId(request.source_id))
+        source_bytes = await source_stream.read()
+        source_array = np.frombuffer(source_bytes, np.uint8)
+        source_bgr = cv2.imdecode(source_array, cv2.IMREAD_COLOR)
+        source_rgb = cv2.cvtColor(source_bgr, cv2.COLOR_BGR2RGB)
+
+        # Read target
+        target_stream = await fs_bucket.open_download_stream(ObjectId(request.target_id))
+        target_bytes = await target_stream.read()
+        target_array = np.frombuffer(target_bytes, np.uint8)
+        target_bgr = cv2.imdecode(target_array, cv2.IMREAD_COLOR)
+        target_rgb = cv2.cvtColor(target_bgr, cv2.COLOR_BGR2RGB)
+
+        # Run pipeline
+        final_img, final_path, err = face_swap_and_enhance(source_rgb, target_rgb)
+        if err:
+            raise HTTPException(status_code=500, detail=err)
+
+        # Store result in GridFS
+        with open(final_path, "rb") as f:
+            final_bytes = f.read()
+        result_id = await fs_bucket.upload_from_stream("enhanced.png", final_bytes)
+
+        return {"result_id": str(result_id)}
 
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# -------- Download Endpoint --------
 @fastapi_app.get("/download/{result_id}")
 async def download_result(result_id: str):
-    logger.info(f"Downloading result: {result_id}")
-    doc = await results_collection.find_one({"_id": ObjectId(result_id)})
-    if not doc:
-        logger.error(f"Result not found: {result_id}")
+    try:
+        stream = await fs_bucket.open_download_stream(ObjectId(result_id))
+        file_data = await stream.read()
+        return Response(
+            content=file_data,
+            media_type="image/png",
+            headers={"Content-Disposition": f"attachment; filename=enhanced.png"}
+        )
+    except Exception:
         raise HTTPException(status_code=404, detail="Result not found")
-    return Response(
-        content=doc["data"],
-        media_type=doc["content_type"],
-        headers={"Content-Disposition": f"attachment; filename={doc['filename']}"}
-    )
 
-# Mount Gradio at /gradio
+# Mount Gradio
 fastapi_app = mount_gradio_app(fastapi_app, demo, path="/gradio")
 
-# Run the app with Uvicorn
 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)

test.py

@@ -1,33 +1,54 @@
-# import requests
+import requests
 
-# response = requests.get("https://logicgoinfotechspaces-faceswap.hf.space/health")
-# print(response.status_code)  # Should be 200
-# print(response.json())  # Should print: {'status': 'healthy'}
 
-# import requests
+#HEALTH CHECK
+response = requests.get("https://logicgoinfotechspaces-faceswap.hf.space/health")
+print(response.status_code)  # Should be 200
+print(response.json())  # Should print: {'status': 'healthy'}
 
-# with open("./source.jpeg", "rb") as f:
-#     files = {"image": ("./source.jpeg", f, "image/jpeg")}
-#     response = requests.post("https://logicgoinfotechspaces-faceswap.hf.space/source", files=files)
-#     print(response.status_code)  # Should be 200
-#     print(response.json())  # Should print: {'source_id': '<id>'}
-#     source_id = response.json().get("source_id")
+#----------------------------------------------------------------------------------------------------#
 
-# import requests
+#UPLOAD SOURCE IMAGE
+with open("./source.jpeg", "rb") as f:
+    files = {"image": ("./source.jpeg", f, "image/jpeg")}
+    response = requests.post("https://logicgoinfotechspaces-faceswap.hf.space/source", files=files)
+    print(response.status_code)  # Should be 200
+    print(response.json())  # Should print: {'source_id': '<id>'}
+    source_id = response.json().get("source_id")
 
-# with open("./target.jpeg", "rb") as f:
-#     files = {"image": ("./target.jpeg", f, "image/jpeg")}
-#     response = requests.post("https://logicgoinfotechspaces-faceswap.hf.space/target", files=files)
-#     print(response.status_code)  # Should be 200
-#     print(response.json())  # Should print: {'target_id': '<id>'}
-#     target_id = response.json().get("target_id")
+#----------------------------------------------------------------------------------------------------# 
 
-import requests
+#UPLOAD TARGET IMAGE 
+with open("./target.jpeg", "rb") as f:
+    files = {"image": ("./target.jpeg", f, "image/jpeg")}
+    response = requests.post("https://logicgoinfotechspaces-faceswap.hf.space/target", files=files)
+    print(response.status_code)  # Should be 200
+    print(response.json())  # Should print: {'target_id': '<id>'}
+    target_id = response.json().get("target_id")
 
+#----------------------------------------------------------------------------------------------------# 
+
+#SWAP IMAGES
 source_id = "68bff17eb3b2e06c24cafa22"  # Replace with actual ID
 target_id = "68bff1a1b3b2e06c24cafa23"  # Replace with actual ID
 payload = {"source_id": source_id, "target_id": target_id}
 response = requests.post("https://logicgoinfotechspaces-faceswap.hf.space/faceswap", json=payload)
 print(response.status_code)  # Should be 200
 print(response.json())  # Should print: {'result_id': '<id>'}
-result_id = response.json().get("result_id")
+result_id = response.json().get("result_id")
+
+#----------------------------------------------------------------------------------------------------# 
+
+
+#DOWNLOAD IMAGES
+result_id = "68bff4800e80c0bb228b5962"  # Replace with actual ID
+response = requests.get(f"https://logicgoinfotechspaces-faceswap.hf.space/download/{result_id}")
+print(response.status_code)  # Should be 200
+if response.status_code == 200:
+    with open("result.png", "wb") as f:
+        f.write(response.content)
+    print("Image downloaded as result.png")
+else:
+    print(response.json())  # Print error details
+
+#----------------------------------------------------------------------------------------------------#