Pose: match seg/normal aesthetic (Helvetica, gradient title, badges, ImageSlider, hero examples, 0.4B-only preload)

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,11 +1,7 @@
 """Sapiens2 pose-estimation Gradio Space.
 
-Top-down 308-keypoint pose: RTMDet finds people, Sapiens2 estimates keypoints
+Top-down 308-keypoint pose: DETR finds people, Sapiens2 estimates keypoints
 on each crop, and we draw skeleton + keypoints with the GOLIATH palette.
-
-All checkpoints are pulled from HuggingFace at startup so this Space repo
-stays small. The eager pre-load below warms the cache for the detector and
-all 4 pose sizes during boot, so user requests are instant.
 """
 
 # Block mmpretrain: mmdet's reid modules try `import mmpretrain` inside
@@ -14,8 +10,7 @@ all 4 pose sizes during boot, so user requests are instant.
 import sys
 sys.modules["mmpretrain"] = None
 
-# Use the vendored sapiens2 source (the upstream package pins requires-python
-# to >=3.12; we run on 3.10 here, so we sys.path-import the source instead).
+# Use the vendored sapiens2 source.
 import os
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
@@ -46,7 +41,6 @@ from pose_render_utils import visualize_keypoints
 ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
 CONFIGS_DIR = os.path.join(ASSETS_DIR, "configs")
 
-# Sapiens2 pose checkpoints — fetched from HF model repos at startup.
 POSE_MODELS = {
     "0.4B": {
         "repo": "facebook/sapiens2-pose-0.4b",
@@ -69,7 +63,7 @@ POSE_MODELS = {
         "config": os.path.join(CONFIGS_DIR, "sapiens2_5b_keypoints308_shutterstock_goliath_3po-1024x768.py"),
     },
 }
-DEFAULT_SIZE = "1B"
+DEFAULT_SIZE = "0.4B"  # iteration mode — only this is preloaded; others lazy-load on click
 
 DETECTOR_MODEL_ID = "facebook/detr-resnet-50"  # COCO person = label 1
 
@@ -79,7 +73,7 @@ NMS_THR = 0.3
 
 
 # -----------------------------------------------------------------------------
-# Model cache (load once, reuse across requests)
+# Model cache
 
 _pose_model_cache: dict = {}
 _detector_cache: dict = {}
@@ -116,13 +110,13 @@ def _get_pose_model(size: str):
     return _pose_model_cache[size]
 
 
-# -----------------------------------------------------------------------------
-# Eager pre-load: download + warm-load detector + all pose sizes at startup so
-# the first user (and every user thereafter) gets an instant response.
-print("[startup] pre-loading detector + all pose sizes ...")
+# Iteration mode: only preload the default (0.4B) for fast Space boot.
+# Re-enable full preload by uncommenting the loop below.
+print("[startup] pre-loading detector + 0.4B (iteration mode) ...")
 _get_detector()
-for _size in POSE_MODELS:
-    _get_pose_model(_size)
+_get_pose_model(DEFAULT_SIZE)
+# for _size in POSE_MODELS:
+#     _get_pose_model(_size)
 print("[startup] ready.")
 
 
@@ -135,14 +129,14 @@ def _detect_persons(image_rgb: np.ndarray) -> np.ndarray:
     inputs = proc(images=pil_img, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         outputs = model(**inputs)
-    target_sizes = torch.tensor([image_rgb.shape[:2]], device=DEVICE)  # (h, w)
+    target_sizes = torch.tensor([image_rgb.shape[:2]], device=DEVICE)
     results = proc.post_process_object_detection(
         outputs, target_sizes=target_sizes, threshold=BBOX_THR
     )[0]
-    person_mask = results["labels"] == 1  # COCO person
-    boxes = results["boxes"][person_mask].cpu().numpy()  # (N, 4) x1,y1,x2,y2
+    person_mask = results["labels"] == 1
+    boxes = results["boxes"][person_mask].cpu().numpy()
     scores = results["scores"][person_mask].cpu().numpy().reshape(-1, 1)
-    bboxes = np.concatenate([boxes, scores], axis=1)  # (N, 5)
+    bboxes = np.concatenate([boxes, scores], axis=1)
     bboxes = bboxes[nms(bboxes, NMS_THR), :4]
     if len(bboxes) == 0:
         h, w = image_rgb.shape[:2]
@@ -161,12 +155,12 @@ def _estimate_pose(image_bgr: np.ndarray, bboxes: np.ndarray, model) -> Tuple[Li
 
     inputs = torch.cat(inputs_list, dim=0)
     with torch.no_grad():
-        pred = model(inputs)  # (B, K, h, w) heatmaps
+        pred = model(inputs)
 
     pred = pred.cpu().numpy()
     keypoints, scores = [], []
     for i, sample in enumerate(samples_list):
-        kpts_i, scr_i = model.codec.decode(pred[i])  # (1, K, 2), (1, K)
+        kpts_i, scr_i = model.codec.decode(pred[i])
         meta = sample["meta"]
         kpts_i = kpts_i / meta["input_size"] * meta["bbox_scale"] + meta["bbox_center"] - 0.5 * meta["bbox_scale"]
         keypoints.append(kpts_i[0])
@@ -182,7 +176,8 @@ def predict(image: Image.Image, size: str, kpt_thr: float):
     if image is None:
         return None, None
 
-    image_rgb = np.array(image.convert("RGB"))
+    image_pil = image.convert("RGB")
+    image_rgb = np.array(image_pil)
     image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
 
     bboxes = _detect_persons(image_rgb)
@@ -202,6 +197,7 @@ def predict(image: Image.Image, size: str, kpt_thr: float):
         kpt_color=meta["keypoint_colors"],
         link_color=meta["skeleton_link_colors"],
     )
+    vis_pil = Image.fromarray(vis_rgb)
 
     instances = [
         {
@@ -215,7 +211,8 @@ def predict(image: Image.Image, size: str, kpt_thr: float):
         json.dump({"instances": instances}, f)
         json_path = f.name
 
-    return Image.fromarray(vis_rgb), json_path
+    # Slider value: (left=input RGB, right=keypoints+skeleton overlay).
+    return (image_pil, vis_pil), json_path
 
 
 # -----------------------------------------------------------------------------
@@ -227,30 +224,84 @@ EXAMPLES = sorted(
     if n.lower().endswith((".jpg", ".jpeg", ".png"))
 )
 
-with gr.Blocks(title="Sapiens2 Pose", theme=gr.themes.Default()) as demo:
-    gr.Markdown(
-        "# Sapiens2: Pose Estimation\n"
-        "### ICLR 2026\n"
-        "Top-down 308-keypoint human pose. RTMDet finds people; Sapiens2 estimates keypoints.\n\n"
-        "[Code](https://github.com/facebookresearch/sapiens2) · "
-        "[Models](https://huggingface.co/facebook/sapiens2) · "
-        "[Paper](https://openreview.net/pdf?id=IVAlYCqdvW)"
-    )
+CUSTOM_CSS = """
+:root, body, .gradio-container, button, input, select, textarea,
+.gradio-container *:not(code):not(pre) {
+    font-family: "Helvetica Neue", Helvetica, Arial, sans-serif !important;
+    -webkit-font-smoothing: antialiased;
+    -moz-osx-font-smoothing: grayscale;
+}
+
+#title { text-align: center; font-size: 44px; font-weight: 700;
+         letter-spacing: -0.01em; margin: 28px 0 4px;
+         background: linear-gradient(90deg, #1d4ed8 0%, #6d28d9 50%, #be185d 100%);
+         -webkit-background-clip: text; -webkit-text-fill-color: transparent;
+         background-clip: text; }
+#subtitle { text-align: center; font-size: 12px; color: #64748b;
+            letter-spacing: 0.18em; margin: 0 0 14px; text-transform: uppercase;
+            font-weight: 500; }
+#badges { display: flex; justify-content: center; flex-wrap: wrap;
+          gap: 8px; margin: 0 0 32px; }
+.pill { display: inline-flex; align-items: center; gap: 6px;
+        padding: 7px 14px; border-radius: 999px;
+        background: #f1f5f9; color: #0f172a !important;
+        font-size: 13px; font-weight: 500; letter-spacing: 0.01em;
+        text-decoration: none !important; border: 1px solid #e2e8f0;
+        transition: background 150ms ease, transform 150ms ease, border-color 150ms ease; }
+.pill:hover { background: #0f172a; color: #f8fafc !important;
+              border-color: #0f172a; transform: translateY(-1px); }
+.pill svg { width: 14px; height: 14px; }
+"""
+
+HEADER_HTML = """
+<div id="title">Sapiens2: Pose</div>
+<div id="subtitle">ICLR 2026</div>
+<div id="badges">
+  <a class="pill" href="https://github.com/facebookresearch/sapiens2" target="_blank" rel="noopener">
+    <svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 .3a12 12 0 0 0-3.8 23.4c.6.1.8-.3.8-.6v-2c-3.3.7-4-1.6-4-1.6-.6-1.4-1.4-1.8-1.4-1.8-1.1-.7.1-.7.1-.7 1.3.1 2 1.3 2 1.3 1.1 1.9 3 1.4 3.7 1 .1-.8.4-1.4.8-1.7-2.7-.3-5.5-1.3-5.5-5.9 0-1.3.5-2.4 1.3-3.2-.1-.4-.6-1.6.1-3.2 0 0 1-.3 3.3 1.2a11.5 11.5 0 0 1 6 0c2.3-1.5 3.3-1.2 3.3-1.2.7 1.6.2 2.8.1 3.2.8.8 1.3 1.9 1.3 3.2 0 4.6-2.8 5.6-5.5 5.9.4.4.8 1.1.8 2.2v3.3c0 .3.2.7.8.6A12 12 0 0 0 12 .3"/></svg>
+    Code
+  </a>
+  <a class="pill" href="https://huggingface.co/facebook/sapiens2" target="_blank" rel="noopener">
+    🤗 Models
+  </a>
+  <a class="pill" href="https://openreview.net/pdf?id=IVAlYCqdvW" target="_blank" rel="noopener">
+    <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"/><polyline points="14 2 14 8 20 8"/><line x1="9" y1="13" x2="15" y2="13"/><line x1="9" y1="17" x2="15" y2="17"/></svg>
+    Paper
+  </a>
+  <a class="pill" href="https://rawalkhirodkar.github.io/sapiens2" target="_blank" rel="noopener">
+    <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="2" y1="12" x2="22" y2="12"/><path d="M12 2a15.3 15.3 0 0 1 4 10 15.3 15.3 0 0 1-4 10 15.3 15.3 0 0 1-4-10 15.3 15.3 0 0 1 4-10z"/></svg>
+    Project
+  </a>
+</div>
+"""
+
+with gr.Blocks(title="Sapiens2 Pose", theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
+    gr.HTML(HEADER_HTML)
+
+    with gr.Row(equal_height=True):
+        inp = gr.Image(label="Input", type="pil", height=640)
+        out_img = gr.ImageSlider(
+            label="Input ↔ Pose",
+            type="pil",
+            height=640,
+            max_height=640,
+            slider_position=50,
+        )
+
     with gr.Row():
-        with gr.Column():
-            inp = gr.Image(label="Input", type="pil")
-            with gr.Row():
-                size = gr.Radio(
-                    choices=list(POSE_MODELS.keys()),
-                    value=DEFAULT_SIZE,
-                    label="Model size",
-                )
-                thr = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Keypoint threshold")
-            run = gr.Button("Run", variant="primary")
-            gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14)
-        with gr.Column():
-            out_img = gr.Image(label="Pose-308 result", type="pil")
-            out_json = gr.File(label="Keypoints (.json)")
+        size = gr.Radio(
+            choices=list(POSE_MODELS.keys()),
+            value=DEFAULT_SIZE,
+            label="Model",
+            scale=2,
+        )
+        thr = gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Keypoint threshold", scale=2)
+        run = gr.Button("Run", variant="primary", size="lg", scale=1)
+
+    gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14)
+
+    with gr.Accordion("Raw Keypoints", open=False):
+        out_json = gr.File(label="Keypoints (.json — bboxes + 308 keypoints + scores per person)")
 
     run.click(predict, inputs=[inp, size, thr], outputs=[out_img, out_json])