OCT B-scan video registration (reg_upernet_convnext_t_bmlines)
A learned, fully unsupervised registration model for OCT B-scan videos, where one B-scan location is imaged repeatedly over time. It predicts the motion of each frame against a reference frame:
dx: one lateral shift per frame;dy(x): one axial shift per A-scan column.
It is the second stage of the OcularRigidity pipeline (choroid segmentation, then registration, then pulsation analysis). The registered videos are used to measure the pulsatile choroid thickness change.
The model has no image encoder of its own. It works on the features of the
frozen encoder of the choroid segmentation model
ClementP/ChoroidSegmentationModule@version-3.0.0
(UPerNet ConvNeXt-T). The two models must be used together, at these revisions.
| Parameters | 4.6 M (plus the frozen 37 M segmentation model) |
| Input | Encoder feature pyramids (channels [96, 192, 384, 768]) of a fixed and a moving frame |
| Frame format | uint8, 1536 rows × 1024 columns, PLEX Elite pixel size (1.95 µm axial × 5.9 µm lateral) |
| Output | dx (B,) and dy (B, W), in pixels of the input frame |
| Checkpoint | epoch 10 of 10, step 167,360 (best validation objective) |
Motion convention
The warped frame is
out(x, y) = moving(x - dx, y + dy(x))
so dy(x) is the axial displacement of the moving frame's tissue relative to
the reference (positive means the tissue sits lower in the moving frame), and
dx is the lateral shift.
Usage
Through the OcularRigidity package. The fused pass runs the shared encoder once per frame for both segmentation and registration:
from ocularrigidity.registration.config import RegistrationConfig
from ocularrigidity.registration.fused import segment_and_register
from ocularrigidity.segmentation.trainer.pl_module import ChoroidSegmentationModule
from ocularrigidity.registration.deep_learning.models.regressor import RegistrationRegressor
seg = ChoroidSegmentationModule.from_pretrained("ClementP/ChoroidSegmentationModule", revision="version-3.0.0").eval().cuda()
reg = RegistrationRegressor.from_pretrained("ClementP/OCTVideoRegistration", revision="reg_upernet_convnext_t_bmlines").eval().cuda()
# frames: (T, H, W) uint8 video at the PLEX pixel size (Spectralis resampled to the
# PLEX grid and padded to 1536x1024)
res = segment_and_register(frames, seg, reg, RegistrationConfig(method="learned"))
res.registered_frames # (T, H, W) uint8
res.transform["dx"] # (T,)
res.transform["dy"] # (T, W)
res.registered_masks # (T, H, W) bool, registered choroid masks
Calling the model directly: reg(fixed_feats, moving_feats) takes the
encoder outputs seg.model.encoder(x)[2:], with x = (frame / 255 - 0.5) / 0.5
of shape (B, 1, H, W).
Reference frame: by default it is picked automatically (the frame closest to
the median choroid area, then the motion medoid). Pass ref_idx= to fix it.
Training data
All the data is unlabelled. Each video contributes 40 frames drawn at random without replacement, with a reproducible seed per video. For PLEX the first 20 and last 10 frames are excluded, the same trim as the pipeline; Spectralis series are used whole:
| Device | Scan | Videos | Patients |
|---|---|---|---|
| Zeiss PLEX Elite 9000 | Macula | 400 | 338 |
| Zeiss PLEX Elite 9000 | Optic nerve head | 400 | 362 |
| Heidelberg Spectralis | Macula | 400 | 400 |
- Total: 1,200 videos and 48,000 frames.
- Spectralis frames: resampled to the PLEX grid.
- Validation: 5% of the videos, split by patient, with 40 pairs per video.
Training procedure
There are no ground-truth displacements and no classical registration targets.
Training samples
- Pairs: within-video pairs, 120 per video per epoch.
- Triplets: 40 per video, for cycle consistency.
- Synthetic shifts: 40 per video, up to ±8 px, for shift equivariance.
- Crops: random 768 × 512 crops.
Objective (weights in brackets):
| Term | Weight | What it does |
|---|---|---|
| Feature | 1.0 | Similarity of frozen encoder features after warping |
| Photometric | 1.0 | Multi-scale local normalised cross-correlation |
| Lateral | 1.0 | Lateral alignment term |
| Bruch's membrane (BM) | 0.1 | Distance between the fixed BM line and the analytically warped moving BM line |
| Smoothness | 0.05 | Penalises jumps in dy(x) between columns |
| Curvature | 0.02 | Penalises curvature in dy(x) |
| Coverage | 0.5 | Keeps the warped choroid inside the frame |
| Cycle | 0.3 | A to B to C must equal A to C on triplets |
| Shift equivariance | 1.0 | A known synthetic shift must be recovered |
The BM lines come from the segmentation model, run once on every full, unaugmented frame and stored with the dataset. They are cropped with the same window as the frames.
Architecture: a cascade of one stage per encoder level, coarse to fine. Each
stage builds a local correlation volume (±6 px axially,
±10 steps of 8.0 px laterally) and refines the
global dx and the per-column dy. The dy is a bulk shift plus a
zero-mean residual. Training also uses a coarse-to-fine progress schedule.
Optimisation
- Optimiser: AdamW, learning rate 0.0005, weight decay 0.0001.
- Batches: batch size 4 per GPU, 10 epochs.
- Hardware: 2 GPUs (PyTorch Lightning DDP), bf16 encoder.
Evaluation
Unsupervised validation objective on held-out patients, at the released checkpoint:
| Metric | Value |
|---|---|
| Total objective | 2.327 |
| BM line error after warping | 3.59 px (about 7 µm) |
| Cycle inconsistency | 0.13 px |
| Synthetic-shift error | 0.07 px |
These are self-consistency measures, not errors against ground-truth motion, which doesn't exist for these videos.
Limitations
- Motion model: only a global lateral shift and a per-column axial shift. No rotation, no lateral warp within the frame, no out-of-plane motion: frames where the eye moved out of the B-scan plane can't be corrected, only flagged.
- Segmentation dependency: the model is tied to the segmentation encoder it was trained with. A different segmentation revision needs retraining.
- Training domain: PLEX Elite (macula and optic nerve head) and Spectralis (macula) only.
- Research use only. This is not a medical device.
Citation
OcularRigidity (reference to come).