Update app.py

app.py

@@ -45,9 +45,6 @@ from process import  SynthradAlgorithm2
 
 from process_1 import SynthradAlgorithm1
 
-# =========================
-# Streamlit UI
-# =========================
 st.set_page_config(page_title="SynthRad (nnUNetv2) Demo", layout="wide")
 st.title("SynthRad — MRI/CBCT + Mask → synthetic CT")
 st.image("./workflow.png",width=800)
@@ -94,32 +91,26 @@ src = st.radio("Source", ["Sample", "Upload"], index=0, horizontal=True)
 def build_sample_map(task_name: str):
     repo_dir = REPO_DIRS[task_name]
     if task_name == "Task 1 (MR → CT)":
-        vol_key = "mri"
-        vol_fname = "mr.mha"   
-        mask_fname = "mask1.mha"   
+        vol_fname = "mr.mha"
+        mask_fname = "mask1.mha"
     else:
-        vol_key = "cbct"
-        vol_fname = "cbct.mha" 
-        mask_fname = "mask2.mha"   
+        vol_fname = "cbct.mha"
+        mask_fname = "mask2.mha"
+
+    def pack(region_dir):
+        vol_path  = os.path.join(repo_dir, region_dir, vol_fname)
+        mask_path = os.path.join(repo_dir, region_dir, mask_fname)
+        gt_path   = os.path.join(repo_dir, region_dir, "ct.mha")  # 约定：GT=ct.mha
+        return {"vol": vol_path, "mask": mask_path, "gt": gt_path}
+
     sample_map = {
-        "Abdomen (sample)": {
-            "region": "Abdomen",
-            "vol":  os.path.join(repo_dir, "Abdomen",  vol_fname),
-            "mask": os.path.join(repo_dir, "Abdomen", mask_fname),
-        },
-        "Head and Neck (sample)": {
-            "region": "Head and Neck",
-            "vol":  os.path.join(repo_dir, "Head and Neck", vol_fname),
-            "mask": os.path.join(repo_dir, "Head and Neck", mask_fname),
-        },
-        "Thorax (sample)": {
-            "region": "Thorax",
-            "vol":  os.path.join(repo_dir, "Thorax", vol_fname),
-            "mask": os.path.join(repo_dir, "Thorax", mask_fname),
-        },
+        "Abdomen (sample)":       {"region": "Abdomen",       **pack("Abdomen")},
+        "Head and Neck (sample)": {"region": "Head and Neck", **pack("Head and Neck")},
+        "Thorax (sample)":        {"region": "Thorax",        **pack("Thorax")},
     }
     return sample_map
 
+
 SAMPLE_MAP = build_sample_map(task)
 
 
@@ -214,40 +205,116 @@ if run_btn:
 if st.session_state.vol_np is None:
     st.info("Select Upload or Sample, then click Run")
 else:
-
+    in_lps  = sitk.DICOMOrient(st.session_state.input_vol, "LPS")
     out_lps = sitk.DICOMOrient(st.session_state.synth_ct, "LPS")
-    vol = sitk.GetArrayFromImage(out_lps).astype(np.float32)
-    D, H, W = vol.shape
 
+    res = sitk.ResampleImageFilter()
+    res.SetReferenceImage(in_lps)
+    res.SetInterpolator(sitk.sitkLinear)
+    res.SetOutputPixelType(out_lps.GetPixelID())
+    out_on_input = res.Execute(out_lps)
+
+    gt_on_input = None
+    if src == "Sample":
+        gt_path = SAMPLE_MAP[sample_key].get("gt", None)
+        if gt_path and os.path.exists(gt_path):
+            gt_img = sitk.DICOMOrient(sitk.ReadImage(gt_path), "LPS")
+            res.SetReferenceImage(in_lps)
+            res.SetInterpolator(sitk.sitkLinear)
+            res.SetOutputPixelType(gt_img.GetPixelID())
+            gt_on_input = res.Execute(gt_img)
+
+    # numpy
+    in_vol  = sitk.GetArrayFromImage(in_lps).astype(np.float32)
+    syn_vol = sitk.GetArrayFromImage(out_on_input).astype(np.float32)
+    gt_vol  = sitk.GetArrayFromImage(gt_on_input).astype(np.float32) if gt_on_input is not None else None
+
+    st.subheader("Input vs Synthetic CT Viewer (Axial only)")
+    n_slices = in_vol.shape[0]
+    idx = st.slider("Slice index (Axial/Z)", 0, n_slices - 1, n_slices // 2)
+
+    def get_axial(arr, k):
+        return arr[k, :, :]
+
+    sl_in  = get_axial(in_vol, idx)
+    sl_syn = get_axial(syn_vol, idx)
+    img_in  = _norm2u8(sl_in)
+    img_syn = _norm2u8(sl_syn)
+    img_gt  = _norm2u8(get_axial(gt_vol, idx)) if gt_vol is not None else None
+
+    overlay_mask = st.checkbox("Overlay mask (red)")
+    alpha = st.slider("Mask opacity", 0.0, 1.0, 0.35, 0.05, disabled=not overlay_mask)
+    mask_slice = None
+    if overlay_mask and st.session_state.input_mask is not None:
+        mask_lps = sitk.DICOMOrient(st.session_state.input_mask, "LPS")
+        res_nn = sitk.ResampleImageFilter()
+        res_nn.SetReferenceImage(in_lps)
+        res_nn.SetInterpolator(sitk.sitkNearestNeighbor)
+        mask_on_input = res_nn.Execute(mask_lps)
+        mask_np = sitk.GetArrayFromImage(mask_on_input)
+        mask_slice = get_axial(mask_np, min(idx, mask_np.shape[0]-1))
+        mask_plot = np.where(mask_slice > 0, 1.0, np.nan)
+    else:
+        mask_plot = None
+
+    import plotly.graph_objects as go
+    from plotly.subplots import make_subplots
+    sx, sy, _ = in_lps.GetSpacing()
+    xs = np.arange(img_in.shape[1]) * sx
+    ys = np.arange(img_in.shape[0]) * sy
+
+    cols = 3 if (src == "Sample" and img_gt is not None) else 2
+    titles = ["Input (MRI/CBCT)", "Synthetic CT"] + (["Ground-Truth CT"] if cols == 3 else [])
+    fig = make_subplots(rows=1, cols=cols, subplot_titles=tuple(titles))
+
+    fig.add_trace(go.Heatmap(z=img_in, x=xs, y=ys, colorscale="gray",
+                             zmin=0, zmax=255, showscale=False, hoverinfo="skip"), row=1, col=1)
+    # synCT
+    fig.add_trace(go.Heatmap(z=img_syn, x=xs, y=ys, colorscale="gray",
+                             zmin=0, zmax=255, showscale=False, hoverinfo="skip"), row=1, col=2)
+    # GT
+    if cols == 3:
+        fig.add_trace(go.Heatmap(z=img_gt, x=xs, y=ys, colorscale="gray",
+                                 zmin=0, zmax=255, showscale=False, hoverinfo="skip"), row=1, col=3)
+
+    # mask overlay
+    if mask_plot is not None:
+        red_scale = [[0.0, "rgba(255,0,0,1.0)"], [1.0, "rgba(255,0,0,1.0)"]]
+        for c in range(1, cols+1):
+            fig.add_trace(go.Heatmap(z=mask_plot, x=xs, y=ys, colorscale=red_scale,
+                                     showscale=False, opacity=alpha, hoverinfo="skip"), row=1, col=c)
+
+    for c in range(1, cols+1):
+        fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False, row=1, col=c)
+        fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False, row=1, col=c)
+
+    fig.update_layout(height=600, margin=dict(l=10, r=10, t=40, b=10))
+    st.plotly_chart(fig, use_container_width=True)
+
+    # Caption
+    if cols == 3:
+        st.caption(f"Axial (Z) slice {idx+1}/{n_slices} — All aligned to input geometry; GT only for samples.")
+    else:
+        st.caption(f"Axial (Z) slice {idx+1}/{n_slices} — Aligned to input geometry.")
     col_d1, col_d2, col_d3 = st.columns(3)
 
-
     with col_d3:
-        _download_sitk_image(
-            st.session_state.synth_ct,
-            file_name="synth_ct.nii.gz",
-            label="Download synthetic CT",
-        )
-
+        _download_sitk_image(st.session_state.synth_ct,
+                             file_name="synth_ct.nii.gz",
+                             label="Download synthetic CT")
 
     with col_d1:
         if st.session_state.input_vol is not None:
             in_name = "input_mr.nii.gz" if task == "Task 1 (MR → CT)" else "input_cbct.nii.gz"
             in_label = "Download input MRI" if task == "Task 1 (MR → CT)" else "Download input CBCT"
-            _download_sitk_image(
-                st.session_state.input_vol,
-                file_name=in_name,
-                label=in_label,
-            )
+            _download_sitk_image(st.session_state.input_vol, file_name=in_name, label=in_label)
         else:
             st.button("Download input", disabled=True)
 
     with col_d2:
         if st.session_state.input_mask is not None:
-            _download_sitk_image(
-                st.session_state.input_mask,
-                file_name="input_mask.nii.gz",
-                label="Download input Mask",
-            )
+            _download_sitk_image(st.session_state.input_mask,
+                                 file_name="input_mask.nii.gz",
+                                 label="Download input Mask")
         else:
             st.button("Download input Mask", disabled=True)