Reconstruction of lung images from ultrasound RF signal data. I designed and trained multiple architectures, ultimately increasing inference depth and improving on initial designs. I worked on this research project at the Anima AI + Science Lab. I was mentored by Dr. Jiayun Wang.
The objective is to learn the mapping from radio frequency (RF) data to ultrasound images. The RF data has shape (time, num_transducers, num_receivers) (1509, 64, 128). At any given moment in time, this means 64 transducers emit a signal, and 128 transducers receive and record the signal. Each time sample RF[t, :, :] represents a "slice" of chest imaged. As time increases, we image deeper, past the chest wall and into the lung.
This video visualizes the RF data over time (1509 samples).
rf_sample.mp4
Input data has simplified shape ~ (time, x, y).
2D U-Net architectures are trained on time slices with shape (x, y) so while they can learn spatial features, they fail to learn temporal features.
I implemented and trained a modified 3D U-Net to learn both spatial and temporal features, trained on the full temporal volume with shape (time, x, y) instead of just each temporal slice with shape (x, y). As seen in results, this insight proved helpful.
nets.unetcontains a full implementation of both the 2D and 3D U-Nets- 4x 3D U-Net architectures (flagship)
- 3x 2D U-Net architectures (2D slice training)
nets.cnn.VGG16an initial baseline with an antiquated approach (mostly a sanity test)nets.densevarious dense models (did not perform well)- next steps: explore ViT & diffusion architectures, SSL for large unlabeled data
From top to bottom (full resolution RF data):
- ground truth: 80 mm lung imaging depth
- 32 mm inference: existing models
- 48-80 mm inference: my trained 3D U-Net
Interestingly, my deepest model (bottom) was able to reconstruct some fine details where more shallow existing models fail. As we increase the depth with which models are trained at, we are able to reconstruct more fine details, retaining accuracy.
Metrics and detailed results in paper.