使用 VTK 将 mask 提取成 mesh

1. 前言

网上使用 VTK 将 Mask 数据生成 Mesh 的方法介绍有很多，但是直接简单的代码的反倒不多。
因此，写一个简单的 Python 片段，批量讲 mask 生成 mesh, 并保存起来。

2. 代码实现的功能

代码片段实现的功能：

• 文件夹中查找特定后缀的文件
• 将 mask 生成 mesh, 并保存下来

3. 代码如下

import os
import time
import vtk


def find_files_with_suffix(target_dir, target_suffix="txt"):
   """ 查找以 target_suffix 为后缀的文件，并返回 """
   find_res = []
   target_suffix_dot = "." + target_suffix
   walk_generator = os.walk(target_dir)
   for root_path, dirs, files in walk_generator:
       if len(files) < 1:
           continue
       for file in files:
           file_name, suffix_name = os.path.splitext(file)
           if suffix_name == target_suffix_dot:
               find_res.append(os.path.join(root_path, file))
   return find_res


def convert_mask_to_mesh_with_marching_cube(input_path):
   """ 将 mask 转成 mesh, mask 必须是 0 - 1 的 uchar 类型 """
   assert os.path.exists(input_path), r'文件必须存在 {}'.format(input_path)

   # 读取 mask 变成 image data
   if input_path.find('.mhd'):
       image_reader = vtk.vtkMetaImageReader()
   elif input_path.find('.nii'):
       image_reader = vtk.vtkNIFTIImageReader()

   image_reader.SetFileName(input_path)
   image_reader.Update()

   # 提取 mesh
   mc = vtk.vtkMarchingCubes()
   mc.SetInputData(image_reader.GetOutput())
   mc.SetValue(0, 0.5)
   mc.ComputeNormalsOn()
   mc.Update()

   # 保存成 obj 格式
   writer = vtk.vtkOBJWriter()
   writer.SetInputData(mc.GetOutput())
   writer.SetFileName(r'{}.obj'.format(input_path))
   writer.Update()


if __name__ == '__main__':

   mask_folder = r'C:\mask'
   files = find_files_with_suffix(mask_folder, target_suffix='mhd')

   for file in files:
       print(r'begin process: {}'.format(file))
       convert_mask_to_mesh_with_marching_cube(file)

   print('end')

4. 总结

上面的代码，只需要装好 VTK 9.0， 基本就可以直接使用了。