ahmed275/medical-annotation-removal

Medical Annotation Removal Pipeline

Fine-tuned SAM3TrackerModel checkpoint and supporting artefacts from the M.Sc. thesis "Bridging the Annotation Distribution Gap in Medical Imaging: A Three-Stage Pipeline for Automated Detection, Segmentation, and Removal of Visual Annotations from Medical Educational Imagery."

Technical University of Munich, M.Sc. Data Engineering and Analytics, 2026.

What this model does

Given an annotated educational image (an arrow, glyph, or freeform contour drawn over an underlying object) and a type-specific geometric prompt derived from the annotation mask, the model predicts a per-instance binary segmentation of the underlying object that the annotation refers to.

The predicted object mask drives Stage 3 of the full pipeline (FLUX.1 Fill inpainting), which removes the annotation while preserving the underlying scene.

Files in this repository

File	Description
best_by_iou.pth	Production checkpoint — epoch 23, val IoU 0.7143. The model deployed for downstream inference.
best_by_loss.pth	Alternative checkpoint — epoch 16, val loss 0.3385. Use if you prefer best-by-loss selection.
RUN_INFO.md	Training metadata snapshot.
loss_iou_curve.png	Training curves.
config.json	Lightweight metadata for downstream loaders.

Architecture

Property	Value
Base model	SAM3TrackerModel (Hiera vision encoder, video-pretrained on SA-V)
Total parameters	458 M
Frozen parameters	454 M (vision encoder)
Trainable parameters	4.2 M (0.9 %) — prompt encoder + mask decoder only
Precision	bfloat16

Training

Property	Value
Dataset	In-house, 9,964 source images expanded to 82,875 annotated samples
Annotation classes	arrow (25 %), number/letter (25 %), freeform_bbox (50 %)
Split	80 / 10 / 10 train/val/test, stratified by source image
Optimizer	AdamW
Learning rate	5e-5
Weight decay	1e-4
Batch size	48
Epochs	30
Loss	20 · focal + Dice + IoU-MSE
Hardware	Single NVIDIA A40 (48 GB), bfloat16
Wall-clock	~12 hours

Headline results on the held-out test split (N = 8,288)

Metric	Value
Micro IoU	0.7086
Boundary F1 (2 px tol)	0.7628

Per-class breakdown:

Class	N	IoU	Boundary F1
Arrow	2,080	0.6543	0.6882
Freeform bbox	4,141	0.7664	0.8523
Number / Letter	2,067	0.6475	0.6585

Type-specific prompt routing

The model is trained with a routing scheme that derives a different geometric prompt for each annotation class:

Class	Prompt
Arrow	Single point at the arrow tip (extracted via PCA on the annotation mask)
Number / Letter	Stratified samples of points on the glyph mask
Freeform bbox	Filled-contour dense mask of the freeform outline

Under a paired comparison against a baseline that forces every annotation through a generic uniform box prompt (N = 5,531 identical basenames), type-specific routing yields:

Metric	Production	Generic	Δ
IoU	0.7226	0.6860	+3.66 pp
Boundary F1	0.7614	0.6946	+6.68 pp

The Boundary F1 gap is 1.83x the IoU gap, and the asymmetry ratio is preserved within every class (arrow 1.83x, freeform 1.80x, number-or-letter 2.07x). This is the empirical signature that type-specific routing recovers boundary-aligned segmentation signal that uniform routing loses.

Intended use and limitations

Intended use. Research on automated cleanup of annotated educational medical imagery for downstream training-data construction (medical vision-language models). The pipeline architecture and per-stage error attribution methodology are reusable in adjacent settings that share the same structural pattern (heterogeneous overlay-style annotations on natural images).

Limitations.

• The evaluation is on a non-clinical proxy dataset that is structurally representative of the medical-annotation problem but is not clinical radiology. Absolute accuracy on clinical imagery remains to be established and is the principal direction for future work.
• All metrics are from a single training seed.
• The full pipeline (Pipeline B) shows a class-asymmetric cost when ground-truth annotation masks are replaced with predictions from GroundedSAM 2: the freeform class degrades by a factor of 4.7x on FID due to a stage-coupling failure in which SAM 2's mask refinement returns the object enclosed by the contour rather than the contour itself. This is the binding constraint identified by the thesis and the proposed fix lives upstream of this checkpoint (in the detection stage).

Inference example

import torch
from sam3_tracker_model import SAM3TrackerModel  # your local loader

# Load
model = SAM3TrackerModel.from_pretrained("facebook/sam2-hiera-large")
ckpt = torch.load("best_by_iou.pth", map_location="cpu")
model.load_state_dict(ckpt["model"], strict=False)
model.eval().to("cuda")

# Type-specific prompt routing (sketch)
# See repository: github.com/USERNAME/medical-annotation-removal
#   /inference/inference_zeroshot_pcs.py for the full prompt-derivation
#   logic per annotation type.

Citation

@mastersthesis{your_thesis_2026,
  title  = {Bridging the Annotation Distribution Gap in Medical Imaging},
  author = {Your Name},
  school = {Technical University of Munich},
  year   = {2026},
  type   = {M.Sc. Thesis}
}

Source code

Full pipeline source at: 👉 https://github.com/USERNAME/medical-annotation-removal

License

This checkpoint is released under the MIT License. Note that it derives from publicly available foundation models (SAM 2 / SAM 3 Tracker, Hiera) — please check the upstream licences of those models before any commercial use.

---

Generated by upload_to_huggingface.py on 2026-06-07 18:17 UTC.