Front page: model card of reg_upernet_convnext_t_bmlines

65281a5 verified 12 days ago

6.68 kB

---
library_name: pytorch
tags:
- medical-imaging
- ophthalmology
- oct
- choroid
- image-registration
- motion-correction
- unsupervised
- video
- pytorch_model_hub_mixin
- model_hub_mixin
---
# 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`](https://huggingface.co/ClementP/ChoroidSegmentationModule/tree/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:

```python
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).

	---
	library_name: pytorch
	tags:
	- medical-imaging
	- ophthalmology
	- oct
	- choroid
	- image-registration
	- motion-correction
	- unsupervised
	- video
	- pytorch_model_hub_mixin
	- model_hub_mixin
	---

	# 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`](https://huggingface.co/ClementP/ChoroidSegmentationModule/tree/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:

	```python
	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).