Segment brain MRI images with a UNet

This example uses an encoder-decoder network called a UNet to mark regions in brain MRI images. For each pixel, the network outputs the probability that it belongs to an abnormality. The notebook draws the outline of the predicted region on the image.

Objectives

• Load a pretrained UNet from torch.hub and convert its weights to bfloat16.
• Build encoder blocks with convolution, batch normalisation, ReLU and max pooling.
• Build decoder blocks with transposed convolutions and skip connections.
• Turn the network output into a segmentation mask.

Before you begin

1. Install the toolkit. The example is written to run on the CPU emulator as well as on hardware.

2. Check that Python can import the toolkit:

   python -c "import qant_native_computing_toolkit as q; print(q.generic.get_available_npus())"

   On the emulator this prints {0: 'cpu_backend'}.

3. Get the examples repository:

   git clone --recurse-submodules https://github.com/Q-ANT-GmbH/qant_native_computing_toolkit_examples.git
   cd qant_native_computing_toolkit_examples/python/image_segmentation

4. Install the example’s dependencies:

   pip install -r requirements.txt

   This example also needs scikit-image, which the requirements file includes.

5. Make sure the machine can reach github.com. The notebook loads the model with torch.hub and downloads a sample MRI image from the same GitHub repository.

Run the example

1. Start Jupyter and open image_segmentation_on_npu.ipynb:

   jupyter notebook image_segmentation_on_npu.ipynb

2. Run all cells. The last cell shows the MRI image with the predicted region outlined in red, under the title “U-Net for brain MRI (NPU)”.

To run it without opening a browser, execute the notebook from the command line:

jupyter nbconvert --to notebook --execute image_segmentation_on_npu.ipynb --output executed.ipynb

How it works

Load the model

The notebook loads the pretrained UNet with torch.hub, takes its state dict, and passes it to QAntUNet from qant_unet.py:

w = torch.hub.load(
    "mateuszbuda/brain-segmentation-pytorch",
    "unet",
    in_channels=3,
    out_channels=1,
    init_features=32,
    pretrained=True,
).state_dict()

As in the ResNet example, the constructor converts every tensor to a bfloat16 NumPy array.

Encoder blocks

Each encoder level applies double_conv: two rounds of convolution, batch normalisation and ReLU. Each level after the first starts with a 2 by 2 max pooling that halves the image size:

x = q_ai.conv_fprop(x, conv1_w, 1, 1, 1)
x = q_ai.batchnorm2d_fprop(
    x, bn1_running_mean, bn1_running_var, bn1_running_weight, bn1_running_bias, 0
)
x = q_ai.relu_fprop(x)

The last argument of batchnorm2d_fprop is eps, which qant_unet.py sets to 0. PyTorch batch normalisation uses 1e-5 by default.

Decoder blocks

Each decoder level doubles the image size with a transposed convolution, adds its bias, and joins the result with the matching encoder output along the channel axis:

d = q_ai.conv_transpose_fprop(x, weight, 0, 2, 1, 0)
if bias is not None:
    d = q_ai.add_bias_fprop(d, bias)

The positional arguments are padding 0, stride 2, dilation 1 and output padding 0. The joined channels then go through another double_conv.

Output layer

A final 1 by 1 convolution reduces the channels to one, and sigmoid_fprop turns each value into a probability:

x = q_ai.conv_fprop(x, self.w[f"{block_name}.weight"], 0, 1, 1)

The notebook rounds the probabilities to a 0 or 1 mask and draws its outline on the image with utils.outline.

Expected output

We have no output from the real model to report. See the warning at the top of this page.

As a code check only, we ran QAntUNet on the CPU emulator with randomly initialised weights of the correct shapes and a random 3 by 256 by 256 input. This confirms that the toolkit calls in qant_unet.py run and produce the expected shape. It says nothing about segmentation quality.

Measure	Value from the synthetic code check
Output shape	(1, 1, 256, 256)
Forward pass time on the CPU emulator	56.56 seconds

Run the tests

The example includes a test that runs the notebook, then runs the original PyTorch UNet on the same image and compares the two masks:

python -m pytest -v test_image_segementation.py

The file name is spelled segementation in the repository. The test rounds both outputs to masks and asserts that more than 95% of pixels agree. It also prints the mean and median differences between the two outputs.

In our sandbox the test failed at the torch.hub.load cell with HTTP Error 400: Bad Request, before any toolkit code ran.

Troubleshooting

If torch.hub.load fails with an HTTP error, check that your machine can reach github.com. The model and the sample image both come from GitHub.

What’s next

• See a network built from the same operations that ran end to end on the emulator in Image recognition.
• Learn how to use several cards at once in Parallel execution.