Update main.py

main.py

@@ -6,85 +6,65 @@ import numpy as np
 from PIL import Image, ImageEnhance, ImageFilter
 from rembg import remove, new_session
 import mediapipe as mp
+from mediapipe.tasks import python
+from mediapipe.tasks.vision import FaceDetector, FaceDetectorOptions, RunningMode
 
-app = FastAPI(title="Pro Passport API")
+app = FastAPI()
 
-# Global variables
+# Global models
 session = None
-face_mesh = None
+detector = None
 
 @app.on_event("startup")
 async def startup_event():
-    global session, face_mesh
+    global session, detector
     print("Loading AI Models...")
-    # Use the portrait-specific model for higher precision
+    # Use rembg with BiRefNet (best for portraits)
     session = new_session("birefnet-portrait")
-    face_mesh = mp.solutions.face_mesh.FaceMesh(
-        static_image_mode=True, 
-        refine_landmarks=True, 
-        max_num_faces=1
-    )
+    
+    # Initialize MediaPipe Tasks API (The only stable way)
+    # Note: You must ensure 'detector.tflite' exists in your /app directory
+    # If not present, download it from Google's official site
+    base_options = python.BaseOptions(model_asset_path='detector.tflite')
+    options = FaceDetectorOptions(base_options=base_options, running_mode=RunningMode.IMAGE)
+    detector = mp.tasks.vision.FaceDetector.create_from_options(options)
     print("AI Models Loaded.")
 
 @app.post("/generate")
-async def generate_passport(
-    file: UploadFile = File(...),
-    suit_mode: bool = Form(True)
-):
-    # 1. Load image
+async def generate_passport(file: UploadFile = File(...)):
+    if detector is None:
+        raise HTTPException(status_code=500, detail="Models not loaded")
+
     contents = await file.read()
     pil_img = Image.open(io.BytesIO(contents)).convert("RGBA")
     np_img = np.array(pil_img.convert("RGB"))
-    h, w, _ = np_img.shape
-
-    # 2. Precise Landmark Analysis
-    results = face_mesh.process(cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))
-    if not results.multi_face_landmarks:
-        raise HTTPException(status_code=400, detail="No face detected")
     
-    landmarks = results.multi_face_landmarks[0].landmark
-    # 152 = Chin, 10 = Top of head
-    chin_y = int(landmarks[152].y * h)
-    crown_y = int(landmarks[10].y * h)
-    head_h = chin_y - crown_y
+    # 1. Detection via Tasks API
+    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_img)
+    results = detector.detect(mp_image)
     
-    # 3. Compliance Crop (The "Passport Math")
-    # US standard: Chin to top of head is 50-69% of image height.
-    # We aim for ~60%.
-    margin_top = int(head_h * 0.8)  # Space above head
-    margin_bottom = int(head_h * 1.5) # Space below chin for shoulders
+    if not results.detections:
+        raise HTTPException(status_code=400, detail="No face detected")
     
-    y_start = max(0, crown_y - margin_top)
-    y_end = min(h, chin_y + margin_bottom)
+    # 2. Extract Box
+    bbox = results.detections[0].bounding_box
+    h, w = np_img.shape[:2]
     
-    # Crop the person
+    # Simple Crop logic based on bounding box
+    y_start = max(0, int(bbox.origin_y - bbox.height * 0.5))
+    y_end = min(h, int(bbox.origin_y + bbox.height * 1.5))
     crop = pil_img.crop((0, y_start, w, y_end))
     
-    # 4. Professional Background Removal
+    # 3. Background Removal
     no_bg = remove(crop, session=session, alpha_matting=True)
     
-    # 5. Composite White BG
+    # 4. White BG + Resize
     final = Image.new("RGBA", (600, 600), (255, 255, 255, 255))
-    # Resize keeping aspect ratio
     no_bg.thumbnail((500, 500), Image.Resampling.LANCZOS)
-    paste_x = (600 - no_bg.width) // 2
-    paste_y = (600 - no_bg.height) + 20 # Offset to center head
-    final.paste(no_bg, (paste_x, paste_y), no_bg)
-    
-    # 6. Formal Suit Enhancer (Cutout.pro style)
-    if suit_mode:
-        # Create a "formal" filter for the lower half
-        lower_half = final.crop((0, 300, 600, 600))
-        lower_half = lower_half.filter(ImageFilter.GaussianBlur(1))
-        enhancer = ImageEnhance.Contrast(lower_half)
-        lower_half = enhancer.enhance(1.2)
-        final.paste(lower_half, (0, 300))
-    
-    # 7. Final Polish (DSLR Look)
-    final = ImageEnhance.Sharpness(final).enhance(1.3)
-    final = ImageEnhance.Color(final).enhance(1.1)
+    final.paste(no_bg, ((600-no_bg.width)//2, (600-no_bg.height)+20), no_bg)
     
+    # 5. Export
     buf = io.BytesIO()
-    final.convert("RGB").save(buf, format="JPEG", quality=95, optimize=True)
+    final.convert("RGB").save(buf, format="JPEG", quality=95)
     buf.seek(0)
     return StreamingResponse(buf, media_type="image/jpeg")