DEV Community

kito2718
kito2718

Posted on

BCT4: 3D U-Net Bottlenecks and Minimal CV Score Improvement (0.4483 -> 0.4578)

#python #machinelearning #deeplearning #computervision

Eye Catch Image
Biohub - Cell Tracking During Development

Abstract

  • Constructed and trained a 3D U-Net model in PyTorch for 3D cell centroid detection.
  • Observed a slight CV score increase from 0.4483 to 0.4578, falling far short of top leaderboard scores(0.982).
  • Analyzed key technical bottlenecks including extreme 3D class imbalance, spatial resolution loss, and tracking limits.

Overview & Background

In our local validation setup, simple intensity thresholding(blob_dog) struggled to separate dark, dim cells from background noise.
To overcome this limitation, we generated 3D Gaussian heatmaps centered on Ground Truth(GT) cell coordinates and trained a 3D Convolutional Neural Network(3D U-Net). However, even after multiple training epochs and hyperparameter tuning, the CV score only marginally improved to 0.4578.

This post analyzes the technical bottlenecks and structural limitations responsible for this plateau.

Main Content

1. Overall 3D U-Net Pipeline

The complete processing pipeline combining 3D U-Net detection and Linear Sum Assignment(LAP) tracking is summarized below.

3D U-Net Processing Flow
Figure 2: Overall pipeline from 3D image input to metric evaluation

2. Quantitative Results & Score Progression

Comparison of predicted node counts and local CV scores across experiments:

Experiment ID Method & Model Detection Setting Predicted Nodes Edge Jaccard Division Jaccard FINAL CV SCORE
Baseline Traditional Processing(blob_dog) Intensity Threshold ~19,800 0.4483 0.0000 0.4483
Exp 1 3D U-Net(3 Ep) Fixed Threshold(0.25) 464,195 0.1223 nan 0.1223
Exp 2 3D U-Net(15 Ep) Fixed Threshold(0.55) 35,000 0.4578 nan 0.4578
Exp 3 Target 3D U-Net Adaptive Threshold + LAP 35,000 0.4589 0.0039 0.4593 
[Local CV Score Progression]
0.50 |
0.45 |-----------------------+----------+ (Exp 2/3: 0.4593)
     |                       |          |
0.40 |  + (Baseline: 0.4483) |          |
0.35 |  |                    |          |
0.30 |  |                    |          |
0.25 |  |                    |          |
0.20 |  |                    |          |
0.15 |  |   + (Exp 1: 0.1223)|          |
0.00 +--+---+----------------+----------+--------->
       Base Exp1            Exp2       Exp3
Enter fullscreen mode Exit fullscreen mode

3. Key Technical Bottlenecks

We identified three fundamental reasons why a standalone 3D U-Net fails to reach high leaderboard performance:

(1) Extreme Class Imbalance in 3D Space(Zero Collapse)

In a 3D image volume(64x256x256 voxels), cell centers occupy less than 0.1% of total voxels. When training with Standard Mean Squared Error(MSE) loss, the network easily collapses into a local minimum where predicting all background(0.0) minimizes total loss.

(2) Loss of Spatial Resolution in Deep Pooling

Max-pooling operations inside 3D U-Net reduce spatial resolution, blurring boundaries between tightly clustered cells and merging adjacent centroids into a single heatmap blob.

(3) Limitations of Frame-to-Frame Distance Tracking

Greedy or LAP matching based solely on frame-to-frame Euclidean distance fails to capture long-term temporal trajectory consistency and cell division events, introducing false-positive and false-negative edges.

Conclusion

A naive 3D U-Net heatmap regression approach plateaus around CV 0.4593 and cannot reach competitive leaderboard performance(0.96+). Rebuilding the architecture with 3D Instance Segmentation(e.g., StarDist 3D) and Multi-frame Min-Cost Flow tracking is necessary for the next major milestone.

Hope this helps!

Japanese Series Links

For Japanese readers, the full series is published on Zenn:

Top comments (0)