Improving Medical Imaging Models Through Robust Data Annotation
In medical AI, model performance is often constrained not by architecture but by data quality. This is especially true in 3D medical imaging — CT scans, MRI volumes, and PET/CT studies — where annotation requires specialized clinical expertise and the cost of a labeling error is not just a worse metric but potentially an incorrect clinical decision.
During the development of machine learning systems for medical imaging, we observed that annotation inconsistencies can degrade segmentation Dice scores by 5-15%, depending on the anatomy and modality. In regulated settings (FDA-cleared devices, CE-marked software), these inconsistencies also create audit risk: you need to demonstrate that your training data meets a defensible quality standard.
This article covers practical strategies for detecting and correcting annotation errors, with code examples and concrete metrics.
The Clinical Context: Why Annotation Quality Is Non-Negotiable
Before diving into the technical approach, it’s worth understanding why annotation errors in medical imaging carry outsized consequences compared to other domains.
False positives in diagnostic models can trigger unnecessary follow-up procedures — biopsies, additional imaging with radiation exposure, patient anxiety. In screening applications (lung cancer CT screening, mammography), the base rate of disease is low, so even a small increase in false positive rate translates to a large number of unnecessary interventions across a population.
False negatives are even more dangerous: a missed lesion means a missed diagnosis. In oncology, delayed detection can mean the difference between Stage I and Stage III disease.
Annotation errors corrupt both directions. A radiologist who inconsistently labels small nodules introduces noise that makes the model’s decision boundary unreliable in exactly the size range where clinical decisions are most difficult (5-10mm nodules, where guidelines differ on follow-up intervals).
In our experience, cleaning 5% of a dataset’s annotations often improves model performance more than doubling the dataset size. The leverage of data quality work is enormous, but it’s consistently underinvested in because it’s less glamorous than architecture research.
Annotation Quality Pipeline
The approach we use is a systematic pipeline that combines automated detection with expert review:
┌──────────────────┐ ┌──────────────────┐ ┌───────────────────┐
│ │ │ │ │ │
│ Raw Annotations │────▶│ Loss Tracking │────▶│ Outlier │
│ (Segmentation │ │ (Per-sample │ │ Detection │
│ masks from │ │ Dice/CE loss │ │ (Statistical │
│ radiologists) │ │ across epochs) │ │ thresholds) │
│ │ │ │ │ │
└──────────────────┘ └──────────────────┘ └─────────┬─────────┘
│
▼
┌──────────────────┐ ┌───────────────────┐
│ │ │ │
│ Clean Dataset │◀────│ Manual Review │
│ (Verified │ │ (Expert │
│ annotations) │ │ re-annotation) │
│ │ │ │
└──────────────────┘ └───────────────────┘
The key insight is that the model itself becomes a tool for finding bad annotations. Samples that the model consistently gets wrong — or that it oscillates on across training epochs — are disproportionately likely to have annotation problems.
Detecting Annotation Errors via Training Dynamics
The core detection mechanism is per-sample loss tracking. During training, we record the loss for every sample at every epoch. Samples with annotation errors exhibit characteristic patterns:
- Persistently high loss: The model never learns to match the annotation, because the annotation is wrong
- High variance across epochs: The model oscillates between fitting the (incorrect) label and fitting the true underlying pattern
- Loss that increases as training progresses: The model initially memorizes the wrong label, then “unlearns” it as it sees more correct examples of similar anatomy
Here is a concrete implementation:
import torch
import torch.nn as nn
import numpy as np
from collections import defaultdict
from dataclasses import dataclass
@dataclass
class SampleRecord:
sample_id: str
losses: list[float] # one per epoch
class AnnotationQualityTracker:
"""Track per-sample loss across training epochs to detect annotation errors.
Usage:
tracker = AnnotationQualityTracker()
for epoch in range(num_epochs):
for batch in dataloader:
loss_per_sample = compute_per_sample_loss(model, batch)
for sid, loss_val in zip(batch["sample_ids"], loss_per_sample):
tracker.record(sid, epoch, loss_val)
suspects = tracker.get_suspects(
min_epochs=10,
loss_percentile=95,
variance_percentile=90,
)
"""
def __init__(self):
self._records: dict[str, SampleRecord] = {}
self._current_epoch: dict[str, list[float]] = defaultdict(list)
self._epoch = -1
def record(self, sample_id: str, epoch: int, loss: float):
"""Record a per-sample loss value for a given epoch."""
if epoch != self._epoch:
# Flush previous epoch: average any duplicate observations
# (from augmentation or multi-crop) into a single value.
self._flush_epoch()
self._epoch = epoch
self._current_epoch[sample_id].append(loss)
def _flush_epoch(self):
for sample_id, losses in self._current_epoch.items():
avg_loss = sum(losses) / len(losses)
if sample_id not in self._records:
self._records[sample_id] = SampleRecord(
sample_id=sample_id, losses=[]
)
self._records[sample_id].losses.append(avg_loss)
self._current_epoch.clear()
def get_suspects(
self,
min_epochs: int = 10,
loss_percentile: float = 95,
variance_percentile: float = 90,
) -> list[dict]:
"""Identify samples likely to have annotation errors.
A sample is flagged if it meets ANY of these criteria:
1. Mean loss is above the loss_percentile threshold
2. Loss variance is above the variance_percentile threshold
3. Loss trend is positive (increasing loss over training)
Returns a sorted list of suspect samples with diagnostic metadata.
"""
self._flush_epoch()
# Filter to samples with enough history
eligible = {
sid: rec for sid, rec in self._records.items()
if len(rec.losses) >= min_epochs
}
if not eligible:
return []
# Compute statistics for each sample
stats = []
for sid, rec in eligible.items():
losses = np.array(rec.losses)
mean_loss = float(np.mean(losses))
loss_var = float(np.var(losses))
# Linear regression slope to detect increasing loss
epochs = np.arange(len(losses), dtype=np.float64)
slope = float(np.polyfit(epochs, losses, 1)[0])
stats.append({
"sample_id": sid,
"mean_loss": mean_loss,
"loss_variance": loss_var,
"loss_trend": slope,
"num_epochs": len(losses),
"final_loss": float(losses[-1]),
})
# Compute thresholds
all_means = np.array([s["mean_loss"] for s in stats])
all_vars = np.array([s["loss_variance"] for s in stats])
loss_threshold = float(np.percentile(all_means, loss_percentile))
var_threshold = float(np.percentile(all_vars, variance_percentile))
# Flag suspects
suspects = []
for s in stats:
reasons = []
if s["mean_loss"] >= loss_threshold:
reasons.append("high_mean_loss")
if s["loss_variance"] >= var_threshold:
reasons.append("high_variance")
if s["loss_trend"] > 0:
reasons.append("increasing_loss")
if reasons:
s["flag_reasons"] = reasons
s["priority_score"] = (
s["mean_loss"] / loss_threshold
+ s["loss_variance"] / max(var_threshold, 1e-8)
+ max(s["loss_trend"], 0) * 10
)
suspects.append(s)
suspects.sort(key=lambda x: x["priority_score"], reverse=True)
return suspects
def compute_per_sample_dice_loss(
predictions: torch.Tensor, # (B, C, D, H, W) — logits
targets: torch.Tensor, # (B, C, D, H, W) — one-hot
smooth: float = 1e-5,
) -> torch.Tensor:
"""Compute Dice loss per sample (not reduced across the batch).
Returns a tensor of shape (B,) with loss values for each sample.
"""
probs = torch.sigmoid(predictions)
# Flatten spatial dimensions
probs_flat = probs.view(probs.size(0), probs.size(1), -1)
targets_flat = targets.view(targets.size(0), targets.size(1), -1)
intersection = (probs_flat * targets_flat).sum(dim=-1)
union = probs_flat.sum(dim=-1) + targets_flat.sum(dim=-1)
dice_per_class = (2.0 * intersection + smooth) / (union + smooth)
dice_loss = 1.0 - dice_per_class.mean(dim=1) # average across classes
return dice_loss # (B,)
The priority_score combines multiple signals into a single ranking. In practice, we review the top 3-5% of flagged samples, which typically catches 60-80% of genuine annotation errors with a manageable review workload for the clinical team.
Inter-Annotator Agreement Metrics
Before building automated detection, it’s critical to understand your baseline annotation quality. Inter-annotator agreement quantifies how consistently multiple annotators label the same data.
Cohen’s Kappa for Classification Tasks
For classification labels (e.g., “lesion present” vs. “no lesion”), Cohen’s kappa measures agreement corrected for chance:
def cohens_kappa(annotator_a: list[int], annotator_b: list[int]) -> float:
"""Compute Cohen's kappa between two annotators.
Values:
< 0.20 Poor agreement
0.21-0.40 Fair
0.41-0.60 Moderate
0.61-0.80 Substantial
0.81-1.00 Almost perfect
In clinical settings, substantial agreement (>0.61) is typically
the minimum acceptable standard for training data.
"""
assert len(annotator_a) == len(annotator_b)
n = len(annotator_a)
# Observed agreement
agreed = sum(a == b for a, b in zip(annotator_a, annotator_b))
p_observed = agreed / n
# Expected agreement by chance
classes = set(annotator_a) | set(annotator_b)
p_expected = 0.0
for c in classes:
p_a = sum(1 for x in annotator_a if x == c) / n
p_b = sum(1 for x in annotator_b if x == c) / n
p_expected += p_a * p_b
if p_expected == 1.0:
return 1.0
return (p_observed - p_expected) / (1.0 - p_expected)
Dice Coefficient for Segmentation Tasks
For volumetric segmentation, the Dice coefficient (also known as F1 score for sets) is the standard agreement metric:
def dice_coefficient(
mask_a: np.ndarray,
mask_b: np.ndarray,
) -> float:
"""Compute Dice coefficient between two binary segmentation masks.
Dice = 2 * |A ∩ B| / (|A| + |B|)
For medical segmentation, we generally require:
- Dice > 0.85 for large structures (liver, lungs)
- Dice > 0.70 for medium structures (kidneys, spleen)
- Dice > 0.50 for small/difficult structures (lymph nodes, small lesions)
Annotations below these thresholds should be reviewed.
"""
intersection = np.logical_and(mask_a, mask_b).sum()
total = mask_a.sum() + mask_b.sum()
if total == 0:
return 1.0 # both masks are empty — perfect agreement
return float(2.0 * intersection / total)
def surface_dice(
mask_a: np.ndarray,
mask_b: np.ndarray,
tolerance_mm: float = 2.0,
spacing_mm: tuple[float, ...] = (1.0, 1.0, 1.0),
) -> float:
"""Normalized Surface Dice (NSD) — measures agreement at boundaries.
Standard Dice can be misleadingly high for large structures where the
volume is dominated by interior voxels. Surface Dice focuses on boundary
accuracy, which is where annotation disagreements actually live.
tolerance_mm: maximum distance (in mm) for a surface point to be
considered correctly annotated. Typically 2mm for CT, 3mm for MRI.
"""
from scipy import ndimage
# Compute distance transforms
dt_a = ndimage.distance_transform_edt(~mask_a, sampling=spacing_mm)
dt_b = ndimage.distance_transform_edt(~mask_b, sampling=spacing_mm)
# Extract surfaces (boundary voxels)
struct = ndimage.generate_binary_structure(mask_a.ndim, 1)
surface_a = np.logical_xor(mask_a, ndimage.binary_erosion(mask_a, struct))
surface_b = np.logical_xor(mask_b, ndimage.binary_erosion(mask_b, struct))
# Count surface points within tolerance
surface_a_close = (dt_b[surface_a] <= tolerance_mm).sum()
surface_b_close = (dt_a[surface_b] <= tolerance_mm).sum()
total_surface = surface_a.sum() + surface_b.sum()
if total_surface == 0:
return 1.0
return float((surface_a_close + surface_b_close) / total_surface)
Surface Dice is often more informative than volumetric Dice for evaluating annotation consistency. Two radiologists might achieve Dice > 0.90 on liver segmentation while disagreeing substantially on the boundary near the portal vein. Standard Dice masks this disagreement because the liver interior dominates the volume. Surface Dice with a 2mm tolerance reveals it.
We compute both metrics across all samples annotated by multiple radiologists. Samples below the structure-specific threshold are automatically routed for adjudication by a senior annotator.
Segmentation Models in Medical Imaging
With clean data in hand, the modeling side is comparatively straightforward. 3D UNet variants remain the workhorse architecture for volumetric medical segmentation. Here is a concise implementation:
import torch
import torch.nn as nn
class ConvBlock3D(nn.Module):
"""Two 3x3x3 convolutions with instance norm and LeakyReLU.
Instance norm (not batch norm) is standard in medical imaging because
batch sizes are typically very small (2-4) due to memory constraints
of 3D volumes, making batch statistics unreliable.
"""
def __init__(self, in_ch: int, out_ch: int):
super().__init__()
self.block = nn.Sequential(
nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
nn.InstanceNorm3d(out_ch, affine=True),
nn.LeakyReLU(0.01, inplace=True),
nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
nn.InstanceNorm3d(out_ch, affine=True),
nn.LeakyReLU(0.01, inplace=True),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.block(x)
class UNet3D(nn.Module):
"""3D UNet for volumetric medical image segmentation.
Design choices reflecting production medical imaging constraints:
- Instance normalization instead of batch norm (small batch sizes)
- LeakyReLU to avoid dead neurons in sparse segmentation targets
- Deep supervision outputs at multiple decoder levels
- Channel progression [32, 64, 128, 256, 512] balances capacity
and memory — a 128x128x128 input at batch size 2 uses ~11GB VRAM
"""
def __init__(self, in_channels: int = 1, num_classes: int = 3):
super().__init__()
channels = [32, 64, 128, 256, 512]
# Encoder path
self.encoders = nn.ModuleList()
self.downsamplers = nn.ModuleList()
ch_in = in_channels
for ch_out in channels:
self.encoders.append(ConvBlock3D(ch_in, ch_out))
self.downsamplers.append(
nn.Conv3d(ch_out, ch_out, kernel_size=2, stride=2)
)
ch_in = ch_out
# Bottleneck
self.bottleneck = ConvBlock3D(channels[-1], channels[-1])
# Decoder path
self.upsamplers = nn.ModuleList()
self.decoders = nn.ModuleList()
for i in range(len(channels) - 1, -1, -1):
ch_out = channels[i]
ch_in = channels[i] * 2 if i < len(channels) - 1 else channels[i]
self.upsamplers.append(
nn.ConvTranspose3d(
ch_in if i == len(channels) - 1 else channels[i + 1],
ch_out,
kernel_size=2,
stride=2,
)
)
# After concatenation with skip connection
self.decoders.append(ConvBlock3D(ch_out * 2, ch_out))
# Deep supervision heads — auxiliary loss at intermediate resolutions
# helps gradients flow through the decoder and improves convergence
self.deep_supervision_heads = nn.ModuleList([
nn.Conv3d(channels[i], num_classes, kernel_size=1)
for i in range(len(channels))
])
# Final output
self.final_conv = nn.Conv3d(channels[0], num_classes, kernel_size=1)
def forward(
self, x: torch.Tensor
) -> tuple[torch.Tensor, list[torch.Tensor]]:
"""Forward pass.
Returns:
- Main output at full resolution: (B, num_classes, D, H, W)
- List of deep supervision outputs at decreasing resolutions
"""
# Encoder
skip_connections = []
for encoder, downsampler in zip(self.encoders, self.downsamplers):
x = encoder(x)
skip_connections.append(x)
x = downsampler(x)
x = self.bottleneck(x)
# Decoder with skip connections
deep_outputs = []
for i, (upsampler, decoder) in enumerate(
zip(self.upsamplers, self.decoders)
):
x = upsampler(x)
skip = skip_connections[-(i + 1)]
# Handle size mismatches from odd input dimensions
if x.shape != skip.shape:
x = nn.functional.pad(
x,
[0, skip.shape[4] - x.shape[4],
0, skip.shape[3] - x.shape[3],
0, skip.shape[2] - x.shape[2]],
)
x = torch.cat([x, skip], dim=1)
x = decoder(x)
# Deep supervision output at this resolution
ds_idx = len(self.decoders) - 1 - i
deep_outputs.append(self.deep_supervision_heads[ds_idx](x))
output = self.final_conv(x)
return output, deep_outputs
A few implementation notes:
Instance normalization rather than batch normalization is essential for medical imaging. With typical batch sizes of 2-4 (limited by the memory footprint of 3D volumes), batch statistics are too noisy to be useful. Instance norm computes statistics per-sample, per-channel, sidestepping this issue entirely.
Deep supervision — computing auxiliary losses at intermediate decoder resolutions — is one of the highest-impact architectural choices for 3D medical segmentation. It addresses the vanishing gradient problem that is particularly acute in deep 3D networks and provides a regularization effect. In practice, we weight the auxiliary losses with an exponential decay (e.g., 0.5, 0.25, 0.125) and sum them with the main loss.
Strided convolutions for downsampling instead of max pooling. Max pooling discards spatial information that the decoder needs to reconstruct precise boundaries. Strided convolutions let the network learn what information to preserve.
The Impact of Data Quality: A Concrete Example
To make this concrete, here are typical results from an annotation cleaning cycle on a liver segmentation project:
| Stage | Dataset Size | Mean Dice | p95 Dice | Annotation Error Rate |
|---|---|---|---|---|
| Before cleaning | 850 volumes | 0.91 | 0.78 | ~8% (estimated) |
| After automated flagging | 850 volumes | — | — | 67 samples flagged |
| After expert review | 850 volumes | — | — | 52 confirmed errors |
| After re-annotation | 850 volumes | 0.94 | 0.87 | <1% (estimated) |
The 3-point improvement in mean Dice is meaningful, but the 9-point improvement at p95 is where the clinical impact lies. The tail of the performance distribution — the cases the model gets most wrong — is exactly where annotation errors concentrate, and it’s exactly where clinical risk is highest.
Annotation Quality in Regulated Settings
For medical devices under FDA or CE regulatory pathways, annotation quality is not just a performance concern — it’s an audit requirement. Regulators expect documentation of:
- Annotator qualifications: Board-certified radiologists with subspecialty expertise
- Annotation protocol: Written guidelines with visual examples of edge cases
- Inter-annotator agreement: Quantified using the metrics described above, with pre-specified acceptability thresholds
- Adjudication process: How disagreements are resolved, by whom, and how those decisions are documented
- Traceability: Every annotation linked to the annotator, timestamp, and protocol version
Building the loss-tracking and agreement-metric infrastructure described in this post isn’t just good engineering practice — in regulated contexts, it’s effectively a compliance requirement. The alternative is manual review of every sample, which doesn’t scale and is itself error-prone.
Conclusion
In medical AI pipelines, data quality engineering is just as critical as model development — and frequently delivers larger performance improvements per engineering hour invested. The combination of automated loss tracking, statistical outlier detection, and targeted expert review creates a scalable annotation quality pipeline that improves with each iteration. Investing in this infrastructure early saves enormous cost downstream, both in model performance and in regulatory overhead.
Mattia Gaggi is an applied machine learning engineer working on production ML systems.
Last updated 2026-04-13