wasserth / TotalSegmentator

Tool for robust segmentation of >100 important anatomical structures in CT and MR images
Apache License 2.0
1.44k stars 242 forks source link

Output file type #261

Closed josephrich98 closed 1 month ago

josephrich98 commented 8 months ago

Is there a way to convert the output nifti file to a BMP binary segmentation mask? I am interested in acquiring a binary mask of a 3D input image, as well as the segmentation mask of each 2D slice comprising this 3D image. Thank you!

wasserth commented 8 months ago

Certainly! You can accomplish this with a Python script using the nibabel library to read the nifti file and Pillow (PIL) or matplotlib to save each slice as a BMP image. Here's a step-by-step guide:

Step 1: Read the nifti file

import numpy as np
import nibabel as nib
from PIL import Image

# Load the nifti file
nifti_file_path = 'path/to/your/file.nii'
nii = nib.load(nifti_file_path)
data = nii.get_fdata()

Step 2: Normalize and save the segmentation mask as BMP

# Assume data is a 3D numpy array with shape (height, width, depth)
for i in range(data.shape[2]):
    # Process each 2D slice
    slice_2d = data[:, :, i]

    # Normalize the slice for visualization and convert to uint8
    slice_normalized = ((slice_2d - slice_2d.min()) * (255.0 / (slice_2d.max() - slice_2d.min()))).astype(np.uint8)

    # Create a binary mask where the segmentation label is 1 and everything else is 0
    # Let's assume label `1` corresponds to the segmentation you're interested in
    binary_mask = (slice_normalized == 1).astype(np.uint8) * 255

    # Save the binary mask as a BMP image
    img = Image.fromarray(binary_mask)
    img.save(f'output_directory/slice_{i}.bmp')

Step 3: Save the combined segmentation mask

If you also need to combine the entire 3D segmentation into one 2D image (where each segment is represented by a unique color or shade), do this:

# Let's assume there are up to 5 different segmentations labeled with values from 1 to 5
combined_mask = np.zeros(data.shape[0:2], dtype=np.uint8)

for i in range(1, 6):  # Labels 1 to 5
    combined_mask[data == i] = i * (255 // 5)  # Assign a unique shade based on the label

# Convert to an image and save
combined_image = Image.fromarray(combined_mask)
combined_image.save('combined_segmentation_mask.bmp')

Make sure to replace path/to/your/file.nii with the actual path of your nifti file and adjust the label numbers according to your specific use case. Also, ensure the output directory where the BMP files are saved exists or is created by your script.

Remember to validate if the label values (assumed as 1 in this script) match the ones in your output nifti file before running this code. Each label you're interested in must be handled correctly to create accurate binary masks.

josephrich98 commented 8 months ago

Thank you so much! This is very helpful for my work. I really appreciate your thoroughness!