How to use MedImageParse healthcare AI models for segmentation of medical images

MedImageParse

Biomedical image analysis is crucial for discovery in fields like cell biology, pathology, and radiology. Traditionally, tasks such as segmentation, detection, and recognition of relevant objects are addressed separately, which can limit the overall effectiveness of image analysis. However, MedImageParse unifies these tasks through image parsing by jointly conducting segmentation, detection, and recognition across numerous object types and imaging modalities. By applying the interdependencies among these subtasks—such as the semantic labels of segmented objects—the model enhances accuracy and enables novel applications. For example, it allows users to segment all relevant objects in an image by using a simple text prompt. This approach eliminates the need to manually specify bounding boxes for each object.

The following image shows the conceptual architecture of the MedImageParse model where an image embedding model is augmented with a task adaptation layer to produce segmentation masks and textual descriptions.

Remarkably, the segmentation masks and textual descriptions are achieved by using only standard segmentation datasets, augmented by natural-language labels, or descriptions harmonized with established biomedical object ontologies. This approach not only improves individual task performance but also offers an all-in-one tool for biomedical image analysis, paving the way for more efficient and accurate image-based biomedical discovery.

MedImageParse 3D

Similar to the MedImageParse model, MedImageParse 3D uses a combination of a text prompt and a medical image to create a segmentation mask. However, unlike MedImageParse, MedImageParse 3D takes in an entire 3D volume—a common way of representing the imaged area for cross-sectional imaging modalities like CT or MRI—and generates the three-dimensional segmentation mask.

import base64
import json
import os

sample_image_xray = os.path.join(image_path)

def read_image(image_path):
    with open(image_path, "rb") as f:
        return f.read()

sample_image =  "sample_image.png"
data = {
    "input_data": {
        "columns": [ "image", "text" ],
        "index": [ 0 ],
        "data": [
            [
                base64.encodebytes(read_image(sample_image)).decode("utf-8"),
                "neoplastic cells in breast pathology & inflammatory cells"
            ]
        ]
    }
}
data_json = json.dumps(data)

# Create request json
request_file_name = "sample_request_data.json"
with open(request_file_name, "w") as request_file:
    json.dump(data, request_file)

response = ml_client_workspace.online_endpoints.invoke(
    endpoint_name=endpoint_name,
    deployment_name=deployment_name,
    request_file=request_file_name,
)

import base64
import json
import urllib.request
import matplotlib.pyplot as plt
import nibabel as nib
import tempfile
import os

# Replace with the path to your NIfTI input file
sample_image = "./examples/amos_0308.nii.gz"
with open(sample_image, "rb") as image_file:
    base64_image = base64.b64encode(image_file.read()).decode('utf-8')

# Prepare data payload
data = {
    "input_data": {
        "columns": [ "image", "text" ],
        "index": [ 0 ],
        "data": [
            [
                base64_image,
                "pancreas"
            ]
        ]
    }
}
data_json = json.dumps(data)
body = str.encode(data_json)

# Add your endpoint URL and API key
url = "<your-endpoint-url>"
api_key = "<your-api-key>"

headers = {
    'Content-Type': 'application/json',
    'Authorization': 'Bearer ' + api_key
}

req = urllib.request.Request(url, body, headers)

# Make sure decode_base64_to_nifti() and plot_segmentation_masks() are defined above
try:
    response = urllib.request.urlopen(req)
    result = response.read()
    result_list = json.loads(result)

    # Extract and decode NIfTI segmentation output
    nifti_file_str = result_list[0]["nifti_file"]
    nifti_file_data = decode_base64_to_nifti(nifti_file_str)

    print(nifti_file_data.shape)
    plot_segmentation_masks(nifti_file_data)

except urllib.error.HTTPError as error:
    print("The request failed with status code: " + str(error.code))
    print(error.info())
    print(error.read().decode("utf8", 'ignore'))

The input_data object contains the following fields:

The input_data object contains the following fields:

Requesting segmentation of all cells in a pathology image

{
  "input_data": {
    "columns": [
      "image",
      "text"
    ],
    "index":[0],
    "data": [
      ["iVBORw0KGgoAAAANSUhEUgAAAAIAAAACCAYAAABytg0kAAAAAXNSR0IArs4c6QAAAARnQU1BAACx\njwv8YQUAAAAJcEhZcwAAFiUAABYlAUlSJPAAAAAbSURBVBhXY/gUoPS/fhfDfwaGJe///9/J8B8A\nVGwJ5VDvPeYAAAAASUVORK5CYII=\n",
      "neoplastic & inflammatory cells "]
    ]
  }
}

Requesting segmentation of pancreas

{
  "input_data": {
    "columns": [
      "image",
      "text"
    ],
    "index":[0],
    "data": [
      ["iVBORw0KGgoAAAAN...",
      "pancreas"]
    ]
  }
}

The response payload is a list of JSON-formatted strings, each corresponding to a submitted image. Each string contains a segmentation_object.

The segmentation_object contains the following fields:

The segmentation_mask contains the following fields:

The response is a list of objects. Each object contains the segmentation result for one input. The segmentation mask is encoded as a Base64 string inside a serialized JSON object under the key nifti_file.

Response to a simple inference requesting segmentation of two objects

[
  {
    "image_features": "{ 
    'data': '4oCwUE5HDQoa...',
    'shape': [2, 1024, 1024], 
    'dtype': 'uint8'}",
    "text_features": ['liver', 'pancreas']
  }
]

[
  {
    "nifti_file": "{\"data\": \"H4sIAAAAAAAE...\"}"
  }
]

The nifti_file field is a stringified JSON object. To decode:

import json
import base64
import tempfile
import nibabel as nib
import os

def decode_base64_to_nifti(base64_string: str):
    """
    Decode a Base64 string back to a NIfTI image.

    Args:
        base64_string (str): Base64 encoded string of NIfTI image, wrapped in a JSON string

    Returns:
        np.ndarray: Decoded NIfTI image data as a NumPy array
    """
    # Parse the inner JSON object to extract the 'data' field
    base64_string = json.loads(base64_string)["data"]

    # Decode the Base64 string to raw bytes
    byte_data = base64.b64decode(base64_string)

    # Write the decoded bytes to a temporary .nii.gz file
    with tempfile.NamedTemporaryFile(suffix='.nii.gz', delete=False) as temp_file:
        temp_file.write(byte_data)     # Write bytes to file
        temp_file.flush()              # Ensure it's fully written to disk
        nifti_image = nib.load(temp_file.name)  # Load the image with nibabel

    # Clean up the temporary file to avoid clutter
    os.unlink(temp_file.name)

    # Return the image as a NumPy array
    return nifti_image.get_fdata()

To plot segmentation mask

import matplotlib.pyplot as plt

def plot_segmentation_masks(segmentation_masks):
    """
    Plot a series of 2D slices from a 3D segmentation mask volume.
    Only slices with non-zero masks are displayed.

    Args:
        segmentation_masks (np.ndarray): A 3D NumPy array of shape (H, W, D),
                                         where D is the number of slices.
    """
    index = 1
    plt.figure(figsize=(15, 15))

    # Loop through each slice (along the depth axis)
    for i in range(segmentation_masks.shape[2]):
        # Only show slices that contain non-zero segmentation
        if segmentation_masks[:, :, i].sum() > 0:
            plt.subplot(4, 4, index)  # Adjust grid size if needed
            plt.imshow(segmentation_masks[:, :, i], cmap='gray')
            plt.axis('off')
            index += 1

    plt.tight_layout()
    plt.show()

The deployed model API supports images encoded in PNG format. For optimal results, we recommend using uncompressed or lossless PNGs with RGB images.

As described in the API specification, the model only accepts images in the resolution of 1024x1024 pixels. You need to resize and pad images if they have a non-square aspect ratio.

For techniques and sample code useful for submitting images of various sizes stored using various biomedical imaging formats, see the Generating Segmentation for a Variety of Imaging Modalities notebook.

The deployed model API supports volumes encoded in NIfTI format.