nn unet批量处理python代码求助

下面给你两套方案:
A) 完整的手动操作步骤(“先跑通一次”)
B) 一键脚本(在 Slicer 的 Python 控制台里运行,自动批量加载 nii 并调用 nnUNet 分割与保存)


A) 先手动跑通一次(图形界面)

假设你已安装好 NNUNet 扩展(扩展名:SlicerNNUNet),Slicer 版本 5.8.x。

  1. 加载数据(nii)
  • 菜单:File → Add Data… → Add File(s)…
  • 进入你的目录:
    /Volumes/T7 Shield/data2412all/12_20/abnormal_t2w
  • 多选所有 .nii.nii.gz 文件(或先选一例用于测试),勾选 “Show Options” 确认 Single File,点击 OK 载入。
  • 载入后可在 Data 模块看到每个 t2w 体积(Scalar Volume)。
  1. 通过 Module Finder 找到 nnUNet 并进入
  • 按快捷键 Ctrl/Cmd + F 打开 Module Finder(或点工具栏放大镜图标)。
  • 搜索:nn UNetSlicerNNUNet
  • 点搜索结果里的 nn UNet,再点 Switch to module 进入模块。
  1. 在 nnUNet 模块中设置参数并运行
  • Input volume:下拉选择你刚载入的某个 t2w 体积(先选一例)

  • Model Path:设置为
    /Users/<你的用户名>/nnuet weight/Dataset850_TotalSegMRI_part1_organs_1088subj

    <你的用户名> 换成实际用户名;路径里有空格没关系。

  • 其他参数(如果模块有显示 folds、plans、gpu/cpu 等),保持默认或按需调整。

  • 点击 Apply。等待预测完成(进度条走完/状态变成完成)。

  • 完成后场景里会新增一个 Segmentation 节点(通常以输入名派生)。你可以到 Segmentations / Segment Editor 查看或保存(File → Save)。

先用 1 例试跑成功,再进行批量更稳妥。


B) 一键脚本:批量加载 nii → 切换 nnUNet → 逐个应用 → 批量保存

用法:打开 Slicer → View → Python Interactor,把下面脚本 一次性 粘贴进去回车。
它会:

  • 扫描并载入 /Volumes/T7 Shield/data2412all/12_20/abnormal_t2w 下所有 .nii/.nii.gz
  • 自动切到 nnUNet 模块
  • Model Path 设为你提供的目录
  • 逐个把体积放到 Input volume 并点击 Apply 运行
  • 监控按钮状态,等待本例结束后再处理下一例
  • 把分割结果保存到你指定的输出目录(同名 .seg.nrrd

注意:不同版本的 SlicerNNUNet UI 的控件 objectName 可能不同。脚本里包含了“自动探测控件”的逻辑——它会在模块界面里寻找带 “Apply” 文本的按钮、包含 “Model” 字样的路径框/按钮,以及 “Input” 的体积选择器;如果第一次运行没找到,会打印可用控件名称,按提示把 LIKELY_* 名称改成你界面里的控件名再运行即可。

# === nnUNet 批处理脚本(Slicer 5.8.x)===
import os, glob, time
import qt, slicer
import slicer.util as su

# 1) 路径配置
data_dir = r"/Volumes/T7 Shield/data2412all/12_20/abnormal_t2w"
model_path = r"/Users/<YOUR_USERNAME>/nnuet weight/Dataset850_TotalSegMRI_part1_organs_1088subj"  # ← 改成你的用户名
out_dir   = os.path.join(data_dir, "_nnunet_out")
os.makedirs(out_dir, exist_ok=True)

# 2) 载入所有 nii / nii.gz
nii_paths = sorted(glob.glob(os.path.join(data_dir, "*.nii"))) + \
            sorted(glob.glob(os.path.join(data_dir, "*.nii.gz")))
if not nii_paths:
    raise RuntimeError(f"未在目录中找到 NIfTI 文件: {data_dir}")

loaded_volumes = []
for p in nii_paths:
    n = su.loadVolume(p)
    if n:
        loaded_volumes.append((p, n))
print(f"已载入 {len(loaded_volumes)} 个体积。")

# 3) 进入 nnUNet 模块
#   模块内部名:SlicerNNUNet(通常属性名为 slicer.modules.slicernnunet)
moduleName = "SlicerNNUNet"
su.selectModule(moduleName)
module = getattr(slicer.modules, moduleName.lower(), None)
if module is None:
    raise RuntimeError("未找到 SlicerNNUNet 模块。请确认已在扩展管理器中安装并重启。")
widget = module.widgetRepresentation()
if widget is None:
    raise RuntimeError("未能获取 SlicerNNUNet 模块的 widgetRepresentation。")

# 4) 帮助:列出子控件名字(若找控件失败,可查看输出后手动替换下面的 LIKELY_* 变量)
def list_children(w):
    out = []
    def _recurse(obj):
        for c in obj.children():
            try:
                name = c.objectName
            except:
                name = ""
            text = ""
            if hasattr(c, "text"):
                try:
                    text = c.text
                except:
                    pass
            out.append((c.__class__.__name__, name, text))
            _recurse(c)
    _recurse(w)
    return out

# 5) 智能查找控件(尽量不依赖固定 objectName)
def find_button_by_text(root, text_contains="Apply"):
    for klass, name, text in list_children(root):
        if klass == "QPushButton" and isinstance(text, str) and text_contains.lower() in text.lower():
            return su.findChild(root, name)
    return None

def find_lineedit_like_for_model(root):
    # 可能是 PathLineEdit / ctkPathLineEdit / QLineEdit 或带 DirectoryButton 的组合
    candidates = []
    for klass, name, text in list_children(root):
        if any(k in klass.lower() for k in ["pathline", "ctkpath", "lineedit"]):
            if "model" in (name or "").lower() or "model" in (text or "").lower():
                c = su.findChild(root, name)
                if c: candidates.append(c)
    return candidates[0] if candidates else None

def find_directory_button_near_model(root):
    # 找一个和 "Model" 相关的目录按钮(QToolButton/ctkDirectoryButton)
    for klass, name, text in list_children(root):
        if any(k in klass.lower() for k in ["toolbutton", "pushbutton", "button"]):
            if "model" in (name or "").lower():
                return su.findChild(root, name)
    return None

def find_input_selector(root):
    # qMRMLNodeComboBox 或其派生,name/text 含 "input"
    for klass, name, text in list_children(root):
        if "qmrmlnodecombobox" in klass.lower():
            if "input" in (name or "").lower() or "input" in (text or "").lower():
                return su.findChild(root, name)
    return None

# 尝试自动找
applyButton = find_button_by_text(widget, "Apply")
inputSelector = find_input_selector(widget)
modelPathEdit = find_lineedit_like_for_model(widget)  # 优先用可 setText 的
modelDirButton = find_directory_button_near_model(widget)

# 若有找不到的,打印所有控件供人工查看
if any(x is None for x in [applyButton, inputSelector]) or (modelPathEdit is None and modelDirButton is None):
    print("\n[诊断] 未能自动识别全部控件。下面列出子控件以便你确认 objectName:")
    for klass, name, text in list_children(widget):
        if name or text:
            print(f"{klass:28s}  name='{name}'  text='{text}'")
    raise RuntimeError("请根据上方列表,修正脚本中控件查找逻辑(见 LIKELY_* 注释)。")

# 设置 Model Path
def set_model_path(path):
    if modelPathEdit and hasattr(modelPathEdit, "setText"):
        modelPathEdit.setText(path)
        qt.QApplication.processEvents()
    elif modelDirButton:
        # 某些界面只提供“选择目录”按钮。弹对话框不易脚本化,提示用户手动设置一次后再继续。
        raise RuntimeError("当前界面通过按钮选择 Model Path,脚本无法直接填入。请手动在界面设置好 Model Path 后重新运行循环部分。")
    else:
        raise RuntimeError("无法设置 Model Path。")

set_model_path(model_path)
print(f"[OK] Model Path 已设为:{model_path}")

# 6) 逐个分割
def wait_until_apply_finished(btn, poll=0.5, timeout_sec=3600):
    # 经验:多数模块运行中会禁用 Apply(或文本变化)。这里以 enabled 状态作为简单判断。
    t0 = time.time()
    while True:
        qt.QApplication.processEvents()
        time.sleep(poll)
        # 条件:按钮再次可用 → 视为完成
        if btn.isEnabled():
            break
        if time.time() - t0 > timeout_sec:
            raise TimeoutError("等待 nnUNet 预测超时。")
    time.sleep(0.2)
    qt.QApplication.processEvents()

# 获取/设置 Input volume
def set_input_volume(node):
    if hasattr(inputSelector, "setCurrentNode"):
        inputSelector.setCurrentNode(node)
    elif hasattr(inputSelector, "setCurrentNodeID"):
        inputSelector.setCurrentNodeID(node.GetID())
    else:
        raise RuntimeError("无法设置 Input volume(未知的选择器类型)。")
    qt.QApplication.processEvents()

# 保存分割(寻找最新的 Segmentation 节点)
def get_newest_segmentation_node():
    segNodes = su.getNodesByClass('vtkMRMLSegmentationNode')
    if not segNodes:
        return None
    # 返回最近创建/修改的(简单用场景内顺序或按名称匹配输入名)
    return segNodes[-1]

def save_segmentation(segNode, refPath):
    base = os.path.basename(refPath)
    stem = os.path.splitext(os.path.splitext(base)[0])[0]  # 兼容 .nii.gz
    outPath = os.path.join(out_dir, f"{stem}.seg.nrrd")
    ok = su.saveNode(segNode, outPath)
    if ok:
        print(f"[Saved] {outPath}")
    else:
        print(f"[Warn] 保存失败:{outPath}")

# 主循环
for i, (p, volNode) in enumerate(loaded_volumes, 1):
    print(f"\n=== [{i}/{len(loaded_volumes)}] {os.path.basename(p)} ===")
    # 设输入体积
    set_input_volume(volNode)

    # 点 Apply(若当前已禁用,先处理事件队列)
    if not applyButton.isEnabled():
        qt.QApplication.processEvents()
        time.sleep(0.2)
    # 触发
    applyButton.click()
    qt.QApplication.processEvents()

    # 等待完成
    # 运行开始后,许多模块会临时禁用 Apply;若未禁用,也可加固定等待,但这里先按 enabled 轮询
    applyButton.setEnabled(False)  # 主动置为 False,以确保 wait 逻辑可生效(有些界面不会变)
    qt.QApplication.processEvents()
    wait_until_apply_finished(applyButton, poll=1.0, timeout_sec=7200)

    # 找到分割并保存
    seg = get_newest_segmentation_node()
    if seg:
        save_segmentation(seg, p)
    else:
        print("[Warn] 未发现新的分割节点,请在场景中检查。")

print("\n[Done] 全部处理完成。输出目录:", out_dir)

小贴士

  • 如果 Model Path 只能通过“选择目录按钮”设置,脚本会提醒你:请手动在界面里先选好一次,然后把 set_model_path(model_path) 换成一个空函数(或注释掉),再重新运行“主循环”那段(从 for i, (p, volNode) in enumerate... 开始)。
  • 如果脚本提示找不到某控件,它会把所有子控件的 class/name/text 打印出来。根据输出把 find_* 函数里匹配规则稍微改一下就行(常见变化是 objectName 不同)。
  • 初次运行 nnUNet 可能会弹出 Install/Download dependencies 的提示(依赖安装)。如果出现,请先在 GUI 里把依赖安装好,再跑脚本。
  • 输出为 .seg.nrrd,你也可以改成 .seg.nii.gz(先用 slicer.util.saveNode 直接保存 .seg.nrrd 最稳)。
  • 如果你想把分割另存为 labelmap模型(surface .ply/.stl),可在保存前用 slicer.modules.segmentations.logic() 做导出,告诉我我可以把这段也补上。

常见问题排查

  • 没找到 nnUNet 模块:打开 Extensions Manager 搜索安装 “NNUNet”(Internal name: SlicerNNUNet),重启 Slicer。
  • Apply 一直灰/没反应:检查 Input volume 是否选择正确、Model Path 是否指向含权重与配置的 nnUNet 模型目录;查看底部 Python Console 是否有缺依赖/找不到权重的报错。
  • 速度慢:macOS 上若无 GPU/或没配置 CUDA,将会用 CPU;可先用少量数据验证流程。
  • 路径含空格:脚本里用原始字符串 r"…",Slicer 内部可以正常处理。

代码如出现问题,需要交互操作几次后,大模型才能给出正确代码。