Skip to main content

nnU-Net

AI Segmentation 📥 Download Script

nn-UNet segmentation tutorial script.

Performs nn-UNet segmentation on the first volume in the opened project, splits a multilabel result into individual masks, keeps the largest label, filters to the largest connected region, generates a 3D preview, converts it to a surface, and performs simple surface cleanup and export.

Complete Example​

# Import necessary modules
import ScriptingApi as api

# Instantiate the Application class
# The Application class serves as the primary entry point for interacting with
# the software. To begin using the API, create an instance of this class.
# Use the original class name `Application` or the pythonic alias `app` or `application`
# as the variable name to improve IntelliSense support. For details, refer to the documentation.
app = api.Application()

# Close any open project
app.close_project()

# Open a specific project
project_path = R'C:\Data\MyProject.vvcx'
app.open_project(project_path)

# Get the name of the first volume
volumes = app.get_all_volume_names()
if len(volumes) == 0:
   raise RuntimeError("No volumes found in project.")
volume_name = volumes[0]
print("Using volume:", volume_name)

# Get AiSegmentation instance
# Use the original class name `AiSegmentation` or the pythonic alias `ai_segmentation`
# as the variable name to improve IntelliSense support.
# Hover over the variable to view available methods and properties.
# If no tooltip appears, IntelliSense does not recognize the variable name.
# For details, refer to the documentation.
AiSegmentation = app.get_ai_segmentation()

# Set the AI segmentation model type
AiSegmentation.set_model_type(api.AiSegmentationModelType.NnUnet)

# Check installation status and install if necessary
installed = AiSegmentation.get_installation_status()
if not installed:
    print('Installation required: The AI model is not installed.')
    print('Please run the application as an administrator to install it.')
    print('Attempting to install the AI model now.')
    installed = AiSegmentation.install_model()
    if not installed:
        raise RuntimeError('An error occurred during the installation of the AI model. The log file has been saved for troubleshooting.')
    else:
        print('The AI model is installed and ready for use.')
else:
    print('The AI model is installed and ready for use.')

# Retrieve and print available tasks
available_tasks = AiSegmentation.get_available_tasks()
print(available_tasks)

# Configure nnU-Net parameters
NnUnetParams = AiSegmentation.get_default_nn_unet_params()
NnUnetParams.models_dir = 'C:/Users/furqa/.totalsegmentator/nnunet/results'
NnUnetParams.dataset = 'Dataset298_TotalSegmentator_total_6mm_1559subj'
NnUnetParams.configuration = '3d_fullres'
NnUnetParams.folds = [0]

# Windows stability recommendation for nnU-Net export workers:
# 0/0 enables sequential mode and avoids multiprocessing IPC failures seen with some models.
NnUnetParams.npp = 0
NnUnetParams.nps = 0

# Prefer GPU when available. nnU-Net may still move large result arrays to CPU if GPU memory is insufficient.
NnUnetParams.device = 'gpu'

# Run nnU-Net segmentation
mask_names = AiSegmentation.run_nn_unet([volume_name], NnUnetParams)
if len(mask_names) == 0:
    raise RuntimeError("Segmentation returned no masks (operation cancelled or failed).")
print("Generated masks:", mask_names)

# Get the name of the currently active mask
active_mask = app.get_active_mask_name()
print("Active mask:", active_mask)

# Get the instance of MaskOperations to perform operations on masks
MaskOperations = app.get_mask_operations()

# Split multi-label mask into individual masks and extract only the largest label by volume
splitted_masks = MaskOperations.split_multi_label_mask(volume_name, active_mask, True, 1)
if len(splitted_masks) == 0:
   raise RuntimeError("Split operation did not produce separate labels.")
print("Splitted masks (largest label retained):", splitted_masks)

# Retain only the largest connected region
MaskOperations.filter_regions(splitted_masks, True, 1, 8)

# Show only the splited mask, hide the original multi-label mask
app.set_masks_visible([active_mask], False)
app.set_masks_visible(mask_names, True)

# Set 3D preview quality and generate preview of all visible masks
app.set_mask_3d_preview_quality(api.Mask3dPreviewQuality.Optimal)
app.generate_mask_3d_preview(app.get_visible_mask_names())

# Convert generated 3D preview to surface objects
surfaces = MaskOperations.convert_3d_preview_to_surface_objects()
if len(surfaces) == 0:
    raise RuntimeError("No surface generated from 3D preview.")
print("Created surfaces:", surfaces)

# Get the name of the currently active surface
surface = app.get_active_surface_name()
print(surface)

# Get the instance of SurfaceOperations to perform operations on surfaces
SurfaceOperations = app.get_surface_operations()

# Filter surface to keep only the largest shell
SurfaceFilterShellsParams = api.SurfaceFilterShellsParams()
SurfaceFilterShellsParams.largest_shells = 1
SurfaceOperations.filter_shells(surface, api.SurfaceFilterShellsMethod.LargestShells, SurfaceFilterShellsParams)

# Remesh the surface for better quality
SurfaceOperations.remesh([surface])

# Fix surface issues using diagnostics checks
SurfaceDiagnosticsChecks = api.SurfaceDiagnosticsChecks()
SurfaceOperations.fix_surface(surface, SurfaceDiagnosticsChecks)

# Hide all volumes and masks to focus on the surface
app.set_volumes_visible(app.get_all_volume_names(), False)
app.set_masks_visible(app.get_all_mask_names(), False)

# Isolate the surface
app.isolate_surfaces([surface])

# Export the surface to disk in STL format
SurfaceOperations.export_surface_to_disk(surface, 'C:/tmp/surface.stl')

# Set surface representation to Solid with Edges
SurfaceRenderPropertiesOperations = SurfaceOperations.get_render_properties_operations()
SurfaceRenderPropertiesOperations.set_representation([surface], api.SurfaceRepresentation.SolidEdges)

print("nn-UNet tutorial completed successfully.")