Explainable Segmentation in Medical Imaging: Building Confidence Through Human-AI Collaboration

Project Description Deep-learning based computer vision systems often perform excellently on benchmarks, but have unpredictable failure modes when applied to real world data. This can lead to a lack of trust from professionals who use these systems in their work. In situations when these systems are good enough to be trusted by practitioners, they will often trust them too much, which can lead to serious and embarrassing errors, when the model inevitably makes a mistake. We propose developing an uncertainty-aware ensemble system for medical image segmentation that directs physician attention to ambiguous regions, reducing review time while maintaining diagnostic accuracy.

This could be incredibly perilous when working in domains that can have serious consequences for an incorrect prediction, such as in medicine. When attempting to diagnose someone with cancer, a false non-match could be life or death. In my experience building computer vision tools, I’ve found that they’re often best used to augment the eye of an expert. We will implement an ensemble of 5-7 state-of-the-art segmentation architectures (some potential examples being various flavors of U-Nets and Segment Anything based medical models) and combine their predictions using uncertainty-weighted voting. Regions where models disagree will be flagged with high uncertainty, directing physician attention to the most ambiguous cases.

I would like to pilot this project by building a proof of concept application that will look at one interesting segmentation problem in medical imaging, segment the image using several of the best performing approaches, and then put together a visual report, identifying the regions of the image that the models suggest might be of interest, visualizing based on how many models identify that region. The interface showing this visualization would then prompt the practitioner for feedback, which we could use to further refine our models. If things go really well, we can try to incorporate some Explainable AI techniques to give the predictions further justifications.

For the pilot, I’m open to a number of potential medical imaging problems depending on student interest, and I could potentially supervise multiple students working together on different problems for this projects. I’ve identified a few areas that could make for interesting and relatively easy starting points for this work. All of the datasets are publicly available, and most of the datasets have associated Kaggle competitions, meaning that there is quite a bit of reference code for the students to review.

Some examples of potential starting points:
Kvasir-SEG Dataset for Colon Polyp Segmentation (https://www.kaggle.com/datasets/debeshjha1/kvasirseg)
LIDC/IDRI database of Lung CT Scans for Lung Cancer Segmentation (https://www.kaggle.com/datasets/avc0706/luna16)
BRaTS 2021 Task 1 Dataset for Brain Tumor Segmentation (https://www.kaggle.com/datasets/dschettler8845/brats-2021-task1)
CBIS-DDSM Breast Cancer Tumor Segmentation
(https://www.kaggle.com/datasets/awsaf49/cbis-ddsm-breast-cancer-image-dataset)

This is an excellent project for an undergraduate, with a very low barrier to entry. As long as they have experience with python, this will be a good entry point for working on a machine learning project, as they will get experience reimplementing a number of models from the literature on a single problem. This makes it very achievable for them to complete over the course of a semester, while also allowing them to get exposure to a wide variety of approaches. Furthermore, the plan will be to build this into an application that we could demonstrate to the appropriate medical professionals, and get feedback for further revisions. This means that the student researcher will build experience with “full-stack” AI development, implementing, training, and deploying a model in a real world environment.