This project answers two problems in healthcare and AI:

  1. COVID-19 chest CT lesion segmentation.
  2. Verified neural networks with out-of-distribution detection

Deep neural networks are seeing widespread use in the healthcare field. Many recent innovations in computer vision have especially advanced technologies regarding medical imaging. However, performance of neural networks in practice is not guaranteed whenever there are subtle differences between new data and the training data. This is known as dataset shift, or out-of-distribution detection.

It would be ideal if we could know a priori whether a given input belongs to the same abstract distribution as the training data. If so, the result from the neural network could be trusted.

What it does

This project has two main components:

  1. COVID-19 chest CT lesion segmentation.
  2. Out-of-distribution detection with OT-Flow, an optimal transport normalizing flow based on Neural Ordinary Differential Equations (Neural ODE).

First, an image is input into the COVID-19 lesion segmenter to produce a segmentation mask. Second, the extracted features are input into the Neural ODE to produce a "likelihood score", a number describing how likely it is that these features come from the same distribution as the training set.

How I built it

The COVID-19 lesion segmentation network is a manually-coded U-Net trained from scratch on the open-source dataset MosMedData Chest CT Scans with COVID-19.

The NeuralODE is adapted from the OT-Flow code and trains on extracted features of the penultimate layer from the frozen U-Net using the same training dataset.

Challenges I ran into

Initially there was a step 3 plan for this project: deploy a webapp to run the two networks client-side using ONNX. However, since ONNX is still in development, there were two main issues I ran into. First, ONNX doesn't support feature map upsampling layers required for the U-Net architecture. Second, even a workaround to this resulted in undefined behavior and my laptop GPU crashing. Thus, the webapp portion of the project was discarded.

Accomplishments that I'm proud of

I'm proud that I got the OT-Flow model to work. It is difficult to train and I don't have experience with these types of networks. I also became significantly more familiar with PyTorch, since most of my background is in TensorFlow.

What I learned

ONNX is promising but still in development.

NeuralODEs aren't as intimidating as they seem, but are still very tricky to train.

What's next for FlowCheck

This framework is not restricted to just COVID-19 segmentation. All neural networks that to be deployed in mission-critical scenarios such as healthcare either need human supervision or automated guarantees. While this method does not provide rigorous mathematical proof of verified neural networks, it is a step forward towards fully-automated systems.

Judging Criteria at a Glance

Polish Code is written in pure Python and Jupyter Notebooks for dissemination. Easily ported to provide command line tools for hospitals / research centers.

Usefulness Successfully segments lesions in chest CT of patients with COVID-19. Addresses important concerns with deployment of AI tools in practical settings.

Creativity Application of Neural ODEs to calculate log-likelihood of extracted features from a pre-trained network to do feature-level anomaly detection rather than image-level.

Technical Skill Normalizing flows are one of the most complicated variants of deep generative models (compared to variational autoencoders and generative adversarial networks), and continuous normalizing flows such as Neural ODEs are even moreso. These networks are fully invertible and involve a differential equation solver within the structure.

In this project, the normalizing flow is overlaid with a traditional deep segmentation task which was also prepared and trained from scratch within the scope of this hackathon.

Built With

Share this project: