Fix Gradio API crash, render_wireframe args, add brow_lift/mentoplasty masks

README.md

@@ -1,11 +1,11 @@
 ---
 title: LandmarkDiff
-emoji: 🔬
+emoji: "\U0001F52C"
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 5.12.0
-python_version: 3.11
+sdk_version: 5.29.0
+python_version: "3.11"
 app_file: app.py
 pinned: true
 license: mit
@@ -52,4 +52,4 @@ GPU modes (ControlNet, img2img) with photorealistic rendering are available in t
 - [Wiki](https://github.com/dreamlessx/LandmarkDiff-public/wiki)
 - [Discussions](https://github.com/dreamlessx/LandmarkDiff-public/discussions)
 
-**Version:** v0.2.0
+**Version:** v0.2.1

app.py

@@ -6,12 +6,12 @@ import cv2
 import gradio as gr
 import numpy as np
 
-from landmarkdiff.landmarks import extract_landmarks, render_landmark_image
 from landmarkdiff.conditioning import render_wireframe
-from landmarkdiff.manipulation import apply_procedure_preset, PROCEDURE_LANDMARKS
+from landmarkdiff.landmarks import extract_landmarks
+from landmarkdiff.manipulation import PROCEDURE_LANDMARKS, apply_procedure_preset
 from landmarkdiff.masking import generate_surgical_mask
 
-VERSION = "v0.2.0"
+VERSION = "v0.2.1"
 
 GITHUB_URL = "https://github.com/dreamlessx/LandmarkDiff-public"
 DOCS_URL = f"{GITHUB_URL}/tree/main/docs"
@@ -31,6 +31,7 @@ PROCEDURE_DESCRIPTIONS = {
 def warp_image_tps(image, src_pts, dst_pts):
     """Thin-plate spline warp (CPU only)."""
     from landmarkdiff.synthetic.tps_warp import warp_image_tps as _warp
+
     return _warp(image, src_pts, dst_pts)
 
 
@@ -55,14 +56,20 @@ def process_image(image_rgb, procedure, intensity):
 
     face = extract_landmarks(image_bgr)
     if face is None:
-        return image_rgb_512, image_rgb_512, image_rgb_512, image_rgb_512, "No face detected. Try a clearer photo with good lighting."
+        return (
+            image_rgb_512,
+            image_rgb_512,
+            image_rgb_512,
+            image_rgb_512,
+            "No face detected. Try a clearer photo with good lighting.",
+        )
 
     # Manipulate landmarks
     manipulated = apply_procedure_preset(face, procedure, float(intensity), image_size=512)
 
-    # Generate wireframe
-    wireframe = render_wireframe(manipulated, (512, 512))
-    wireframe_rgb = cv2.cvtColor(wireframe, cv2.COLOR_BGR2RGB) if wireframe.ndim == 3 else wireframe
+    # Generate wireframe (pass width and height as separate keyword args)
+    wireframe = render_wireframe(manipulated, width=512, height=512)
+    wireframe_rgb = cv2.cvtColor(wireframe, cv2.COLOR_GRAY2RGB)
 
     # Generate mask
     mask = generate_surgical_mask(face, procedure, 512, 512)
@@ -77,9 +84,7 @@ def process_image(image_rgb, procedure, intensity):
     side_by_side = np.hstack([image_rgb_512, composited_rgb])
 
     # Displacement stats
-    displacement = np.mean(np.linalg.norm(
-        manipulated.pixel_coords - face.pixel_coords, axis=1
-    ))
+    displacement = np.mean(np.linalg.norm(manipulated.pixel_coords - face.pixel_coords, axis=1))
 
     info = (
         f"Procedure: {procedure}\n"
@@ -156,7 +161,8 @@ HEADER_MD = f"""
 
 **Anatomically-conditioned facial surgery outcome prediction from standard clinical photography**
 
-Upload a face photo, select a procedure, and adjust intensity to see a predicted surgical outcome in real time.
+Upload a face photo, select a procedure, and adjust intensity to see a predicted
+surgical outcome in real time.
 This demo runs TPS (thin-plate spline) warping on CPU. The full package also supports
 GPU-accelerated ControlNet and img2img inference modes.
 
@@ -219,7 +225,10 @@ with gr.Blocks(
                     label="Surgical Procedure",
                 )
                 intensity = gr.Slider(
-                    minimum=0, maximum=100, value=50, step=1,
+                    minimum=0,
+                    maximum=100,
+                    value=50,
+                    step=1,
                     label="Intensity (%)",
                     info="0 = no change, 100 = maximum effect",
                 )
@@ -262,7 +271,10 @@ with gr.Blocks(
                             proc_idx = row_idx * 3 + col_idx
                             if proc_idx < len(PROCEDURES):
                                 cmp_outputs.append(
-                                    gr.Image(label=PROCEDURES[proc_idx].replace("_", " ").title(), height=200)
+                                    gr.Image(
+                                        label=PROCEDURES[proc_idx].replace("_", " ").title(),
+                                        height=200,
+                                    )
                                 )
 
         cmp_btn.click(
@@ -285,7 +297,9 @@ with gr.Blocks(
                 )
                 sweep_btn = gr.Button("Generate Sweep", variant="primary", size="lg")
             with gr.Column(scale=2):
-                sweep_gallery = gr.Gallery(label="Intensity Sweep (0% - 100%)", columns=3, height=400)
+                sweep_gallery = gr.Gallery(
+                    label="Intensity Sweep (0% - 100%)", columns=3, height=400
+                )
 
         sweep_btn.click(
             fn=intensity_sweep,

landmarkdiff/masking.py

@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING
 import cv2
 import numpy as np
 
-from landmarkdiff.landmarks import FaceLandmarks, LANDMARK_REGIONS
+from landmarkdiff.landmarks import FaceLandmarks
 
 if TYPE_CHECKING:
     from landmarkdiff.clinical import ClinicalFlags
@@ -21,39 +21,188 @@ if TYPE_CHECKING:
 MASK_CONFIG: dict[str, dict] = {
     "rhinoplasty": {
         "landmark_indices": [
-            1, 2, 4, 5, 6, 19, 94, 141, 168, 195, 197, 236, 240,
-            274, 275, 278, 279, 294, 326, 327, 360, 363, 370, 456, 460,
+            1,
+            2,
+            4,
+            5,
+            6,
+            19,
+            94,
+            141,
+            168,
+            195,
+            197,
+            236,
+            240,
+            274,
+            275,
+            278,
+            279,
+            294,
+            326,
+            327,
+            360,
+            363,
+            370,
+            456,
+            460,
         ],
         "dilation_px": 30,
         "feather_sigma": 15.0,
     },
     "blepharoplasty": {
         "landmark_indices": [
-            33, 7, 163, 144, 145, 153, 154, 155, 157, 158, 159, 160, 161, 246,
-            362, 382, 381, 380, 374, 373, 390, 249, 263, 466, 388, 387, 386,
-            385, 384, 398,
+            33,
+            7,
+            163,
+            144,
+            145,
+            153,
+            154,
+            155,
+            157,
+            158,
+            159,
+            160,
+            161,
+            246,
+            362,
+            382,
+            381,
+            380,
+            374,
+            373,
+            390,
+            249,
+            263,
+            466,
+            388,
+            387,
+            386,
+            385,
+            384,
+            398,
         ],
         "dilation_px": 15,
         "feather_sigma": 10.0,
     },
     "rhytidectomy": {
         "landmark_indices": [
-            10, 21, 54, 58, 67, 93, 103, 109, 127, 132, 136, 150, 162, 172,
-            176, 187, 207, 213, 234, 284, 297, 323, 332, 338, 356, 361, 365,
-            379, 389, 397, 400, 427, 454,
+            10,
+            21,
+            54,
+            58,
+            67,
+            93,
+            103,
+            109,
+            127,
+            132,
+            136,
+            150,
+            162,
+            172,
+            176,
+            187,
+            207,
+            213,
+            234,
+            284,
+            297,
+            323,
+            332,
+            338,
+            356,
+            361,
+            365,
+            379,
+            389,
+            397,
+            400,
+            427,
+            454,
         ],
         "dilation_px": 40,
         "feather_sigma": 20.0,
     },
     "orthognathic": {
         "landmark_indices": [
-            0, 17, 18, 36, 37, 39, 40, 57, 61, 78, 80, 81, 82, 84, 87, 88,
-            91, 95, 146, 167, 169, 170, 175, 181, 191, 200, 201, 202, 204,
-            208, 211, 212, 214,
+            0,
+            17,
+            18,
+            36,
+            37,
+            39,
+            40,
+            57,
+            61,
+            78,
+            80,
+            81,
+            82,
+            84,
+            87,
+            88,
+            91,
+            95,
+            146,
+            167,
+            169,
+            170,
+            175,
+            181,
+            191,
+            200,
+            201,
+            202,
+            204,
+            208,
+            211,
+            212,
+            214,
         ],
         "dilation_px": 35,
         "feather_sigma": 18.0,
     },
+    "brow_lift": {
+        "landmark_indices": [
+            70,
+            63,
+            105,
+            66,
+            107,  # left brow
+            300,
+            293,
+            334,
+            296,
+            336,  # right brow
+            9,
+            8,
+            10,  # forehead midline
+            109,
+            67,
+            103,  # upper face left
+            338,
+            297,
+            332,  # upper face right
+        ],
+        "dilation_px": 25,
+        "feather_sigma": 15.0,
+    },
+    "mentoplasty": {
+        "landmark_indices": [
+            148,
+            149,
+            150,
+            152,
+            171,
+            175,
+            176,
+            377,
+        ],
+        "dilation_px": 25,
+        "feather_sigma": 12.0,
+    },
 }
 
 
@@ -62,7 +211,7 @@ def generate_surgical_mask(
     procedure: str,
     width: int | None = None,
     height: int | None = None,
-    clinical_flags: "ClinicalFlags | None" = None,
+    clinical_flags: ClinicalFlags | None = None,
     image: np.ndarray | None = None,
 ) -> np.ndarray:
     """Generate a feathered surgical mask for a procedure.
@@ -134,15 +283,20 @@ def generate_surgical_mask(
     if clinical_flags is not None:
         # Vitiligo: reduce mask over depigmented patches to preserve them
         if clinical_flags.vitiligo and image is not None:
-            from landmarkdiff.clinical import detect_vitiligo_patches, adjust_mask_for_vitiligo
+            from landmarkdiff.clinical import adjust_mask_for_vitiligo, detect_vitiligo_patches
+
             patches = detect_vitiligo_patches(image, face)
             mask = adjust_mask_for_vitiligo(mask, patches)
 
         # Keloid: soften transitions in keloid-prone regions
         if clinical_flags.keloid_prone and clinical_flags.keloid_regions:
-            from landmarkdiff.clinical import get_keloid_exclusion_mask, adjust_mask_for_keloid
+            from landmarkdiff.clinical import adjust_mask_for_keloid, get_keloid_exclusion_mask
+
             keloid_mask = get_keloid_exclusion_mask(
-                face, clinical_flags.keloid_regions, w, h,
+                face,
+                clinical_flags.keloid_regions,
+                w,
+                h,
             )
             mask = adjust_mask_for_keloid(mask, keloid_mask)
 

requirements.txt

@@ -3,3 +3,4 @@ mediapipe>=0.10.9
 opencv-python-headless>=4.9.0
 numpy>=1.26.0
 Pillow>=10.0.0
+pydantic<2.11