nn unet批量处理python代码求助

软件版本
问题描述:我想使用3dslicer软件中的nn UNet插件来对一批数据进行膀胱MRI分割。我的膀胱数据路径为/Volumes/T7 Shield/data2412all/12_20/abnormal_t2w。 我的膀胱mri的数据格式为nii。请先帮我把数据加载到slicer软件中。

然后通过module finder的搜索框查找nn UNet,查到nn UNet后,选择nn UNet后再点击switch to module。成功进入nn UNet后,input volume中选择加载的t2w图像文件。然后使nn UNet 的Model Path为/Users/**/nnuet weight/Dataset850_TotalSegMRI_part1_organs_1088subj。然后点击apply。等到slicer软件对图像进行分割。

此外我还发现以下内容

nnUNet

Category: Segmentation

This extension is meant to streamline the integration of nnUnet based models into 3D Slicer.

It allows for quick and reliable nnUNet dependency installation in 3D Slicer environment and provides simple logic to launch nnUNet prediction on given directories.

The installation steps are based on the work done in the Slicer Total Segmentator extension

This module was originally co-financed by the Fédération Française d’Orthodontie (FFO) as part of the Dental Segmentator developments and the Cure Overgrowth Syndromes (COSY) RHU Project.

Contributors: Thibault Pelletier (Kitware SAS)

Internal name: SlicerNNUNet

Type: Python Scripted Loadable

Location: /Applications/Slicer.app/Contents/Extensions-33241/NNUNet/lib/Slicer-5.8/qt-scripted-modules/SlicerNNUNet.py

segment statisticscsv/Volumes/T7 Shield/data2412all/12_20/result

下面给你两套方案:
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 内部可以正常处理。

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