commit

app.py

@@ -17,6 +17,11 @@ from motor.motor_asyncio import AsyncIOMotorClient
 from bson.objectid import ObjectId
 from gradio import mount_gradio_app
 import uvicorn
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # -------------------------------------------------
 # Paths
@@ -35,6 +40,7 @@ os.makedirs(MODELS_DIR, exist_ok=True)
 # Download models once
 # -------------------------------------------------
 def download_models():
+    logger.info("Downloading models...")
     inswapper_path = hf_hub_download(
         repo_id=REPO_ID,
         filename="models/inswapper_128.onnx",
@@ -55,6 +61,7 @@ def download_models():
             repo_type="model",
             local_dir=MODELS_DIR
         )
+    logger.info("Models downloaded successfully")
     return inswapper_path
 
 inswapper_path = download_models()
@@ -62,10 +69,12 @@ inswapper_path = download_models()
 # -------------------------------------------------
 # Initialize face analysis and swapper
 # -------------------------------------------------
-providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']  # Prefer CUDA
+providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+logger.info(f"Initializing FaceAnalysis with providers: {providers}")
 app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=providers)
 app.prepare(ctx_id=0, det_size=(640, 640))
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=providers)
+logger.info("FaceAnalysis and swapper initialized")
 
 # -------------------------------------------------
 # CodeFormer setup
@@ -74,11 +83,13 @@ CODEFORMER_PATH = "CodeFormer/inference_codeformer.py"
 
 def ensure_codeformer():
     if not os.path.exists("CodeFormer"):
+        logger.info("Cloning CodeFormer repository...")
         subprocess.run("git clone https://github.com/sczhou/CodeFormer.git", shell=True, check=True)
         subprocess.run("pip install -r CodeFormer/requirements.txt", shell=True, check=True)
         subprocess.run("python CodeFormer/basicsr/setup.py develop", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py facelib", shell=True, check=True)
         subprocess.run("python CodeFormer/scripts/download_pretrained_models.py CodeFormer", shell=True, check=True)
+        logger.info("CodeFormer setup complete")
 
 ensure_codeformer()
 
@@ -94,6 +105,7 @@ 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")
 
 # -------------------------------------------------
 # Lock for face swap
@@ -104,6 +116,7 @@ swap_lock = threading.Lock()
 # Pipeline Function
 # -------------------------------------------------
 def face_swap_and_enhance(src_img, tgt_img):
+    logger.info("Starting face swap and enhancement")
     try:
         with swap_lock:
             shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
@@ -111,38 +124,56 @@ def face_swap_and_enhance(src_img, tgt_img):
             os.makedirs(UPLOAD_DIR, exist_ok=True)
             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"
+
             src_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
             tgt_bgr = cv2.cvtColor(tgt_img, cv2.COLOR_RGB2BGR)
 
+            logger.info("Detecting faces...")
             src_faces = app.get(src_bgr)
             tgt_faces = app.get(tgt_bgr)
             if not src_faces or not tgt_faces:
-                return None, None, "❌ Face not detected in one of the images."
+                logger.error("Face not detected in one of the images")
+                return None, None, "❌ Face not detected in one of the images"
 
             unique_name = f"swapped_{uuid.uuid4().hex[:8]}.jpg"
             swapped_path = os.path.join(UPLOAD_DIR, unique_name)
+            logger.info("Performing face swap...")
             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):
-                return None, None, "❌ CodeFormer did not produce final results."
+                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:
-                return None, None, "❌ No enhanced image found in final results."
+                logger.error("No enhanced image found in final results")
+                return None, None, "❌ No enhanced image found in final results"
 
             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)}"
 
 # -------------------------------------------------
@@ -182,6 +213,7 @@ async def health():
 
 @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,
@@ -189,10 +221,12 @@ async def upload_source(image: UploadFile = File(...)):
         "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)}
 
 @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,
@@ -200,6 +234,7 @@ async def upload_target(image: UploadFile = File(...)):
         "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)}
 
 class FaceSwapRequest(BaseModel):
@@ -208,24 +243,34 @@ class FaceSwapRequest(BaseModel):
 
 @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:
@@ -239,12 +284,15 @@ async def perform_faceswap(request: FaceSwapRequest):
         "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)}
 
 @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}")
         raise HTTPException(status_code=404, detail="Result not found")
     return Response(
         content=doc["data"],