Update app.py

app.py

@@ -4,69 +4,112 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 import cv2
 import numpy as np
-from huggingface_hub import hf_hub_download
-from realesrgan import RealESRGAN
+from realesrgan import RealESRGANer
+from basicsr.archs.rrdbnet_arch import RRDBNet
 from mediapipe import solutions as mp_solutions
 
-# ✅ Cache dir for Hugging Face
-CACHE_DIR = "/tmp/hf_cache"
-os.environ["HF_HOME"] = CACHE_DIR
-os.makedirs(CACHE_DIR, exist_ok=True)
-
+# Initialize FastAPI app
 app = FastAPI(title="Face Beautification API")
 
-# ✅ Load Mediapipe face detector
-mp_face = mp_solutions.face_detection.FaceDetection(
-    model_selection=1, min_detection_confidence=0.5
-)
+# --- Model Loading ---
+
+# Load Mediapipe face detector
+try:
+    mp_face = mp_solutions.face_detection.FaceDetection(
+        model_selection=1, min_detection_confidence=0.5)
+except Exception as e:
+    print(f"Error loading Mediapipe face detector: {e}")
+    mp_face = None
+
+# Correctly load the Real-ESRGAN model for CPU inference
+# The model weights are pre-downloaded in the Dockerfile to the 'weights' directory
+model_path = os.path.join("weights", "RealESRGAN_x4plus.pth")
+upsampler = None
+
+if os.path.exists(model_path) and mp_face:
+    try:
+        # Define the model architecture
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        
+        # Initialize the RealESRGANer
+        # `half=False` is important for CPU-only execution
+        upsampler = RealESRGANer(
+            scale=4,
+            model_path=model_path,
+            dni_weight=None,
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=False,
+            gpu_id=None, # Explicitly set to None for CPU
+        )
+    except Exception as e:
+        print(f"Error loading Real-ESRGAN model: {e}")
+        upsampler = None
+else:
+    print("Model weights not found or face detector failed to load.")
 
-# ✅ Download ESRGAN weights from Hugging Face (free)
-model_path = hf_hub_download(
-    repo_id="eugenesiow/real-esrgan",
-    filename="RealESRGAN_x4plus.pth",
-    cache_dir=CACHE_DIR
-)
 
-# ✅ Load ESRGAN model
-device = "cuda" if cv2.cuda.getCudaEnabledDeviceCount() > 0 else "cpu"
-model = RealESRGAN(device, scale=4)
-model.load_weights(model_path)
+# --- API Endpoints ---
 
 @app.get("/")
 async def root():
+    """Root endpoint to check if the API is running."""
     return {"message": "Free Face Beautification API is running!"}
 
 @app.post("/beautify")
 async def beautify(image: UploadFile = File(...)):
+    """
+    Receives an image, detects faces, enhances them, and returns the result.
+    """
+    if not upsampler or not mp_face:
+        return JSONResponse({"error": "Model not loaded, API is not operational."}, status_code=503)
+
     try:
-        # Read image
+        # 1. Read and decode the uploaded image
         contents = await image.read()
         npimg = np.frombuffer(contents, np.uint8)
         img = cv2.imdecode(npimg, cv2.IMREAD_COLOR)
 
-        # Detect faces
+        if img is None:
+            return JSONResponse({"error": "Invalid image file."}, status_code=400)
+
+        # 2. Detect faces using Mediapipe
         results = mp_face.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
         if not results.detections:
-            return JSONResponse({"error": "No face detected"}, status_code=400)
+            return JSONResponse({"error": "No face detected in the image."}, status_code=400)
 
+        # 3. Process each detected face
         for detection in results.detections:
             bbox = detection.location_data.relative_bounding_box
             h, w, _ = img.shape
-            x1 = int(bbox.xmin * w)
-            y1 = int(bbox.ymin * h)
-            x2 = int((bbox.xmin + bbox.width) * w)
-            y2 = int((bbox.ymin + bbox.height) * h)
+            
+            # Ensure coordinates are within image bounds
+            x1 = max(0, int(bbox.xmin * w))
+            y1 = max(0, int(bbox.ymin * h))
+            x2 = min(w, int((bbox.xmin + bbox.width) * w))
+            y2 = min(h, int((bbox.ymin + bbox.height) * h))
+
+            # Crop the face
             face = img[y1:y2, x1:x2]
 
-            # Beautify (enhance + smooth)
-            face_upscaled = model.predict(face)
-            face_smooth = cv2.bilateralFilter(face_upscaled, 9, 75, 75)
+            if face.size == 0:
+                continue
+            
+            # 4. Enhance the face using Real-ESRGAN
+            # The `enhance` method returns the upscaled image and its mode
+            face_upscaled, _ = upsampler.enhance(face)
 
-            # Blend face back
-            face_smooth = cv2.resize(face_smooth, (x2 - x1, y2 - y1))
-            img[y1:y2, x1:x2] = face_smooth
+            # 5. Apply a smoothing filter for the "beautification" effect
+            face_smooth = cv2.bilateralFilter(face_upscaled, d=9, sigmaColor=75, sigmaSpace=75)
+            
+            # 6. Resize the enhanced face back to its original dimensions and blend it in
+            face_smooth_resized = cv2.resize(face_smooth, (x2 - x1, y2 - y1))
+            img[y1:y2, x1:x2] = face_smooth_resized
 
-        # Encode final image as Base64
+        # 7. Encode the final image to Base64 to send in the JSON response
         _, buffer = cv2.imencode(".jpg", img)
         img_base64 = base64.b64encode(buffer).decode("utf-8")
 
@@ -75,9 +118,11 @@ async def beautify(image: UploadFile = File(...)):
             "message": "Beautification complete!",
             "image_base64": img_base64
         })
+
     except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
+        return JSONResponse({"error": f"An unexpected error occurred: {str(e)}"}, status_code=500)
 
 @app.get("/health")
 async def health():
-    return {"ready": True}
+    """Health check endpoint."""
+    return {"ready": bool(upsampler and mp_face)}