feat: image gen

app/util/passport_photo_engine/haar_face_detector.py

@@ -0,0 +1,10 @@
+import cv2
+
+
+class HaarFaceDetector:
+    def __init__(self):
+        self.detector = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")
+    def detect(self, img_rgb):
+        gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
+        faces = self.detector.detectMultiScale(gray, 1.1, 5, minSize=(80,80))
+        return max(faces, key=lambda f: f[2]*f[3]) if len(faces) > 0 else None

app/util/passport_photo_engine/manual_face_extractor.py

@@ -0,0 +1,13 @@
+import numpy as np
+import math
+
+class ManualFaceExtractor:
+    def extract_face(self, img_rgb, mask):
+        if mask is None: return {"chin_angle": 0.0}
+        y_idxs, x_idxs = np.where(mask > 30)
+        if len(y_idxs) < 100: return {"chin_angle": 0.0}
+        sort_idx = np.argsort(y_idxs)[-int(len(y_idxs)*0.2):] # Bottom 20%
+        try:
+            m, c = np.polyfit(x_idxs[sort_idx], y_idxs[sort_idx], 1)
+            return {"chin_angle": float(np.clip(math.degrees(math.atan(m)), -30, 30))}
+        except: return {"chin_angle": 0.0}

app/util/passport_photo_engine/passport_cropper.py

@@ -0,0 +1,57 @@
+import cv2
+import numpy as np
+
+class PassportCropper:
+    def __init__(self, output_size=(600,800), bg_color=(255,255,255)):
+        self.out_w, self.out_h = output_size
+        self.bg_color = tuple(int(x) for x in bg_color)
+        self.target_aspect = self.out_w / self.out_h
+
+    def composite(self, img_rgb, mask):
+        bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+        m = cv2.GaussianBlur(mask, (5,5), 0)
+        alpha = (m.astype(float) / 255.0)[:,:,None]
+        bg = np.full_like(bgr, np.array(self.bg_color, dtype=np.uint8))
+        return (bgr.astype(float) * alpha + bg.astype(float) * (1.0 - alpha)).astype(np.uint8)
+
+    def rotate_and_expand_face(self, img_bgr, angle_deg, raw_face_box):
+        h0, w0 = img_bgr.shape[:2]
+        x1, y1, x2, y2 = raw_face_box
+        # Expansion Logic: 0.4 sides, 0.6 top, 1.2 bottom
+        fw, fh = x2-x1, y2-y1
+        hx1, hy1 = max(0, x1-int(fw*0.4)), max(0, y1-int(fh*0.6))
+        hx2, hy2 = min(w0-1, x2+int(fw*0.4)), min(h0-1, y2+int(fh*1.2))
+        
+        # Rotation
+        M = cv2.getRotationMatrix2D((w0/2, h0/2), -angle_deg, 1.0)
+        cos, sin = np.abs(M[0,0]), np.abs(M[0,1])
+        nW, nH = int((h0*sin)+(w0*cos)), int((h0*cos)+(w0*sin))
+        M[0,2] += (nW/2) - w0/2; M[1,2] += (nH/2) - h0/2
+        
+        rot_img = cv2.warpAffine(img_bgr, M, (nW, nH), borderValue=self.bg_color)
+        
+        # Rotate Box Points
+        pts = np.array([[hx1,hy1,1],[hx2,hy1,1],[hx2,hy2,1],[hx1,hy2,1]]).T
+        rot_pts = M @ pts
+        rx, ry = rot_pts[0,:], rot_pts[1,:]
+        return rot_img, (int(rx.min()), int(ry.min()), int(rx.max()), int(ry.max()))
+
+    def crop_to_ratio(self, img, box):
+        bx1, by1, bx2, by2 = box
+        bw, bh = bx2-bx1, by2-by1
+        if bw/bh > self.target_aspect: # Too wide
+            new_h = int(bw/self.target_aspect)
+            by1 -= (new_h - bh)//2; by2 = by1 + new_h
+        else: # Too tall
+            new_w = int(bh*self.target_aspect)
+            bx1 -= (new_w - bw)//2; bx2 = bx1 + new_w
+        
+        # Canvas Crop
+        H, W = img.shape[:2]
+        canvas = np.full((by2-by1, bx2-bx1, 3), self.bg_color, dtype=np.uint8)
+        sx1, sy1 = max(0, bx1), max(0, by1)
+        sx2, sy2 = min(W, bx2), min(H, by2)
+        dx1, dy1 = max(0, sx1-bx1), max(0, sy1-by1)
+        if sx2>sx1 and sy2>sy1: canvas[dy1:dy1+(sy2-sy1), dx1:dx1+(sx2-sx1)] = img[sy1:sy2, sx1:sx2]
+        return cv2.resize(canvas, (self.out_w, self.out_h), interpolation=cv2.INTER_AREA)
+

app/util/passport_photo_engine/segmenter_rmbg.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+from transformers import AutoModelForImageSegmentation
+
+class SegmenterRMBG:
+    def __init__(self, device=None, model_input_size=(1024,1024)):
+        self.device = torch.device(device) if device else torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True).to(self.device)
+        self.modenl.eval()
+        self.model_input_size = list(model_input_size)
+
+    def _preprocess(self, img_np):
+        t = torch.tensor(img_np, dtype=torch.float32).permute(2,0,1).unsqueeze(0)
+        t = F.interpolate(t, size=self.model_input_size, mode="bilinear") / 255.0
+        return normalize(t, [0.5]*3, [1.0]*3).to(self.device)
+
+    def _postprocess(self, result, orig_size):
+        H, W = orig_size
+        if result.dim() == 4: result = result
+        elif result.dim() == 3: result = result.unsqueeze(0)
+        result = F.interpolate(result, size=(H, W), mode="bilinear").squeeze(0).squeeze(0)
+        r_min, r_max = result.min(), result.max()
+        result = (result - r_min) / (r_max - r_min + 1e-8)
+        # Gamma correction for hair
+        result = torch.pow(result, 2.5)
+        result[result < 0.05] = 0
+        return (result * 255).cpu().numpy().astype(np.uint8)
+
+    def segment(self, img_rgb):
+        inp = self._preprocess(img_rgb)
+        with torch.no_grad(): out = self.model(inp)
+        if isinstance(out, (list, tuple)): out = out[0]
+        return self._postprocess(out[0], img_rgb.shape[:2])

requirements.txt

@@ -17,4 +17,13 @@ pandas
 SQLAlchemy
 psycopg2-binary
 
-boto3
+boto3
+
+numpy
+opencv-python-headless
+pillow
+torch
+torchvision
+transformers
+huggingface_hub
+accelerate

server.py

@@ -3,23 +3,44 @@ import tempfile
 os.environ["PLAYWRIGHT_BROWSERS_PATH"] = "/home/user/.cache/ms-playwright"
 
 import logging
-from flask import Flask, request, jsonify, send_file
 from dotenv import load_dotenv
+import io
+import cv2
+import numpy as np
+from PIL import Image
 import json
 import requests
 import uuid
 import importlib
 import io
+from flask import Flask, request, jsonify, send_file
+
 
 from app.util.gen_ai_base import GenAIBaseClient
 from app.util.browser_agent import BrowserAgent
-# from app.util.japan_multientry_visa_letter_generator   import JapanMultiEntryVisaLetterGenerator
+from app.util.passport_photo_engine.haar_face_detector import HaarFaceDetector
+from app.util.passport_photo_engine.manual_face_extractor import ManualFaceExtractor
+from app.util.passport_photo_engine.passport_cropper import PassportCropper
+from app.util.passport_photo_engine.segmenter_rmbg import SegmenterRMBG
 from app.util.parameter_utils import init_secret
 import sys
 sys.stdout.reconfigure(line_buffering=True)
 API = "https://api-dev.spun.global"
 
-
+print("--- Loading Passport AI Models (This happens once) ---")
+passport_models = {
+    "segmenter": SegmenterRMBG(),     # Heavy model (GPU/CPU)
+    "detector": HaarFaceDetector(),   # Fast model
+    "extractor": ManualFaceExtractor()
+}
+print("--- Passport Models Ready ---")
+
+PASSPORT_COLORS = {
+    "white": (255, 255, 255),
+    "id_red": (0, 0, 219),
+    "id_blue": (219, 0, 0),
+    "light_blue": (235, 206, 135)
+}
 
 def create_app() -> Flask:
     load_dotenv() 
@@ -206,6 +227,53 @@ def create_app() -> Flask:
         except Exception as e:
             print(f"Error in /generate/{visa_type}: {e}")
             return jsonify({"error": str(e)}), 500
+        
+    @app.route("/generate-passport-photo", methods=["POST"])
+    def generate_passport():
+        data = request.get_json()
+        bg_color_name = data.get('bg_color_name', 'white')
+        response = requests.get(data['raw_photo'], stream=True)
+        if response.status_code != 200:
+            return jsonify({"error": f"Failed to download image from S3. Status: {response.status_code}"}), 400
+        try:
+            # Read image
+            in_memory_file = io.BytesIO(response.content)
+            pil_image = Image.open(in_memory_file).convert("RGB")
+            img_rgb = np.array(pil_image)
+
+            # Get models
+            seg = passport_models["segmenter"]
+            det = passport_models["detector"]
+            ext = passport_models["extractor"]
+
+            # 1. Segment
+            mask = seg.segment(img_rgb)
+
+            # 2. Detect
+            face_rect = det.detect(img_rgb)
+            if face_rect is None:
+                return jsonify({"error": "No face detected"}), 400
+            x, y, w, h = face_rect
+
+            # 3. Angle
+            info = ext.extract_face(img_rgb, mask)
+            angle = info.get("chin_angle", 0.0)
+
+            # 4. Process
+            selected_bg = PASSPORT_COLORS.get(bg_color_name, (255, 255, 255))
+            cropper = PassportCropper(output_size=(600, 800), bg_color=selected_bg)
+            
+            img_clean = cropper.composite(img_rgb, mask)
+            img_rot, rot_box = cropper.rotate_and_expand_face(img_clean, angle, (x,y,x+w,y+h))
+            final_passport = cropper.crop_to_ratio(img_rot, rot_box)
+
+            # Return result
+            is_success, buffer = cv2.imencode(".jpg", final_passport)
+            return send_file(io.BytesIO(buffer), mimetype='image/jpeg')
+
+        except Exception as e:
+            print(f"Passport Error: {e}")
+            return jsonify({"error": str(e)}), 500
 
     @app.route('/', methods=['GET'])
     def hello_world():