Improve DDF heatmap contrast controls

app.py

@@ -236,6 +236,34 @@ def ddf_magnitude(ddf_xyz: np.ndarray) -> np.ndarray:
     return np.linalg.norm(ddf_xyz.astype(np.float32, copy=False), axis=-1)
 
 
+def ddf_color_limits(mag_xyz: np.ndarray, scale_mode: str) -> Tuple[float, float]:
+    values = mag_xyz[np.isfinite(mag_xyz)]
+    if values.size == 0:
+        return 0.0, 1.0
+
+    if scale_mode == "高位移增强":
+        vmin, vmax = np.percentile(values, [50.0, 99.5])
+    elif scale_mode == "绝对范围":
+        vmin, vmax = np.min(values), np.max(values)
+    else:
+        vmin, vmax = np.percentile(values, [1.0, 99.0])
+
+    vmin = float(vmin)
+    vmax = float(vmax)
+    if not np.isfinite(vmin) or not np.isfinite(vmax) or vmax <= vmin:
+        vmax = vmin + 1e-6
+    return vmin, vmax
+
+
+def format_mm(value: float) -> str:
+    value = float(value)
+    if abs(value) < 0.01:
+        return f"{value:.4f}"
+    if abs(value) < 1.0:
+        return f"{value:.3f}"
+    return f"{value:.2f}"
+
+
 def render_image(image: np.ndarray, caption: str, cmap: str = "gray") -> None:
     fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
     ax.imshow(orient_for_display(image), cmap=cmap, interpolation="nearest")
@@ -254,11 +282,27 @@ def render_rgb(image: np.ndarray, caption: str) -> None:
     plt.close(fig)
 
 
-def render_heatmap_overlay(base_slice: np.ndarray, mag_slice: np.ndarray, alpha: float, caption: str) -> None:
+def render_heatmap_overlay(
+    base_slice: np.ndarray,
+    mag_slice: np.ndarray,
+    alpha: float,
+    caption: str,
+    vmin: float,
+    vmax: float,
+    show_base: bool,
+) -> None:
     fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
+    if show_base:
         ax.imshow(orient_for_display(normalize_slice(base_slice)), cmap="gray", interpolation="nearest")
     heat = orient_for_display(mag_slice)
-    im = ax.imshow(heat, cmap="inferno", alpha=alpha, interpolation="nearest")
+    im = ax.imshow(
+        heat,
+        cmap="turbo",
+        alpha=alpha if show_base else 1.0,
+        interpolation="nearest",
+        vmin=vmin,
+        vmax=vmax,
+    )
     ax.set_title(caption, fontsize=10)
     ax.axis("off")
     fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
@@ -300,14 +344,28 @@ def render_overlay_views(fixed_xyz: np.ndarray, warped_xyz: np.ndarray) -> None:
 def render_ddf_views(fixed_xyz: np.ndarray, ddf_xyz: np.ndarray) -> None:
     fixed_xyz, ddf_xyz = crop_to_common_shape(fixed_xyz, ddf_xyz)
     mag = ddf_magnitude(ddf_xyz)
-    alpha = st.slider("热力图透明度", 0.0, 1.0, 0.55, 0.05)
 
     stats = ddf_summary(ddf_xyz)
     c1, c2, c3, c4 = st.columns(4)
-    c1.metric("平均位移 mm", f"{stats['ddf_mean']:.3f}")
-    c2.metric("P95 mm", f"{stats['ddf_p95']:.3f}")
-    c3.metric("最大 mm", f"{stats['ddf_max']:.3f}")
-    c4.metric("标准差", f"{stats['ddf_std']:.3f}")
+    c1.metric("平均位移 mm", format_mm(stats["ddf_mean"]))
+    c2.metric("P95 mm", format_mm(stats["ddf_p95"]))
+    c3.metric("最大 mm", format_mm(stats["ddf_max"]))
+    c4.metric("标准差", format_mm(stats["ddf_std"]))
+
+    if stats["ddf_max"] < 0.1:
+        st.info(f"当前最大位移只有 {format_mm(stats['ddf_max'])} mm，通常说明这是 smoke 权重或模型还没有充分训练。")
+
+    controls = st.columns(3)
+    with controls[0]:
+        display_mode = st.radio("显示方式", ["仅位移热力图", "CT叠加热力图"], horizontal=True)
+    with controls[1]:
+        scale_mode = st.selectbox("色阶范围", ["相对增强", "高位移增强", "绝对范围"])
+    with controls[2]:
+        alpha = st.slider("叠加透明度", 0.0, 1.0, 0.55, 0.05)
+
+    vmin, vmax = ddf_color_limits(mag, scale_mode)
+    show_base = display_mode == "CT叠加热力图"
+    st.caption(f"当前色阶范围：{format_mm(vmin)} - {format_mm(vmax)} mm")
 
     columns = st.columns(3)
     planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")]
@@ -319,6 +377,9 @@ def render_ddf_views(fixed_xyz: np.ndarray, ddf_xyz: np.ndarray) -> None:
                 get_slice(fixed_xyz, plane, index),
                 get_slice(mag, plane, index),
                 alpha=alpha,
+                vmin=vmin,
+                vmax=vmax,
+                show_base=show_base,
                 caption=f"{cn_name} 形变强度 #{index}",
             )