Seg: add hero banner + 2-col layout (input | annotated overlay) + hover-pop CSS + 29 classes

.gitattributes

@@ -34,3 +34,4 @@ 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
+*.gif filter=lfs diff=lfs merge=lfs -text

app.py

@@ -30,6 +30,7 @@ _ = SegEstimator
 
 ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
 CONFIGS_DIR = os.path.join(ASSETS_DIR, "configs")
+HERO_GIF = os.path.join(ASSETS_DIR, "sapiens2.gif")
 
 SEG_MODELS = {
     "0.4B": {
@@ -57,13 +58,12 @@ DEFAULT_SIZE = "1B"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Class id → readable label, RGB color tuple, "#rrggbb" hex (for gr.AnnotatedImage color_map).
+# All 29 named classes (Background is class 0).
 _CLASS_LABELS = {cid: meta["name"].replace("_", " ") for cid, meta in DOME_CLASSES_29.items()}
 _CLASS_COLORS_RGB = {cid: meta["color"] for cid, meta in DOME_CLASSES_29.items()}
 _CLASS_COLORS_HEX = {
     _CLASS_LABELS[cid]: "#{:02x}{:02x}{:02x}".format(*meta["color"])
     for cid, meta in DOME_CLASSES_29.items()
-    if cid != 0  # background skipped in AnnotatedImage
 }
 
 
@@ -105,14 +105,14 @@ def _segment(image_bgr: np.ndarray, model) -> np.ndarray:
 
 
 def _label_map_to_annotations(label_map: np.ndarray) -> List[Tuple[np.ndarray, str]]:
-    """Convert a (H, W) class-id map into AnnotatedImage's [(mask, label), ...] format.
+    """Convert (H, W) class-id map → AnnotatedImage's [(bool_mask, label), ...] list.
 
-    Each non-background class present in the prediction yields one boolean mask.
+    Includes all 29 named classes (Background as well) so the legend is complete.
     """
     annotations: List[Tuple[np.ndarray, str]] = []
     for cid in np.unique(label_map):
         cid = int(cid)
-        if cid == 0 or cid not in _CLASS_LABELS:  # skip Background
+        if cid not in _CLASS_LABELS:
             continue
         mask = (label_map == cid)
         if not mask.any():
@@ -146,7 +146,7 @@ def predict(image: Image.Image, size: str, opacity: float):
     label_map = _segment(image_bgr, model)               # (H, W)
 
     annotations = _label_map_to_annotations(label_map)
-    annotated = (image_pil, annotations)                 # gr.AnnotatedImage value tuple
+    annotated = (image_pil, annotations)
 
     overlay_bgr = _label_map_to_overlay(image_bgr, label_map, float(opacity))
     overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
@@ -166,37 +166,71 @@ EXAMPLES = sorted(
     if n.lower().endswith((".jpg", ".jpeg", ".png"))
 )
 
-with gr.Blocks(title="Sapiens2 Seg", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        "# Sapiens2: Body-Part Segmentation\n"
-        "### ICLR 2026\n"
-        "29-class per-pixel body-part segmentation. Hover over the predicted regions to highlight each part.\n\n"
-        "[Code](https://github.com/facebookresearch/sapiens2) · "
-        "[Models](https://huggingface.co/facebook/sapiens2) · "
-        "[Paper](https://openreview.net/pdf?id=IVAlYCqdvW)"
-    )
+CUSTOM_CSS = """
+#hero { text-align: center; padding: 24px 0 12px; }
+#hero img { max-width: 360px; width: 100%; border-radius: 12px;
+            box-shadow: 0 8px 28px rgba(0,0,0,0.18); }
+#title { text-align: center; font-size: 44px; font-weight: 700; margin: 14px 0 0;
+         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: 16px; color: #64748b;
+            letter-spacing: 0.08em; margin: 4px 0 12px; text-transform: uppercase; }
+#tagline { text-align: center; font-size: 16px; color: #475569;
+           max-width: 720px; margin: 4px auto 18px; line-height: 1.55; }
+#badges { text-align: center; margin: 0 0 28px; }
+#badges img { display: inline-block; margin: 0 4px; height: 28px; vertical-align: middle; }
+
+/* AnnotatedImage hover pop: when a region is highlighted, lift + brighten it */
+#seg-out .layer-wrap canvas { transition: transform 180ms ease, filter 180ms ease; }
+#seg-out .layer-wrap.active canvas { transform: translateY(-3px) scale(1.01);
+                                     filter: brightness(1.18) drop-shadow(0 6px 12px rgba(0,0,0,0.25)); }
+#seg-out .legend-item { transition: transform 140ms ease; }
+#seg-out .legend-item:hover { transform: translateY(-1px); }
+"""
+
+HEADER_HTML = """
+<div id="hero"><img src="/gradio_api/file=assets/sapiens2.gif" alt="Sapiens2"/></div>
+<div id="title">Sapiens2: Body-Part Segmentation</div>
+<div id="subtitle">ICLR 2026 · 29-class semantic segmentation</div>
+<div id="tagline">Hover any predicted region to see its class name. Pick a model size and an
+overlay opacity, then click Run.</div>
+<div id="badges">
+  <a href="https://github.com/facebookresearch/sapiens2"><img src="https://img.shields.io/badge/Code-GitHub-181717?logo=github&logoColor=white"></a>
+  <a href="https://huggingface.co/facebook/sapiens2"><img src="https://img.shields.io/badge/Models-HuggingFace-FF9D00?logo=huggingface&logoColor=white"></a>
+  <a href="https://openreview.net/pdf?id=IVAlYCqdvW"><img src="https://img.shields.io/badge/Paper-OpenReview-94DD15?logo=adobeacrobatreader&logoColor=white"></a>
+  <a href="https://rawalkhirodkar.github.io/sapiens2"><img src="https://img.shields.io/badge/Project-Page-007FFF?logo=googlechrome&logoColor=white"></a>
+</div>
+"""
+
+with gr.Blocks(title="Sapiens2 Seg", theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
+    gr.HTML(HEADER_HTML)
+
+    with gr.Row(equal_height=True):
+        inp = gr.Image(label="Input (RGB)", type="pil", height=540)
+        out_annot = gr.AnnotatedImage(
+            label="Body parts — hover to highlight",
+            color_map=_CLASS_COLORS_HEX,
+            height=540,
+            show_legend=True,
+            elem_id="seg-out",
+        )
+
     with gr.Row():
-        with gr.Column(scale=1):
-            inp = gr.Image(label="Input", type="pil", height=420)
-            with gr.Row():
-                size = gr.Radio(
-                    choices=list(SEG_MODELS.keys()),
-                    value=DEFAULT_SIZE,
-                    label="Model size",
-                )
-                opacity = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Overlay opacity")
-            run = gr.Button("Run", variant="primary")
-            gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14)
-        with gr.Column(scale=2):
-            out_annot = gr.AnnotatedImage(
-                label="Body parts (hover to highlight)",
-                color_map=_CLASS_COLORS_HEX,
-                height=560,
-                show_legend=True,
-            )
-            with gr.Accordion("Static overlay + raw labels", open=False):
-                out_img = gr.Image(label="Color overlay (PNG)", type="pil")
-                out_npy = gr.File(label="Raw labels (.npy uint8, class indices 0–28)")
+        size = gr.Radio(
+            choices=list(SEG_MODELS.keys()),
+            value=DEFAULT_SIZE,
+            label="Model size",
+            scale=2,
+        )
+        opacity = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Overlay opacity", 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("Static overlay + raw labels", open=False):
+        out_img = gr.Image(label="Color overlay (PNG)", type="pil")
+        out_npy = gr.File(label="Raw labels (.npy uint8, class indices 0–28)")
 
     run.click(predict, inputs=[inp, size, opacity], outputs=[out_annot, out_img, out_npy])
 
@@ -205,4 +239,4 @@ if __name__ == "__main__":
     if torch.cuda.is_available():
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cudnn.allow_tf32 = True
-    demo.launch(share=False)
+    demo.launch(share=False, allowed_paths=[ASSETS_DIR])