A 3-stage pipeline for fast, lightweight microbubble sizing and counting via knowledge distillation.
⚠️ IMPORTANT BUG FIX: The original train_student.py in this repo uses BCEWithLogitsLoss on binary foreground/background masks. This fails catastrophically because microbubble foreground is only ~0.2% of pixels — the model learns to predict ALL background and achieves 99.8% accuracy while detecting zero bubbles. The fixed script train_mse_distill.py uses MSE distillation on the teacher's raw cell_prob LOGITS instead, which gives gradients on ALL pixels (background pixels have informative negative logits ~-6). See train_mse_distill.py for the corrected implementation.
The Problem
Cellpose-SAM is excellent for cell/bubble segmentation, but at ~300M params (1.1 GB) it's expensive at inference. For lab settings where your slides look similar and you're "just detecting circles", this is massive overkill. You're paying for the ability to also segment dogs, neurons, and a thousand other things — capacity you don't need.
The Solution: Distill Into a Tiny Specialist
Model
Params
Size
256×256 GPU
256×256 CPU
FPS (GPU)
Cellpose-SAM
~300M
1.1 GB
~100 ms
seconds
~10
TinyBubbleNet (base_ch=16)
389K
1.5 MB
3 ms
45 ms
337
TinyBubbleNet (base_ch=32)
1.5M
5.8 MB
~5 ms
~80 ms
~200
~33× faster, ~750× smaller. And when your domain is narrow (similar-looking lab slides), the accuracy loss is minimal because the student only needs to learn one visual distribution.
Architecture
TinyBubbleNet is a depthwise-separable U-Net (inspired by PicoSAM2) with a 4-channel output:
Channel
Name
What it encodes
0
dY
Vertical gradient flow (Cellpose-compatible)
1
dX
Horizontal gradient flow (Cellpose-compatible)
2
cell_prob
Foreground/background probability
3
dist_transform
Distance transform (peak = bubble radius)
Instance masks are reconstructed via Euler integration of the flow field — identical to Cellpose post-processing. This means the student is fully compatible with the Cellpose ecosystem.
The distance transform head is the key addition for sizing: the peak value within each detected instance directly gives you the bubble radius.
The Bug and The Fix
The Bug (original train_student.py / losses.py)
python
1# BAD: BCE on binary masks2prob_loss = BCEWithLogitsLoss(pred_prob, binary_mask)
With foreground at only ~0.2% of pixels, the model's dominant gradient signal is "predict all background". Even after 300 epochs with "best val loss 0.0008", the model predicts zero bubbles everywhere.
The Fix (train_mse_distill.py)
python
1# GOOD: MSE on teacher's raw logits2prob_loss = MSE(pred_prob_logits, teacher_cell_prob_logits)
The teacher outputs cell_prob as raw logits (range roughly -9 to +5). Every pixel has an informative value — background pixels should reproduce ~-6, foreground pixels should reproduce ~+5. MSE on logits gives strong gradients everywhere, and the student successfully learns to segment bubbles.
Use --no_depthwise for standard convolutions (more params, possibly better accuracy on complex images).
Key Design Decisions
Why Cellpose flows instead of direct mask prediction? Flows handle overlapping/touching bubbles via convergence — each pixel flows toward its instance center. Direct mask prediction can't separate touching instances.
Why distance transform head? For circles, the DT peak = radius. This gives you sizing "for free" without post-processing the mask.
Why depthwise-separable convs? ~8× fewer params than standard convs. For a narrow domain (your lab slides), this compression is lossless.
Why MSE on logits instead of BCE on masks? See "The Bug and The Fix" section above. BCE on sparse binary masks fails due to extreme class imbalance. MSE on teacher logits gives gradients everywhere.
When to Re-train
The student is specialized to your current lab setup. Re-train when:
Microscope/camera settings change significantly
Bubble preparation protocol changes
Image resolution changes
Re-training is fast: ~30 min for 400 epochs on 50 images with a GPU.
