Front page: model card of reg_upernet_convnext_t_bmlines

README.md

@@ -0,0 +1,166 @@
+---
+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).