Enables batched input for heatmap generation

Extends the `GenerateHeatmap` transform to support batched inputs,
allowing for more efficient processing of multiple landmark sets.

This change modifies the transform to handle inputs with a batch dimension (B, N, spatial_dims) in addition to single-point inputs (N, spatial_dims).
It also includes a demonstration of 3D heatmap generation using PyVista for visualization.

Author: eclipse0922

2d_mdtest.ipynb

@@ -91,6 +91,7 @@
                 "# Heatmap helper using GenerateHeatmap\n",
                 "# NOTE: GenerateHeatmap expects points in (y, x) == (row, col) order matching array indexing.\n",
                 "# This wrapper accepts user-friendly (x, y) and internally reorders to (y, x).\n",
+                "# It now supports batched inputs.\n",
                 "\n",
                 "sigma = 3.0\n",
                 "\n",
@@ -99,8 +100,12 @@
                 "    s = float(sigma_override) if sigma_override is not None else float(sigma)\n",
                 "    tr = GenerateHeatmap(sigma=s, spatial_shape=(H, W))\n",
                 "    # Reorder (x,y) -> (y,x) for the transform\n",
-                "    pts_yx = np.array([[float(y), float(x)]], dtype=np.float32)\n",
-                "    return tr(pts_yx)  # (N,H,W) where pts interpreted as (row, col)"
+                "    # Support batched and non-batched inputs\n",
+                "    pts = np.array(list(zip(y, x)), dtype=np.float32)\n",
+                "    if pts.ndim == 2:\n",
+                "        pts = pts[np.newaxis, ...]  # Add batch dimension: (N, 2) -> (1, N, 2)\n",
+                "    pts_yx = pts[..., [1, 0]]\n",
+                "    return tr(pts_yx)  # (B, N, H, W)\n"
             ]
         },
         {
@@ -125,10 +130,14 @@
                 "        spatial_shape=None if use_ref else (H, W),\n",
                 "        sigma=s,\n",
                 "    )\n",
-                "    pts_yx = np.array([[float(y), float(x)]], dtype=np.float32)\n",
+                "    # Support batched and non-batched inputs\n",
+                "    pts = np.array(list(zip(y, x)), dtype=np.float32)\n",
+                "    if pts.ndim == 2:\n",
+                "        pts = pts[np.newaxis, ...]  # Add batch dimension: (N, 2) -> (1, N, 2)\n",
+                "    pts_yx = pts[..., [1, 0]]\n",
                 "    data = {\"points\": pts_yx, \"ref\": ref_img}\n",
                 "    out = tr(data)\n",
-                "    return out[\"heatmap\"]"
+                "    return out[\"heatmap\"]\n"
             ]
         },
         {
@@ -179,6 +188,7 @@
                 "num_points = 3  # number of random landmarks\n",
                 "sigma_demo = 3.0  # Gaussian sigma\n",
                 "combine_mode = \"max\"  # or 'sum'\n",
+                "batched_input = True  # Set to True to test batched input\n",
                 "\n",
                 "# Sample random (x,y) points within image bounds (user-friendly)\n",
                 "points_xy = np.array(\n",
@@ -189,10 +199,15 @@
                 "\n",
                 "# Convert to (y,x) for the transform\n",
                 "yx_points = points_xy[:, [1, 0]].copy()\n",
+                "if batched_input:\n",
+                "    yx_points = yx_points[np.newaxis, ...]  # Add a batch dimension\n",
                 "\n",
                 "array_tr = GenerateHeatmap(sigma=sigma_demo, spatial_shape=(H, W))\n",
                 "heatmaps = array_tr(yx_points)  # now correct orientation\n",
                 "\n",
+                "if batched_input:\n",
+                "    heatmaps = heatmaps.squeeze(0)  # Remove batch dim for plotting\n",
+                "\n",
                 "if combine_mode == \"max\":\n",
                 "    combined = heatmaps.max(axis=0)\n",
                 "elif combine_mode == \"sum\":\n",
@@ -230,7 +245,7 @@
                 "    ax.set_title(f\"Point {i}: (x={points_xy[i,0]:.1f}, y={points_xy[i,1]:.1f})\")\n",
                 "    ax.set_axis_off()\n",
                 "plt.tight_layout()\n",
-                "plt.show()"
+                "plt.show()\n"
             ]
         }
     ],