Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,229 @@
+"""LandmarkDiff Hugging Face Spaces Demo - TPS-only (CPU)."""
+
+from __future__ import annotations
+
+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.masking import generate_surgical_mask
+
+
+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)
+
+
+def mask_composite(warped, original, mask):
+    """Alpha blend warped into original using mask."""
+    mask_3 = np.stack([mask] * 3, axis=-1) if mask.ndim == 2 else mask
+    return (warped * mask_3 + original * (1.0 - mask_3)).astype(np.uint8)
+
+
+PROCEDURES = list(PROCEDURE_LANDMARKS.keys())
+
+
+def process_image(image_rgb, procedure, intensity):
+    """Process a single image through the TPS pipeline."""
+    if image_rgb is None:
+        blank = np.zeros((512, 512, 3), dtype=np.uint8)
+        return blank, blank, blank, blank, "Upload a face photo to begin."
+
+    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+    image_bgr = cv2.resize(image_bgr, (512, 512))
+    image_rgb_512 = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+
+    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."
+
+    # 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 mask
+    mask = generate_surgical_mask(face, procedure, 512, 512)
+    mask_vis = (mask * 255).astype(np.uint8)
+
+    # TPS warp + composite
+    warped = warp_image_tps(image_bgr, face.pixel_coords, manipulated.pixel_coords)
+    composited = mask_composite(warped, image_bgr, mask)
+    composited_rgb = cv2.cvtColor(composited, cv2.COLOR_BGR2RGB)
+
+    # Side by side
+    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
+    ))
+
+    info = (
+        f"Procedure: {procedure}\n"
+        f"Intensity: {intensity:.0f}%\n"
+        f"Landmarks: {len(face.landmarks)}\n"
+        f"Avg displacement: {displacement:.1f} px\n"
+        f"Confidence: {face.confidence:.2f}\n"
+        f"Mode: TPS (CPU)"
+    )
+
+    return wireframe_rgb, mask_vis, composited_rgb, side_by_side, info
+
+
+def compare_procedures(image_rgb, intensity):
+    """Compare all procedures at the same intensity."""
+    if image_rgb is None:
+        blank = np.zeros((512, 512, 3), dtype=np.uint8)
+        return [blank] * len(PROCEDURES)
+
+    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+    image_bgr = cv2.resize(image_bgr, (512, 512))
+
+    face = extract_landmarks(image_bgr)
+    if face is None:
+        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        return [rgb] * len(PROCEDURES)
+
+    results = []
+    for proc in PROCEDURES:
+        manip = apply_procedure_preset(face, proc, float(intensity), image_size=512)
+        mask = generate_surgical_mask(face, proc, 512, 512)
+        warped = warp_image_tps(image_bgr, face.pixel_coords, manip.pixel_coords)
+        comp = mask_composite(warped, image_bgr, mask)
+        results.append(cv2.cvtColor(comp, cv2.COLOR_BGR2RGB))
+
+    return results
+
+
+def intensity_sweep(image_rgb, procedure):
+    """Generate intensity sweep from 0 to 100."""
+    if image_rgb is None:
+        return []
+
+    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+    image_bgr = cv2.resize(image_bgr, (512, 512))
+
+    face = extract_landmarks(image_bgr)
+    if face is None:
+        return []
+
+    steps = [0, 20, 40, 60, 80, 100]
+    results = []
+    for val in steps:
+        if val == 0:
+            results.append((cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB), "0%"))
+            continue
+        manip = apply_procedure_preset(face, procedure, float(val), image_size=512)
+        mask = generate_surgical_mask(face, procedure, 512, 512)
+        warped = warp_image_tps(image_bgr, face.pixel_coords, manip.pixel_coords)
+        comp = mask_composite(warped, image_bgr, mask)
+        results.append((cv2.cvtColor(comp, cv2.COLOR_BGR2RGB), f"{val}%"))
+
+    return results
+
+
+with gr.Blocks(
+    title="LandmarkDiff - Surgical Outcome Prediction",
+    theme=gr.themes.Soft(),
+) as demo:
+    gr.Markdown(
+        "# LandmarkDiff\n"
+        "**Anatomically-conditioned facial surgery outcome prediction**\n\n"
+        "Upload a face photo, select a procedure, and adjust intensity. "
+        "This demo uses TPS warping (CPU) for real-time preview. "
+        "GPU-accelerated ControlNet/img2img modes are available in the full package.\n\n"
+        "[GitHub](https://github.com/dreamlessx/LandmarkDiff-public) | "
+        "[Paper](https://github.com/dreamlessx/LandmarkDiff-public/tree/main/paper)"
+    )
+
+    with gr.Tab("Single Procedure"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(label="Upload Face Photo", type="numpy", height=350)
+                procedure = gr.Radio(
+                    choices=PROCEDURES,
+                    value="rhinoplasty",
+                    label="Surgical Procedure",
+                )
+                intensity = gr.Slider(
+                    minimum=0, maximum=100, value=50, step=1,
+                    label="Intensity (%)",
+                    info="0 = no change, 100 = maximum effect",
+                )
+                run_btn = gr.Button("Generate Preview", variant="primary", size="lg")
+                info_box = gr.Textbox(label="Info", lines=6, interactive=False)
+
+            with gr.Column(scale=2):
+                with gr.Row():
+                    out_wireframe = gr.Image(label="Deformed Wireframe", height=256)
+                    out_mask = gr.Image(label="Surgical Mask", height=256)
+                with gr.Row():
+                    out_result = gr.Image(label="Predicted Result", height=256)
+                    out_sidebyside = gr.Image(label="Before / After", height=256)
+
+        run_btn.click(
+            fn=process_image,
+            inputs=[input_image, procedure, intensity],
+            outputs=[out_wireframe, out_mask, out_result, out_sidebyside, info_box],
+        )
+        for trigger in [procedure, intensity]:
+            trigger.change(
+                fn=process_image,
+                inputs=[input_image, procedure, intensity],
+                outputs=[out_wireframe, out_mask, out_result, out_sidebyside, info_box],
+            )
+
+    with gr.Tab("Compare Procedures"):
+        gr.Markdown("Compare all procedures side by side at the same intensity.")
+        with gr.Row():
+            with gr.Column(scale=1):
+                cmp_image = gr.Image(label="Upload Face Photo", type="numpy", height=300)
+                cmp_intensity = gr.Slider(0, 100, 50, step=1, label="Intensity (%)")
+                cmp_btn = gr.Button("Compare All", variant="primary", size="lg")
+            with gr.Column(scale=2):
+                cmp_outputs = []
+                rows_needed = (len(PROCEDURES) + 2) // 3
+                for row_idx in range(rows_needed):
+                    with gr.Row():
+                        for col_idx in range(3):
+                            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)
+                                )
+
+        cmp_btn.click(
+            fn=compare_procedures,
+            inputs=[cmp_image, cmp_intensity],
+            outputs=cmp_outputs,
+        )
+
+    with gr.Tab("Intensity Sweep"):
+        gr.Markdown("See how a procedure looks across intensity levels.")
+        with gr.Row():
+            with gr.Column(scale=1):
+                sweep_image = gr.Image(label="Upload Face Photo", type="numpy", height=300)
+                sweep_procedure = gr.Radio(
+                    choices=PROCEDURES,
+                    value="rhinoplasty",
+                    label="Procedure",
+                )
+                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_btn.click(
+            fn=intensity_sweep,
+            inputs=[sweep_image, sweep_procedure],
+            outputs=[sweep_gallery],
+        )
+
+if __name__ == "__main__":
+    demo.launch()