Quantizing 3D Segmentation Model¶
This example demonstrates the use of the Post-training Optimization Tool API for the task of quantizing a 3D segmentation model.
The Brain Tumor Segmentation model from PyTorch is used for this purpose. A custom
DataLoader is created to load images in NIfTI format from the Medical Segmentation Decathlon BRATS 2017 dataset for 3D semantic segmentation task and the implementation of the Dice Index metric is used for the model evaluation. In addition, this example demonstrates how one can use image metadata obtained during image reading and preprocessing to post-process the model raw output. The code of the example is available on GitHub.
How to Prepare the Data¶
To run this example, you will need to download the Brain Tumors 2017 part of the Medical Segmentation Decathlon image database http://medicaldecathlon.com/.
3D MRI data in NIfTI format can be found in the
imagesTr folder, and segmentation masks are in
How to Run the Example¶
Launch Model Downloader tool to download
brain-tumor-segmentation-0002model from the Open Model Zoo repository.
omz_downloader --name brain-tumor-segmentation-0002
Launch Model Converter tool to generate Intermediate Representation (IR) files for the model:
omz_converter --name brain-tumor-segmentation-0002
Launch the example script from the example directory:
python3 ./3d_segmentation_example.py -m <PATH_TO_IR_XML> -d <BraTS_2017/imagesTr> --mask-dir <BraTS_2017/labelsTr>
Optional: you can specify .bin file of IR directly using the