From 519541867e296be86b93e31a19c5731eb0013898 Mon Sep 17 00:00:00 2001 From: stirling-image Date: Tue, 14 Apr 2026 16:18:17 +0800 Subject: [PATCH] fix(ai): support both old and new mediapipe APIs for airgapped Docker (#69) MediaPipe >= 0.10.30 removed the mp.solutions namespace. This broke face blur, face enhance, red-eye removal, and photo restoration for users running newer mediapipe versions (closes #43). All 5 Python scripts that use mediapipe now try the legacy mp.solutions API first and fall back to the new mp.tasks API on AttributeError. Model files (blaze_face_short_range.task, face_landmarker.task) are pre-downloaded during Docker build into /opt/models/mediapipe/ so the image works fully airgapped. Local dev auto-downloads to .models/. Co-authored-by: stirling-image --- docker/download_models.py | 75 ++++++++++++++++++++++++--- packages/ai/python/detect_faces.py | 19 ++++--- packages/ai/python/enhance_faces.py | 19 ++++--- packages/ai/python/face_landmarks.py | 5 +- packages/ai/python/red_eye_removal.py | 19 ++++--- packages/ai/python/restore.py | 19 ++++--- 6 files changed, 117 insertions(+), 39 deletions(-) diff --git a/docker/download_models.py b/docker/download_models.py index cfa2af6c..a5fbcee5 100644 --- a/docker/download_models.py +++ b/docker/download_models.py @@ -65,6 +65,14 @@ NAFNET_MODEL_URL = ( NAFNET_MODEL_PATH = os.path.join(NAFNET_MODEL_DIR, "NAFNet-SIDD-width64.pth") NAFNET_MIN_SIZE = 60_000_000 # ~67 MB +MEDIAPIPE_MODEL_DIR = "/opt/models/mediapipe" +FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task" +FACE_DETECT_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "blaze_face_short_range.task") +FACE_DETECT_MIN_SIZE = 100_000 # ~200 KB +FACE_LANDMARKER_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" +FACE_LANDMARKER_MODEL_PATH = os.path.join(MEDIAPIPE_MODEL_DIR, "face_landmarker.task") +FACE_LANDMARKER_MIN_SIZE = 1_000_000 # ~7 MB + REMBG_MODELS = [ "u2net", "isnet-general-use", @@ -310,19 +318,60 @@ def download_nafnet_model(): print(f" NAFNet model downloaded: {size:,} bytes") +def download_mediapipe_task_models(): + """Download MediaPipe tasks API model files for face detection and landmarks. + + These models are used by the mp.tasks fallback when mp.solutions is + unavailable (mediapipe >= 0.10.30). Pre-downloading ensures the Docker + image works fully airgapped. + """ + print("=== Downloading MediaPipe task models ===") + os.makedirs(MEDIAPIPE_MODEL_DIR, exist_ok=True) + + for url, path, name, min_size in [ + (FACE_DETECT_MODEL_URL, FACE_DETECT_MODEL_PATH, + "blaze_face_short_range", FACE_DETECT_MIN_SIZE), + (FACE_LANDMARKER_MODEL_URL, FACE_LANDMARKER_MODEL_PATH, + "face_landmarker", FACE_LANDMARKER_MIN_SIZE), + ]: + print(f" Downloading {name}...") + urllib.request.urlretrieve(url, path) + size = os.path.getsize(path) + assert size > min_size, ( + f"{name} model too small: {size} bytes (expected > {min_size})" + ) + print(f" {name} downloaded ({size / 1_000_000:.1f} MB)") + print("MediaPipe task models downloaded.\n") + + def verify_mediapipe(): """Verify MediaPipe face detection models are bundled in the wheel.""" print("=== Verifying MediaPipe models ===") import mediapipe as mp - for selection in [0, 1]: - label = "short-range" if selection == 0 else "full-range" - print(f" Verifying {label} model (selection={selection})...") - detector = mp.solutions.face_detection.FaceDetection( - model_selection=selection, min_detection_confidence=0.5 + try: + for selection in [0, 1]: + label = "short-range" if selection == 0 else "full-range" + print(f" Verifying {label} model (selection={selection})...") + detector = mp.solutions.face_detection.FaceDetection( + model_selection=selection, min_detection_confidence=0.5 + ) + detector.close() + print(f" {label} model OK") + except AttributeError: + # mediapipe >= 0.10.30 removed mp.solutions; verify tasks API instead + print(" mp.solutions unavailable, verifying mp.tasks API...") + options = mp.tasks.vision.FaceDetectorOptions( + base_options=mp.tasks.BaseOptions( + model_asset_path=FACE_DETECT_MODEL_PATH + ), + running_mode=mp.tasks.vision.RunningMode.IMAGE, + min_detection_confidence=0.5, ) + detector = mp.tasks.vision.FaceDetector.create_from_options(options) detector.close() - print(f" {label} model OK") + print(" mp.tasks FaceDetector OK") + print("MediaPipe models verified.\n") @@ -423,6 +472,19 @@ def smoke_test(): assert os.path.isdir(vl_dir), f"PaddleOCR-VL model missing: {vl_dir}" print(" PaddleOCR-VL model verified") + # MediaPipe task models must exist (for mp.tasks fallback) + assert os.path.exists(FACE_DETECT_MODEL_PATH), ( + f"MediaPipe face detector model missing: {FACE_DETECT_MODEL_PATH}" + ) + assert os.path.getsize(FACE_DETECT_MODEL_PATH) > FACE_DETECT_MIN_SIZE + print(" MediaPipe face detector model verified") + + assert os.path.exists(FACE_LANDMARKER_MODEL_PATH), ( + f"MediaPipe face landmarker model missing: {FACE_LANDMARKER_MODEL_PATH}" + ) + assert os.path.getsize(FACE_LANDMARKER_MODEL_PATH) > FACE_LANDMARKER_MIN_SIZE + print(" MediaPipe face landmarker model verified") + print("Smoke test passed.\n") @@ -439,6 +501,7 @@ def main(): download_paddleocr_vl_model() download_scunet_model() download_nafnet_model() + download_mediapipe_task_models() verify_mediapipe() smoke_test() print("All models downloaded and verified.") diff --git a/packages/ai/python/detect_faces.py b/packages/ai/python/detect_faces.py index bb3a9c0b..6cf36244 100644 --- a/packages/ai/python/detect_faces.py +++ b/packages/ai/python/detect_faces.py @@ -12,19 +12,22 @@ def emit_progress(percent, stage): # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task" -_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") -_FACE_DETECT_MODEL_PATH = os.path.join(_MODEL_DIR, "blaze_face_short_range.task") +_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.task" +_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") +_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.task") def _ensure_face_detect_model(): - """Download the face detector model if not present.""" - if os.path.exists(_FACE_DETECT_MODEL_PATH): - return _FACE_DETECT_MODEL_PATH - os.makedirs(_MODEL_DIR, exist_ok=True) + """Resolve face detector model. Docker path first, then local dev.""" + if os.path.exists(_DOCKER_MODEL_PATH): + return _DOCKER_MODEL_PATH + if os.path.exists(_LOCAL_MODEL_PATH): + return _LOCAL_MODEL_PATH + os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True) import urllib.request emit_progress(15, "Downloading face detection model") - urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH) - return _FACE_DETECT_MODEL_PATH + urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH) + return _LOCAL_MODEL_PATH def _detect_with_solutions(img_array, min_confidence): diff --git a/packages/ai/python/enhance_faces.py b/packages/ai/python/enhance_faces.py index aa0de284..92a6caef 100644 --- a/packages/ai/python/enhance_faces.py +++ b/packages/ai/python/enhance_faces.py @@ -40,19 +40,22 @@ CODEFORMER_MODEL_PATH = os.environ.get( # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task" -_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") -_FACE_DETECT_MODEL_PATH = os.path.join(_MODEL_DIR, "blaze_face_short_range.task") +_DOCKER_MODEL_PATH = "/opt/models/mediapipe/blaze_face_short_range.task" +_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") +_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "blaze_face_short_range.task") def _ensure_face_detect_model(): - """Download the face detector model if not present.""" - if os.path.exists(_FACE_DETECT_MODEL_PATH): - return _FACE_DETECT_MODEL_PATH - os.makedirs(_MODEL_DIR, exist_ok=True) + """Resolve face detector model. Docker path first, then local dev.""" + if os.path.exists(_DOCKER_MODEL_PATH): + return _DOCKER_MODEL_PATH + if os.path.exists(_LOCAL_MODEL_PATH): + return _LOCAL_MODEL_PATH + os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True) import urllib.request emit_progress(15, "Downloading face detection model") - urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH) - return _FACE_DETECT_MODEL_PATH + urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _LOCAL_MODEL_PATH) + return _LOCAL_MODEL_PATH def detect_faces_mediapipe(img_array, sensitivity): diff --git a/packages/ai/python/face_landmarks.py b/packages/ai/python/face_landmarks.py index 9e3d8108..16351a67 100644 --- a/packages/ai/python/face_landmarks.py +++ b/packages/ai/python/face_landmarks.py @@ -79,12 +79,15 @@ def detect_with_solutions(img_array): # ── New API: mp.tasks (mediapipe >= 0.10.30) ─────────────────────── MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" +_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task" MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") MODEL_PATH = os.path.join(MODEL_DIR, "face_landmarker.task") def ensure_model(): - """Download the face landmarker model if not present.""" + """Resolve face landmarker model. Docker path first, then local dev.""" + if os.path.exists(_DOCKER_MODEL_PATH): + return _DOCKER_MODEL_PATH if os.path.exists(MODEL_PATH): return MODEL_PATH os.makedirs(MODEL_DIR, exist_ok=True) diff --git a/packages/ai/python/red_eye_removal.py b/packages/ai/python/red_eye_removal.py index ab767166..e4443530 100644 --- a/packages/ai/python/red_eye_removal.py +++ b/packages/ai/python/red_eye_removal.py @@ -12,19 +12,22 @@ def emit_progress(percent, stage): # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_MESH_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task" -_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") -_FACE_MESH_MODEL_PATH = os.path.join(_MODEL_DIR, "face_landmarker.task") +_DOCKER_MODEL_PATH = "/opt/models/mediapipe/face_landmarker.task" +_LOCAL_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") +_LOCAL_MODEL_PATH = os.path.join(_LOCAL_MODEL_DIR, "face_landmarker.task") def _ensure_face_mesh_model(): - """Download the face landmarker model if not present.""" - if os.path.exists(_FACE_MESH_MODEL_PATH): - return _FACE_MESH_MODEL_PATH - os.makedirs(_MODEL_DIR, exist_ok=True) + """Resolve face landmarker model. Docker path first, then local dev.""" + if os.path.exists(_DOCKER_MODEL_PATH): + return _DOCKER_MODEL_PATH + if os.path.exists(_LOCAL_MODEL_PATH): + return _LOCAL_MODEL_PATH + os.makedirs(_LOCAL_MODEL_DIR, exist_ok=True) import urllib.request emit_progress(15, "Downloading face mesh model") - urllib.request.urlretrieve(_FACE_MESH_MODEL_URL, _FACE_MESH_MODEL_PATH) - return _FACE_MESH_MODEL_PATH + urllib.request.urlretrieve(_FACE_MESH_MODEL_URL, _LOCAL_MODEL_PATH) + return _LOCAL_MODEL_PATH def _mesh_with_solutions(img_array, max_faces=10, min_confidence=0.5): diff --git a/packages/ai/python/restore.py b/packages/ai/python/restore.py index fda50fdc..06cbed45 100644 --- a/packages/ai/python/restore.py +++ b/packages/ai/python/restore.py @@ -220,19 +220,22 @@ def _get_codeformer_path(): # ── Model path for new mp.tasks API ───────────────────────────────── _FACE_DETECT_MODEL_URL = "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.task" -_FACE_DETECT_MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") -_FACE_DETECT_MODEL_PATH = os.path.join(_FACE_DETECT_MODEL_DIR, "blaze_face_short_range.task") +_FACE_DETECT_DOCKER_PATH = "/opt/models/mediapipe/blaze_face_short_range.task" +_FACE_DETECT_LOCAL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", ".models") +_FACE_DETECT_LOCAL_PATH = os.path.join(_FACE_DETECT_LOCAL_DIR, "blaze_face_short_range.task") def _ensure_face_detect_model(): - """Download the face detector model if not present.""" - if os.path.exists(_FACE_DETECT_MODEL_PATH): - return _FACE_DETECT_MODEL_PATH - os.makedirs(_FACE_DETECT_MODEL_DIR, exist_ok=True) + """Resolve face detector model. Docker path first, then local dev.""" + if os.path.exists(_FACE_DETECT_DOCKER_PATH): + return _FACE_DETECT_DOCKER_PATH + if os.path.exists(_FACE_DETECT_LOCAL_PATH): + return _FACE_DETECT_LOCAL_PATH + os.makedirs(_FACE_DETECT_LOCAL_DIR, exist_ok=True) import urllib.request emit_progress(15, "Downloading face detection model") - urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_MODEL_PATH) - return _FACE_DETECT_MODEL_PATH + urllib.request.urlretrieve(_FACE_DETECT_MODEL_URL, _FACE_DETECT_LOCAL_PATH) + return _FACE_DETECT_LOCAL_PATH def enhance_faces(img_bgr, fidelity=0.7):