debug

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip 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

app.py

@@ -10,25 +10,21 @@ import streamlit as st
 from PIL import Image, ImageDraw
 from huggingface_hub import snapshot_download
 
-# =========================
-# 配置：两个任务的模型仓库 & 本地路径
-# =========================
-# 你可以将两个任务分别指向不同的 HF repo；如果都在同一个，也可以都填同一个。
+
 HF_REPOS = {
-    "Task 1 (MR → CT)": "aehrc/Synthrad2025",
-    "Task 2 (CBCT → CT)": "aehrc/Synthrad2025",  # 如有专门CBCT→CT的repo可在此替换
+    "Task 1 (MR → CT)": "aehrc/Synthrad2025_task1",
+    "Task 2 (CBCT → CT)": "aehrc/Synthrad2025_task2",  
 }
 LOCAL_WEIGHTS_DIRS = {
     "Task 1 (MR → CT)": os.path.abspath("weights/task1"),
     "Task 2 (CBCT → CT)": os.path.abspath("weights/task2"),
 }
 
-# 环境设置（Token）
+
 token = os.getenv("HF_TOKEN")
 if token is None:
     print("[Warn] HF_TOKEN not set. If the model repo is private, set it in Settings → Variables and secrets.")
 
-# 先下载两个任务的权重（如需按需下载，可在选择任务后再下载）
 REPO_DIRS = {}
 for task_name, repo in HF_REPOS.items():
     repo_dir = snapshot_download(
@@ -40,13 +36,11 @@ for task_name, repo in HF_REPOS.items():
     )
     REPO_DIRS[task_name] = repo_dir
 
-# nnUNet 环境变量（指向“当前任务”的 results，会在用户切换任务时动态更新）
+
 os.environ.setdefault("nnUNet_raw", "./nnunet_raw")
 os.environ.setdefault("nnUNet_preprocessed", "./nnunet_preprocessed")
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
-# 从 process.py 导入两个任务的算法类
-# 确保你在 process.py 中定义了 SynthradAlgorithm1（MR→CT）和 SynthradAlgorithm2（CBCT→CT）
 from process import  SynthradAlgorithm2
 
 from process_1 import SynthradAlgorithm1
@@ -56,21 +50,19 @@ from process_1 import SynthradAlgorithm1
 # =========================
 st.set_page_config(page_title="SynthRad (nnUNetv2) Demo", layout="wide")
 st.title("SynthRad — MRI/CBCT + Mask → synthetic CT")
-
-# 任务选择
+st.image("/home/head_neck/sub_2/Synthrad2025/workflow.png",width=800)
 TASKS = ["Task 1 (MR → CT)", "Task 2 (CBCT → CT)"]
 task = st.radio("Select Task", TASKS, index=0, horizontal=True)
 
-# 根据任务设置标题/提示
+
 if task == "Task 1 (MR → CT)":
     vol_label = "MRI volume (.nii/.nii.gz/.mha)"
 else:
     vol_label = "CBCT volume (.nii/.nii.gz/.mha)"
 
-# 动态切换 nnUNet 的 results（不同任务使用不同 results 目录）
+
 os.environ["nnUNet_results"] = REPO_DIRS[task]
 
-# session_state 初始化
 if "algos" not in st.session_state:
     st.session_state.algos = {}
 if "synth_ct" not in st.session_state:
@@ -84,7 +76,6 @@ if "input_vol" not in st.session_state:
 if "input_mask" not in st.session_state:
     st.session_state.input_mask = None
 
-# 懒加载对应任务的算法实例
 def get_algo(task_name: str):
     if task_name not in st.session_state.algos:
         if task_name == "Task 1 (MR → CT)":
@@ -98,25 +89,18 @@ algo = get_algo(task)
 st.subheader("Input")
 src = st.radio("Source", ["Sample", "Upload"], index=0, horizontal=True)
 
-# =========================
-# 样例映射（两任务可共用同一份样例，也可按需区分）
-# 这里假设 repo_dir 下有如下结构：
-#   repo_dir/Abdomen/{cbct.mha, mask.mha} 或 {mri.mha, mask.mha}
-#   repo_dir/Head and Neck/{cbct.mha or mri.mha, mask.mha}
-#   repo_dir/Thorax/{cbct.mha or mri.mha, mask.mha}
-# 如果你的文件名不同，请按需调整。
-# =========================
+
 
 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"    # 如果你的样例文件名不是 mri.mha，请改成实际名称
-        mask_fname = "mask1.mha"    # 如果你的样例文件名不是 mri.mha，请改成实际名称
+        vol_fname = "mr.mha"   
+        mask_fname = "mask1.mha"   
     else:
         vol_key = "cbct"
-        vol_fname = "cbct.mha"   # 如果你的样例文件名不是 cbct.mha，请改成实际名称
-        mask_fname = "mask2.mha"    # 如果你的样例文件名不是 mri.mha，请改成实际名称
+        vol_fname = "cbct.mha" 
+        mask_fname = "mask2.mha"   
     sample_map = {
         "Abdomen (sample)": {
             "region": "Abdomen",
@@ -138,9 +122,7 @@ def build_sample_map(task_name: str):
 
 SAMPLE_MAP = build_sample_map(task)
 
-# =========================
-# 小工具函数
-# =========================
+
 def _download_sitk_image(img: sitk.Image, file_name: str, label: str):
     with tempfile.NamedTemporaryFile(suffix=".nii.gz", delete=False) as tmp:
         sitk.WriteImage(img, tmp.name)
@@ -181,9 +163,6 @@ def _norm2u8(slice2d):
     s = np.clip(s, 0, 1)
     return (s * 255).astype(np.uint8)
 
-# =========================
-# 输入区域（Upload or Sample）
-# =========================
 c1, c2, c3 = st.columns([2, 2, 1])
 
 if src == "Upload":
@@ -208,9 +187,6 @@ else:
 
 run_btn = st.button("Run", type="primary", disabled=not inputs_ready)
 
-# =========================
-# 推理
-# =========================
 if run_btn:
     with st.spinner(f"Running nnUNetv2 {('SynthradAlgorithm1' if task=='Task 1 (MR → CT)' else 'SynthradAlgorithm2')}..."):
         if src == "Upload":
@@ -221,15 +197,13 @@ if run_btn:
             in_vol_img = _read_sitk_from_path(sample["vol"])
             mask_img = _read_sitk_from_path(sample["mask"])
 
-        # 保存原始元信息
+
         st.session_state.orig_meta = (
             in_vol_img.GetSpacing(),
             in_vol_img.GetOrigin(),
             in_vol_img.GetDirection(),
         )
 
-        # 调用不同任务的算法
-        # 约定：算法统一使用 dict 输入：{"image": <sitk.Image>, "mask": <sitk.Image>, "region": <str>}
         out_img = algo.predict({"image": in_vol_img, "mask": mask_img, "region": region_for_run})
 
         st.session_state.synth_ct = out_img
@@ -237,20 +211,17 @@ if run_btn:
         st.session_state.input_vol = in_vol_img
         st.session_state.input_mask = mask_img
 
-# =========================
-# 结果与下载
-# =========================
 if st.session_state.vol_np is None:
-    st.info("请先选择任务与输入（Upload 或 Sample），然后点击 Run")
+    st.info("Select Upload or Sample, then click Run")
 else:
-    # 将输出转为 LPS 方向做显示（可选）
+
     out_lps = sitk.DICOMOrient(st.session_state.synth_ct, "LPS")
     vol = sitk.GetArrayFromImage(out_lps).astype(np.float32)
     D, H, W = vol.shape
 
     col_d1, col_d2, col_d3 = st.columns(3)
 
-    # 下载合成CT
+
     with col_d3:
         _download_sitk_image(
             st.session_state.synth_ct,
@@ -258,7 +229,7 @@ else:
             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"
@@ -271,7 +242,6 @@ else:
         else:
             st.button("Download input", disabled=True)
 
-    # 下载掩膜
     with col_d2:
         if st.session_state.input_mask is not None:
             _download_sitk_image(

process_1.py

@@ -24,10 +24,11 @@ import shutil
 import os
 
 os.environ["OPENBLAS_NUM_THREADS"] = "1"
+force_cpu = os.getenv("FORCE_CPU", "0") == "1"
 device = torch.device("cuda:0" if torch.cuda.is_available() and not force_cpu else "cpu")
 
 
-force_cpu = os.getenv("FORCE_CPU", "0") == "1"
+
 class SynthradAlgorithm1(BaseSynthradAlgorithm):
     """
     This class implements a simple synthetic CT generation algorithm that segments all values greater than 2 in the input image.