https://github.com/vanvalenlab/torch-dynamicnuclearnet
DynamicNuclearNet implemented in PyTorch
https://github.com/vanvalenlab/torch-dynamicnuclearnet
Last synced: 24 days ago
JSON representation
DynamicNuclearNet implemented in PyTorch
- Host: GitHub
- URL: https://github.com/vanvalenlab/torch-dynamicnuclearnet
- Owner: vanvalenlab
- Created: 2026-05-07T21:04:11.000Z (3 months ago)
- Default Branch: main
- Last Pushed: 2026-06-29T14:57:09.000Z (26 days ago)
- Last Synced: 2026-06-29T16:32:17.046Z (26 days ago)
- Language: Python
- Homepage:
- Size: 530 KB
- Stars: 0
- Watchers: 2
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# Torch Mesmer
A PyTorch implementation of DynamicNuclearNet for segmenting live nuclei for cell tracking.
## The Network
DynamicNuclearNet (DNN) is built on a Panoptic network, consisting of an EfficientNetV2BL backbone connected to a [feature pyramid network](https://arxiv.org/abs/1612.03144). The levels of the backbone and feature pyramid network can be selected, but for the pre-trained DNN model, we use backbone levels `C1-C5`, and pyramid levels `P1-P7`. Pyramid levels are then upsampled to match the input resolution (256 x 256 px) and delivered to the semantic heads of the model.
In the pre-trained DNN model, there are three semantic heads:
```markdown
Head 1 (1, 256, 256)
└─ Inner distance transform for the nucleus
Head 2 (1, 256, 256)
└─ Outer distance transform for the nucleus
Head 3 (2, 256, 256)
├─ Foreground pixels for nucleus
└─ Background pixels for nucleus
```
After softmax on the semantic head convolutions, the model concatenates all predictions into an output tensor of shape `(4, 256, 256)` and returns it.
## The Dataset
Each dataset contains one nuclear and one cytosol channel, as well as the labeled ground truth mask for each channel, and metadata that contains the source tissue and experiment of each image.
- **Training**: 4950 square images (512 x 512)
- **Validation**: 1417 square images (512 x 512 px)
- **Test**: 717 square images (512 x 512 px)
## The Training
### Data loaders
The training and validation data are loaded into a PyTorch `Dataset` object, which conducts preprocessing under the hood for each batch. This `Dataset` is then used in the construction of a `Dataloder`. The preprocessing and augmentation pipeline is outlined below:
1. Each item (one image and ground truth mask) is selected from the full dataset.
2. The image is normalized with two steps:
1. Images are thresholded in order to reduce the influence of bright pixels.
2. Imaes are then normalized using Contrast Limited Adaptive Histogram Equalization (CLAHE) with the `equalize_adapthist` function from [scikit-image](https://scikit-image.org/).
3. The labels are then transformed to generate the inner distance transform, outer distance transform, and foreground/background pixels.
4. The normalized images and mask transformations are then augmented using a random combination of rotations, flips, crops and zooms.
5. The images and masks are then returned to the model.
### Optimizer, Learning Rate, and Loss Function
We use the Adam optimizer with a learning rate of 0.0001. Upon a plateau in validation loss that lasted longer than 5 epochs, the model's learning rate is reduced by a factor of 10.
The loss function is a combination of weighted categorical cross entropy (WCCE) and mean squared error (MSE) loss. For continuous predictions (inner distance transforms), MSE loss was used. For categorical predictions (foreground and background), WCCE loss was used with class weights calculated for each batch. Loss from continuous heads was weighted with 0.01 to increase stability during training. The loss calculated from each head was summed and then used in backpropagation.
We used a batch size of 10 images, and an augmented version of each images was seen only once during each epoch. The model was trained for 50 epochs, and we test the model that returned the lowest validation loss.
## The Testing
The model was used to segment 717 test images. These segmentations were then compared to the ground truth using a custom metrics pipeline that analyzes the following:
- **Recall**
- **Precision**
- **Jaccard index (IoU)** - The index of overlap between the ground truth and the prediction
- **Gained detections** - objects segmented but not present in the ground truth
- **Missed detections** - objects present in the ground truth that were missed by the model
- **Splits** - number of "one to many" errors
- **Merges** - number of "many to one" errors
- **Catastrophes** - number of "many to many" errors
Each of these metrics was calculated for every image, allowing us to identify areas of weakness in each trained model.
## How to use DNN
### Instantiate the model
```python
from tifffile import imread
import matplotlib.pyplot as plt
import numpy as np
from torch_dnn.dnn import DNN
model = DNN(
model_path='../.deepcell/dnn/saved_model_best_dict.pth',
device='cuda:0'
)
```
### Read in the example image and view
```python
frame = imread('example/example_image.tiff')
fig, ax = plt.subplots()
ax.imshow(frame)
ax.set_axis_off()
ax.set_title('Nuclear Image')
plt.show()
```
### Predict
```python
frame = frame[np.newaxis, np.newaxis] # model requires a time and channel axis
mask = model.predict(frame)
```
### Plot the results
```python
fig, ax = plt.subplots(1,2)
ax[0].imshow(frame.squeeze())
ax[0].set_title('Nuclear Image')
ax[1].imshow(mask.squeeze())
ax[1].set_title('Predicted segmentation')
for axis in ax:
axis.set_axis_off()
fig.tight_layout()
plt.show()
```
