Working with MultiLabelSegmentation

This notebook shows how to create, transfer, and modify multi-label segmentations between Python and a running MITK Workbench.

You will learn how to:

  • Build a MultiLabelSegmentation from scratch (groups, labels, pixel data)

  • Upload it via node.set_data() and download it back via node.get_data()

  • Inspect and edit label metadata (name, color, visibility)

  • Dilate a label region using a SimpleITK filter (optional)

  • Remove a label including its pixels via remove_label()

  • Validate consistency and push the modified segmentation back

Prerequisites: A running MITK Workbench instance with the REST API enabled.

1. Connect to a running Workbench

[1]:
import numpy as np
import mitk_workbench_remote as mw
from mitk_workbench_remote import Label, Image, MultiLabelSegmentation, LABEL_DTYPE

instances = mw.discover(timeout=2) #longer timeouts to make the example more robust for latencies.
if not instances:
    raise SystemExit("No running MITK Workbench found. Start one and try again.")

wb = instances[0]
print(f"Connected to: {wb.info.name} (MITK {wb.info.mitk_version}) at {wb.url}")
Connected to: MITK Workbench REST API (MITK 2025.12.99-23b4e629) at http://localhost:8080

2. Build a synthetic segmentation

MultiLabelSegmentation organises labels into named groups. Label values are globally unique across all groups; value 0 is reserved as the unlabeled background.

Here we create two groups:

  • Anatomy — two anatomical structures (Liver, Spleen)

  • Findings — one pathological structure (Tumor)

Each group’s pixel data is a separate 3-D integer array stored as an Image. Voxels that belong to a label carry its integer value; unlabeled voxels are 0.

[2]:
SHAPE = (80, 80, 80)      # voxels
SPACING = (2.0, 2.0, 2.0) # mm per voxel

ref_img = Image(np.zeros(SHAPE, dtype=LABEL_DTYPE), spacing=SPACING)
seg = MultiLabelSegmentation.create(shape=SHAPE, spacing=SPACING)

# --- Group 0: Anatomy ---
g_anatomy = seg.add_group("Anatomy")
liver  = seg.add_label(Label(None, "Liver",  color=(0.82, 0.40, 0.12)), group=g_anatomy)
spleen = seg.add_label(Label(None, "Spleen", color=(0.30, 0.60, 0.80)), group=g_anatomy)

# --- Group 1: Findings ---
g_findings = seg.add_group("Findings")
tumor = seg.add_label(Label(None, "Tumor", color=(1.0, 0.1, 0.1), locked=True), group=g_findings)

print(seg)
print(f"  Liver  value={liver.value}")
print(f"  Spleen value={spleen.value}")
print(f"  Tumor  value={tumor.value}")
<MultiLabelSegmentation groups=2 labels=3>
  Liver  value=1
  Spleen value=2
  Tumor  value=3
[3]:
seg
[3]:
ValueNameColorVisibleLocked
Anatomy
1LiverTrueFalse
2SpleenTrueFalse
Findings
3TumorTrueTrue
[4]:
def draw_sphere(arr: np.ndarray, center: tuple, radius: float, value: int) -> None:
    """Paint a filled sphere into arr at center with the given integer label value."""
    cx, cy, cz = center
    x, y, z = np.ogrid[:arr.shape[0], :arr.shape[1], :arr.shape[2]]
    arr[(x - cx)**2 + (y - cy)**2 + (z - cz)**2 <= radius**2] = value

# --- Paint ref_img ---
draw_sphere(ref_img.array, center=(40, 38, 40), radius=40, value=100)

# --- Paint Anatomy group ---
arr_anatomy = np.zeros(SHAPE, dtype=LABEL_DTYPE)
draw_sphere(arr_anatomy, center=(40, 38, 40), radius=18, value=liver.value)
draw_sphere(arr_anatomy, center=(58, 50, 35), radius=10, value=spleen.value)
seg.set_group_image(g_anatomy, arr_anatomy)

# --- Paint Findings group ---
arr_findings = np.zeros(SHAPE, dtype=LABEL_DTYPE)
draw_sphere(arr_findings, center=(36, 44, 42), radius=6, value=tumor.value)
seg.set_group_image(g_findings, arr_findings)

issues = seg.validate()
print("Consistency check:", issues if issues else "OK")

liver_px  = int((arr_anatomy == liver.value).sum())
spleen_px = int((arr_anatomy == spleen.value).sum())
tumor_px  = int((arr_findings == tumor.value).sum())
print(f"Voxel counts — Liver: {liver_px}, Spleen: {spleen_px}, Tumor: {tumor_px}")
Consistency check: OK
Voxel counts — Liver: 23795, Spleen: 4169, Tumor: 925

3. Upload to the Workbench

MITK’s Segmentation view expects a segmentation to be a child node of its reference image in the DataStorage tree. We therefore upload in two steps:

  1. Create a node for the reference image and upload it.

  2. Create the segmentation node as a child of the image node by passing parent=img_node.

set_data() serialises each object to NRRD and transfers it. The transfer mode — direct HTTP body or local file reference for large data on localhost — is negotiated automatically.

After uploading, wb.reinit([img_node]) fits the render windows to the image geometry.

[5]:
# Upload the reference image
img_node = wb.storage.create("Demo CT")
img_node.set_data(ref_img)

# Upload the segmentation as a child of the image node.
# MITK uses the parent-child relationship to link a segmentation to its reference image,
# which enables the Segmentation view to display both together correctly.
seg_node = wb.storage.create("Demo Segmentation", parent=img_node)
seg_node.set_data(seg)

wb.reinit()

print(f"Uploaded '{img_node.name}' (uid={img_node.uid}, type={img_node.data_type})")
print(f"Uploaded '{seg_node.name}' (uid={seg_node.uid}, type={seg_node.data_type})")
print(f"Children of img_node: {[n._name for n in img_node.children]}")
Uploaded 'Demo CT' (uid=node_1, type=Image)
Uploaded 'Demo Segmentation' (uid=node_2, type=MultiLabelSegmentation)
Children of img_node: ['Demo Segmentation']
[6]:
seg_node
[6]:
uidnode_2
nameDemo Segmentation
data_typeMultiLabelSegmentation
path/Demo CT/Demo Segmentation
parent_uidnode_1

4. Download and inspect the segmentation

get_data() returns a MultiLabelSegmentation with complete pixel data and label metadata as stored in the Workbench. The rich HTML representation below lists every group and label with its value, color swatch, visibility, and lock state — a quick sanity check that the roundtrip preserved all metadata correctly.

[7]:
seg = seg_node.get_data()
print(seg)

for i, group in enumerate(seg.groups):
    print(f"\nGroup {i}: {group.name!r}")
    for lbl in seg.get_group_labels(i):
        img = seg.get_group_image(i)
        voxels = int((img.array == lbl.value).sum())
        print(f"  value={lbl.value:2d}  name={lbl.name!r:10s}  "
              f"color={lbl.color}  locked={lbl.locked}  voxels={voxels}")
<MultiLabelSegmentation groups=2 labels=3>

Group 0: 'Anatomy'
  value= 1  name='Liver'     color=(0.8199999928474426, 0.4000000059604645, 0.11999999731779099)  locked=False  voxels=23795
  value= 2  name='Spleen'    color=(0.30000001192092896, 0.6000000238418579, 0.800000011920929)  locked=False  voxels=4169

Group 1: 'Findings'
  value= 3  name='Tumor'     color=(1.0, 0.10000000149011612, 0.10000000149011612)  locked=True  voxels=925
[8]:
seg
[8]:
ValueNameColorVisibleLocked
Anatomy
1LiverTrueFalse
2SpleenTrueFalse
Findings
3TumorTrueTrue

5. Edit label metadata

Label properties such as color, visibility, and name are stored in the MultiLabelSegmentation object in Python. They are only written back to the Workbench when you call node.set_data() again (see step 8). Unlike node display properties, there is no live REST setter for individual label fields.

[9]:
# Retrieve labels from the downloaded segmentation by value
liver_lbl  = seg.get_label(1)
spleen_lbl = seg.get_label(2)
tumor_lbl  = seg.get_label(3)

# Adjust Liver color and make Spleen semi-transparent in MITK's property sense
liver_lbl.color   = (0.10, 0.85, 0.0)   # greenish
spleen_lbl.visible = False                 # hide Spleen for now
tumor_lbl.name    = "Primary Tumor"        # rename

print(f"Liver  color:   {liver_lbl.color}")
print(f"Spleen visible: {spleen_lbl.visible}")
print(f"Tumor  name:    {tumor_lbl.name!r}")
Liver  color:   (0.1, 0.85, 0.0)
Spleen visible: False
Tumor  name:    'Primary Tumor'
[10]:
seg
[10]:
ValueNameColorVisibleLocked
Anatomy
1LiverTrueFalse
2SpleenFalseFalse
Findings
3Primary TumorTrueTrue

6. Dilate the tumor label with SimpleITK (optional)

get_group_image(index) returns the group’s pixel data as an Image. We convert it to a SimpleITK image, apply a binary dilation (growing the tumor mask by 2 voxels = 4 mm), and store the result back with set_group_image().

This pattern — download a group image, apply any filter, write back — is the general way to edit pixel content from Python.

If SimpleITK is not installed, this cell is skipped gracefully.

[11]:
try:
    import SimpleITK
except ImportError:
    SimpleITK = None
    print("SimpleITK not installed — skipping dilation.")
    print("Install with: pip install SimpleITK")

if SimpleITK is not None:
    findings_img = seg.get_group_image(g_findings)
    arr = findings_img.array

    # Build a binary mask: 1 where tumor, 0 elsewhere
    binary = (arr == tumor_lbl.value).astype(np.uint8)
    sitk_mask = SimpleITK.GetImageFromArray(binary)
    sitk_mask.SetSpacing([float(s) for s in seg.spacing])

    # Dilate by 2 voxels in each direction
    dilated = SimpleITK.BinaryDilate(sitk_mask, kernelRadius=[5, 5, 5])
    dilated_arr = SimpleITK.GetArrayFromImage(dilated)

    # Write dilated region back as tumor label value; leave other labels intact
    result = arr.copy()
    result[dilated_arr > 0] = tumor_lbl.value
    seg.set_group_image(g_findings, result)

    before = int(binary.sum())
    after  = int(dilated_arr.sum())
    print(f"Tumor before dilation: {before:5d} voxels")
    print(f"Tumor after dilation:  {after:5d} voxels  (+{after - before})")
Tumor before dilation:   925 voxels
Tumor after dilation:   6127 voxels  (+5202)

7. Remove a label including its pixels

remove_label(value, clear_pixels=True) removes a label from the registry and zeroes out every voxel carrying that value in the group’s image. This is the cleanest way to delete a label — no orphaned pixel values remain.

With clear_pixels=False the metadata is removed but the pixel values stay; useful when you intend to reassign them manually.

Here we delete Spleen (value 2) from the Anatomy group.

[12]:
print("Before removal:")
print(f"  Labels in Anatomy: {[(l.value, l.name) for l in seg.get_group_labels(g_anatomy)]}")

spleen_value = spleen_lbl.value  # save before removal
seg.remove_label(spleen_value, clear_pixels=True)

print("\nAfter removal:")
print(f"  Labels in Anatomy: {[(l.value, l.name) for l in seg.get_group_labels(g_anatomy)]}")

# Confirm pixels are gone
arr = seg.get_group_image(g_anatomy).array
remaining = int((arr == spleen_value).sum())
print(f"  Spleen pixels remaining in image: {remaining}  (expected 0)")

# Lookup after removal should return None
print(f"  get_label({spleen_value}) -> {seg.get_label(spleen_value)}")
Before removal:
  Labels in Anatomy: [(1, 'Liver'), (2, 'Spleen')]

After removal:
  Labels in Anatomy: [(1, 'Liver')]
  Spleen pixels remaining in image: 0  (expected 0)
  get_label(2) -> None
[13]:
seg
[13]:
ValueNameColorVisibleLocked
Anatomy
1LiverTrueFalse
Findings
3Primary TumorTrueTrue

8. Validate consistency and push changes back

validate() checks internal consistency — that every pixel value in each group array maps to a registered label, that no value appears in multiple groups, and that the label-value space is free of gaps or overlaps. Run it before set_data() to catch mistakes early.

After validation, set_data(seg) re-serializes the entire segmentation (all pixel groups plus label metadata) and uploads it. There is no partial-update endpoint — any label or pixel edit requires a full re-upload.

[14]:
issues = seg.validate()
if issues:
    print("Consistency warnings:")
    for w in issues:
        print(f"  {w}")
else:
    print("Validation OK — no consistency issues.")

seg_node.set_data(seg)
wb.reinit([img_node])
print(f"\nPushed modified segmentation back to '{seg_node.name}'.")
print("The Workbench now shows:")
for i, group in enumerate(seg.groups):
    labels = seg.get_group_labels(i)
    print(f"  Group {i} '{group.name}': {[l.name for l in labels]}")
Validation OK — no consistency issues.

Pushed modified segmentation back to 'Demo Segmentation'.
The Workbench now shows:
  Group 0 'Anatomy': ['Liver']
  Group 1 'Findings': ['Primary Tumor']

9. Save the segmentation to disk

save_data(path) downloads the raw serialized bytes from the Workbench and writes them directly to a file without any parsing or re-encoding. The result is a standard MITK multilabel NRRD file containing the 4-D pixel array and the full label metadata, readable by MITK Workbench and any compatible tool.

[15]:
from pathlib import Path

output_path = Path("demo_segmentation.nrrd")
saved = seg_node.save_data(output_path)
print(f"Saved to: {saved}")
print(f"File size: {saved.stat().st_size / 1024:.1f} kB")
Saved to: demo_segmentation.nrrd
File size: 9.7 kB

10. Clean up

Remove both nodes from the DataStorage. Passing recursive=True to the parent node (img_node) also removes seg_node, since it is registered as a child in the DataStorage tree.

[16]:
img_node.remove(recursive=True)  # removes img_node and its child seg_node
output_path.unlink(missing_ok=True)
print("Done.")
Done.