AI Assisted Medical Image Segmentation of At-Risk Organs in Cancer Patients
In cancer research, the segmentation of CT and MRI Scans is an important part of effectively providing treatment. Current technologies require manual annotation of these scans which is time-consuming and laborious. A team of researchers from Canadian universities tackled this challenge and worked with the UBC CIC to develop a machine-learning model to automate the medical image annotation process for head and neck organs on CT and MRI Scans.
Researchers from the University of Calgary, University of Saskatchewan, and McGill University collaborated with the UBC Cloud Innovation Centre (UBC CIC) on a solution to address the challenge of manual medical image annotation, a task often required for diagnosis, prognosis, and treatment planning.
The team participated in the post-challenge phase of the Head and Neck (H&N) Segmentation Grand Challenge and achieved state-of-the-art results. Their performance serves as a benchmark for automatic image segmentation in H&N MR and CT images internationally, providing a baseline for future comparisons and research extensions, such as self-supervised learning. This project also led to a journal paper which is currently under scholarly review and can be found here.
This article explores the team’s motivation and technical solutions, providing insight into their experiences working with the UBC CIC. It includes an interview with the researchers, with quotes referenced throughout the discussion and the full interview posted at the end.
Motivation
Head and Neck (H&N) cancer ranks as the seventh most prevalent cancer globally. Radiotherapy, one of the most common treatments for H&N cancers, works by delivering a high radiation dose to tumorous tissues while sparing the surrounding healthy tissues.
Segmentation is the process of contouring and distinguishing specific regions or structures within medical images and is crucial for identifying Organs At Risk (OAR) in radiotherapy treatment planning. CT and MRI scans provide 3D images of the head and neck. Physicians must review these medical images digitally and label any abnormalities on the images themselves. In each 3D scan, there can be hundreds of individual slices, making this a laborious process. Deep Learning (DL) has been widely explored for automating OAR segmentation. However, most studies have focused on a single modality, either CT or MRI, and not both simultaneously. This project aimed to create a novel DL model, trained on MRI and CT data concurrently, that can be utilized for the automatic segmentation of organs at risk in the Head and Neck.
Computational Challenges and AWS Services
Building a solution to this problem requires multiple iterations of model training and performance evaluation. This requires a platform to process and create models on the fly, and record results. Training models on images requires a lot of computational power and leveraging multiple GPUs in order to train models in reasonable time.
This solution primarily used Amazon S3 to store MRI and CT data, and Amazon SageMaker to train and evaluate the models used in this project. SageMaker distributed data parallelism (SMDDP) was used to improve training times. Within Sagemaker, deep learning models were built using pytorch. For the grand-challenge submission, the challenge rules required a docker image to upload the model.
“… they taught us how to use these resources and make the most out of them. We also met regularly with students and staff from the CIC to discuss our project direction. We collaborated and shared ideas, and there were insights about what kind of resources were available. The teamwork helped us speed up our research work, which was really nice…….. It’s a steep learning curve if you don’t know how the services interact, so it definitely helps having people on the inside explaining to you how it works and speed up the real experiment and get it running”
Technical Solution
Architecture Diagram
Data Pre-Processing and Model Architecture
See the items below to learn more about the AWS architecture
Data
The data for this task is part of the online Head and Neck Challenge 2023, which can be found here. The training dataset consists of 42 CT and MRI volumes, each with 30 organs at risk contoured. For testing, an independent set of data from 14 new patients was used.
Model
The team’s approach enhances the pipeline of the nnU-net architecture, widely used in 2D and 3D image segmentation, by incorporating advanced pre-processing steps such as rigid and non-rigid registration, modality dropout, tailored cropping strategies, and tissue-specific data processing. The proposed pipeline extends the application of nnU-Net to multimodal image segmentation, accommodating scenarios with missing data modalities, which are common in medical imaging. Full details of the model and implementation can be found in their paper here.
Training and Evaluation
The model is evaluated based on its segmentation performance, which involves differentiating various parts of the image by contouring each organ. Two metrics—mean Dice Score (DS) and Hausdorff Distance (HD)—were calculated for each organ at risk (OAR) across the patients to assess the similarity of the predicted contour masks to the ground truth. The ground truth refers to the manually annotated labels created by experts, which serve as the benchmark for the model’s accuracy. Training the model required compute-intensive ml.p4d.24xlarge instances to process thousands of images and build the machine learning model.
Results
The proposed pipeline can effectively segment a wide range of organs at risk in head and neck cancer patients from CT and MRI scans. The team and their model achieved the highest mean Dice Score and lowest mean Hausdorff Distance among all participants in the post-challenge phase of the Head and Neck Segmentation Challenge, setting a new benchmark for this type of image segmentation worldwide.
“We collaborated and shared ideas, and there were insights about what kind of resources were available. The teamwork helped us speed up our research work, which was really nice.”
Future Work
Future work involves utilizing publicly available datasets to pre-train a model using self-supervised learning (SSL). SSL is a machine learning paradigm where the model generates supervisory signals from the data itself, eliminating the need for external human-provided labels. This approach is particularly important as it leverages vast amounts of unlabeled data, making the training process more scalable and cost-effective. After pre-training, the model will be fine-tuned with labeled data and compared to the state-of-the-art models previously developed in this project. The goal is to significantly enhance model performance through the application of SSL.
The collaborative efforts of the research team, alongside the UBC Cloud Innovation Centre, have resulted in significant advancements in automating medical image annotation for organs at risk in head and neck cancer patients. By developing a cutting-edge machine-learning model that effectively segments organs at risk from CT and MRI scans, the team has set a new benchmark in the field. This innovation streamlines the labor-intensive process of manual annotation and enhances the accuracy and efficiency of radiotherapy treatment planning. With plans for future work focusing on self-supervised learning, this research opens new avenues for further improving model performance and expanding its applications in the medical imaging landscape. AI applications in medicine can profoundly improve patient care, revolutionize diagnostic processes, and support healthcare professionals in delivering the best treatments possible.
“… because we got some pretty good results, we decided to compete in the HaN-Seg (The Head and Neck Organ-at-Risk CT & MR Segmentation Challenge) challenge. The initial expected outcome was to create pre-trained models on medical imaging, but along the way we built an application that automatically segments organs at risk in radiotherapy.”
The Interview
Recognizing an opportunity to leverage artificial intelligence in supporting clinical decision-making, a team of researchers and students from the University of Calgary, University of Saskatchewan, and McGill University collaborated with the UBC Cloud Innovation Centre (UBC CIC) on a solution to address the challenge of manual and time-consuming medical image annotation, a task often required for diagnosis, prognosis, and treatment planning. The team worked on developing a machine learning model to automate the medical image annotation process for head and neck organs on CT and MRI scans.
This interview sheds some insight on the hard work and contributions of the student team, and shares what their experience was like working with the UBC CIC.
Note: responses were edited for brevity and clarity
The Project
Q: Thank you for joining us today. Before we start, could you introduce yourselves?
Andrew Heschl: I’m an undergrad currently studying computer science at the University of Calgary. Throughout my second year and during the summer after my first year, I worked at the Vision Research Lab under the supervision of Dr. Farhad Maleki. I’ve been working primarily on the segmentation of medical imaging and have also done some work in the field of precision agriculture with image segmentation tasks.
Mauricio Murillo: I’m finishing my second year of undergrad in computer science at the University of Calgary as well. I’m an international student from Bolivia, and I’ve been working with the Vision Research Lab starting from November 2023. My role in this project has been working on the segmentation of the organs at risk in the head and neck region.
Sébastien Quetin: I’m a PhD student currently at McGill University, under the co-supervision of Dr. Shirin Abbasi Nejad Enger and Dr. Maleki. At the Vision Research Lab, I’m collaborating with Mauricio and Andrew. My main PhD objective is to automate the full workflow of brachytherapy, which is a form of radiation treatment for cancer.
Keyhan Najafian: I’m from the University of Saskatchewan and in the third year of my PhD. I’ve been working mostly on aerial image analysis for precision agriculture images, in which the main focus is algorithm development. I am generally working on projects that use self-supervised learning approaches for developing robust and generalizable models, but I’ve been involved in other medical imaging projects with other universities.
Q: Tell us about the project and the outcomes you hoped to achieve when you initially began working on this project.
Andrew: I was working with Sébastien and Mauricio on the head and neck segmentation of given CTs and MRIs. The goal of this project was to segment 30 different organs at risk to specifically aid clinicians and doctors with their work. We trained a neural network which will use CTs and MRIs from the same patient to do these segmentations.
Mauricio: This was the first project I worked on at the Vision Research Lab, so it was an especially important project to me. I hoped to learn about how research was done with the team, and I’ve learned so much from Dr. Maleki, Sébastien, and everyone involved. It has been a really great experience.
Sébastien: To provide a broader context for this project, the overall idea is to pretrain models on medical imaging, so that it’s easier to further train segmentation or classification models from small labeled datasets. We do this because annotated data, especially in this medical domain, is very scarce. It takes time for doctors to manually annotate data, the job is labour-intensive, and patient privacy concerns often do not allow sharing of data publicly. All of these are reasons why we need to be able to train from small labelled datasets. One way that we propose to tackle this issue is to pre-train models on medical imaging. We started by segmenting images and using the public dataset that Andrew and Mauricio were working on. And because we got some pretty good results, we decided to compete in the HaN-Seg (The Head and Neck Organ-at-Risk CT & MR Segmentation Challenge) challenge. The initial expected outcome was to create pre-trained models on medical imaging, but along the way, we built an application that automatically segments organs at risk in radiotherapy.
Keyhan: I’ll add to what Sébastien said. Collecting data, especially for medical imaging, is challenging, not just in the gathering but in labelling as well. Labelling is a significant bottleneck because it requires domain experts, which are not easy to find. We needed to develop a model that can be easily adapted to different data domains using small sets of data. We usually capture medical images using different devices – you don’t see too many differences with the naked eye, but still, they’re slightly different, so generalizable models that can perform on different data distributions are important. Additionally, because of the scarcity of data, we need to develop models using other approaches, like self-supervised learning which operates effectively with even a limited number of samples. This is why I’ve recently been involved in projects with developing models using self-supervised learning and data simulations instead of data collections or data labelling, making everything automatic and adaptable to different areas of study.
Q: How did you first get involved with the UBC CIC and what did that collaboration look like?
Sébastien: We had the chance to have Dr. Farhad Maleki put us in contact with the UBC CIC. Without him, I don’t think we would have reached out or had a chance to collaborate. For us, it was a lot of help to speed up our research work, which was mainly sharing infrastructure and using software readily available at the UBC CIC. At first, there were a lot of presentations on available AWS resources, and they taught us how to use these resources and make the most out of them. We also met regularly with students and staff from the CIC to discuss our project direction. We collaborated and shared ideas, and there were insights about what kind of resources were available. The teamwork helped us speed up our research work, which was really nice.
Keyhan: I wasn’t on the project from the beginning, but I got involved a few months ago and have been using the facilities. To be honest, considering the large models we need to develop, it is not easy to work on local systems these days, so having these resources helps a lot. The reason is that the speed of the experiment is important. We need to test and evaluate many components, and this isn’t feasible without fast computation.
Sébastien: AWS resources were also useful for us to do things that we normally couldn’t do with regular university resources. I had experience in the past using EC2 machines, but I now had access to Sagemaker and the CIC team showed us applications like MONAI (Medical Open Network for Artificial Intelligence) Label. It’s a steep learning curve if you don’t know how the services interact, so it definitely helps to have people on the inside explaining to you how it works and speed up the real experiment and get it running.
The Experience
Q: Can you talk about a moment or milestone in this project that was particularly rewarding for you as a student?
Andrew: It was rewarding for me when we finally got a submission for the competition. It was a lot of work training this network and building a working Docker image which we were able to submit, and it was an accumulation of all our efforts, time, and research which we put in. The second most rewarding part of this particular research for me was that we were able to get a higher DICE score than the other teams in the competition, which told us that our work was going in the right direction.
Mauricio: There were a ton of moments in this project which were very difficult but also very rewarding. I remember on the final day of the challenge, we were training the model like crazy. I was in the university using the computers here, while Andrew and Sébastien were both using some AWS resources. Everyone was training the model until the very last minute – I think Sébastien stayed up until 3 am to get in one more epoch. We were all trying to get the best submission we could. When it paid off and we saw our results, it was very rewarding.
Sébastien: One more thing to add – these competitions, the way they work is that we have to containerize what we’ve built into an app with Docker. Now that they have the Docker image, anyone can deploy it and use it in the clinic and so on. We created something end-to-end that’s actually working.
Q: What challenges did you face and how did you overcome it, especially with the support of the CIC?
Sébastien: The main thing is getting the experiment running. It’s easy to use resources based on, let’s say a tutorial online, but it was a whole ecosystem we had access to. For example, we have regions, and machines, but before that, you have to ask for quotas and you have to know how these work. We definitely couldn’t have done it without the CIC team. Sometimes we try to run something and we have an error, but we don’t understand it because it’s specific to the SageMaker or any other environment, so it’s nice to have received training from the CIC team from the get-go. We exchanged a lot of messages at first to get everything set up in the beginning, but we were soon able to transform that into a more equal collaboration and brainstorm ideas to move the project forward.
Keyhan: The main challenge that I had was working with SageMaker because of the regions, since I was working in the Canadian region. Some limitations came with the servers that I could access, but gaining quick access to EC2 instances that I could directly connect to through the terminal saved a lot of time and effort.
The Future
Q: Moving forward, how has this experience of working on this project and with the CIC influenced your view of the future of healthcare and the role that AI will play in it?
Andrew: I didn’t interact directly with the CIC’s resources, but the CIC is what allowed me to work on my projects, so I’m grateful. This project showed me that with AI, particularly with these resources becoming available to researchers, we’re going to be able to create more and more applications that can be deployed in the medical setting. Groups like the CIC providing resources to researchers can influence healthcare in a direction where AI will be incorporated as an efficient diagnostic tool and improve patient treatment.
Mauricio: It’s a very exciting time for AI and healthcare, and I’m just grateful to be in this position where we have AWS available to keep learning and keep doing research, and doing the best we can do. This wouldn’t be possible without our universities and the CIC.
Keyhan: We see that these days, AI is used everywhere, and for sure healthcare is one of them. Healthcare is one of the most important ones. For example, we can have a non-medical model, which makes mistakes sometimes; it’s fine, but in the area of healthcare, the margin of error needs to be negligible. Doctors often rely on these outputs, so we need to ensure the models we develop are highly accurate and working without issues. With all the data limitations and other challenges involved, specifically for universities that don’t have access to computational resources, collaborations with organizations like the CIC help a lot by giving us not only the opportunity to grow our knowledge but also both the educational and computational resources to develop such models more efficiently. This will benefit doctors and other professionals in this area of study.
Sébastien: I completely agree with what Keyhan is saying. Because all these infrastructures are present, and it’s easier to build scalable applications. It’s really nice to do research and have the resources to do the best research possible, to create services that are publicly available.
Q: As we wrap up, do you have any advice for other students interested in pursuing similar community-focused and collaborative projects?
Andrew: As an undergrad, some advice I’d give to students like myself is to reach out to your professors and clubs relevant to your interests. For example, we founded an AI club recently here at the university, with the goal of fostering interest in AI and creating a space for people who are interested to get together. Get on a research team, too. Reaching out to Dr. Maleki was the best thing I could have done because it has led to me having the opportunities to work with Sébastien, learn a lot, and collaborate with Keyhan on other projects. So for people in their undergrad, I’d recommend them to reach out and ask for positions, because you need to advocate for yourself if you want to do research in your undergrad years.
Mauricio: Just look for opportunities and don’t be scared. This was one of my first projects. I was very scared in the beginning, but going in you learn a lot and pick up a lot of things really fast.
If you would like access to the code for this project, please submit a request here.
Acknowledgements
Sébastien Quetin, Andrew Heschl, Mauricio Murillo, and Keyhan Najafian supervised by Dr. Farhad Maleki were the team that developed this machine-learning model and set a new milestone towards automating the medical image annotation process for head and neck organs on CT and MRI scans. Rohit Murali, a PhD student at UBC and part of the CIC team, assisted students with following best practices when using AWS infrastructure for cost-effective model training and deployment.
Photo Source: segmentation prediction of a medical image in the grand challenge 2023 H&N dataset visualized using 3D Slicer.
About the University of British Columbia Cloud Innovation Centre (UBC CIC)
The UBC CIC is a public-private collaboration between UBC and Amazon Web Services (AWS). A CIC identifies digital transformation challenges, the problems or opportunities that matter to the community, and provides subject matter expertise and CIC leadership.
Using Amazon’s innovation methodology, dedicated UBC and AWS CIC staff work with students, staff and faculty, as well as community, government or not-for-profit organizations to define challenges, to engage with subject matter experts, to identify a solution, and to build a Proof of Concept (PoC). Through co-op and work-integrated learning, students also have an opportunity to learn new skills which they will later be able to apply in the workforce.
