RABR-Net Refines Cell Boundaries That Dice Scores Miss

Why a great Dice score can still fail you

White blood cell morphology matters in real medicine. Clinicians and laboratory specialists read the appearance of these cells to judge cell size, cytoplasmic extent, nuclear shape, nuclear lobulation, and the ratio of nucleus to cytoplasm. Those cues feed downstream tasks such as cell classification, abnormal cell detection, and leukemia related screening. Getting the cell outline right is not a computer vision nicety here. It is the structural foundation for the measurements that follow.

Here is where it gets interesting. The biologically meaningful information sits at the boundaries, at the edge between cytoplasm and background and between cytoplasm and nucleus. A small contour error barely moves a global overlap score like Dice, because Dice is dominated by the large, easy interior of the cell that any decent model already gets right. But that same small error can substantially distort nuclear area, cytoplasmic volume, and the nucleus to cytoplasm ratio. A model can post a high global Dice and still be unreliable for the morphology it was supposed to support.

There is a second problem that runs alongside the first. Deep segmentation models, especially U-Net variants trained to maximize region overlap, tend to be overconfident in exactly the ambiguous pixels where they should hesitate. In blood smear microscopy that ambiguity is everywhere, at the cytoplasm background interface and around irregular or multilobed nuclei, made worse by staining and illumination differences. A standard pipeline outputs one final mask and says nothing about where it is unsure or whether the boundary it drew can be trusted. So boundary errors become hard to even find, let alone correct. Related work on cell and nucleus segmentation shows how much downstream analysis rides on getting these contours right.

The idea, treat reliability as its own goal

The paper, by Anima Kujur at the Interdisciplinary Centre for Scientific Computing at Heidelberg University, argues that a segmentation system meant for morphology should answer more than what is the predicted mask. It should also answer where is the model uncertain, are the boundary errors localized, and can a refinement correct meaningful contour errors without damaging the regions that were already right. Those questions motivate a reliability aware view of the task, where uncertainty, boundary quality, calibration, robustness, and interpretability are all evaluated together rather than hidden behind one accuracy number.

The answer is a Reliability-Aware Boundary Refinement Network, or RABR-Net, built as a two stage pipeline. First a strong base segmenter produces an initial cytoplasm and nucleus prediction along with its pixel wise probabilities. Second, uncertainty and boundary derived cues are combined into a reliability representation that guides a gated residual refiner. The refiner does not redraw the whole mask. It predicts small local corrections and applies them only through a learned spatial gate, which keeps the process focused on the boundary and, just as important, keeps it interpretable.

Stage one, a strong base and a cache

The base segmenter is a UNet++ architecture with an EfficientNet-B4 encoder pretrained on ImageNet, trained on the three class problem of background, cytoplasm, and nucleus. It is not a new backbone, and that is deliberate. The point of the paper is not to invent a better segmenter but to characterize and correct the reliability of a good one. The model was even selected with boundary quality in mind, using a validation score that adds a quarter weight of Boundary Dice on top of ordinary Dice, so the base already tries to balance overlap against contour accuracy.

A neat engineering choice makes the whole study reproducible. Rather than retrain the base model when developing the refiner, the base predictions were cached, storing the image, the ground truth, the logits, the class probabilities, and the test time augmentation variance for each image. The base can be trained once, its predictions cached once, and the refiner trained and evaluated many times without rerunning the expensive base. It is the kind of modularity that makes an idea easy for others to build on.

Stage two, a map of what to distrust

The heart of the method is the reliability representation. The refiner does not just see the image and the base probabilities. It sees a stack of signals that together say where the base prediction is shaky. Five of these are uncertainty and boundary maps, and they combine with the image and probabilities into an eleven channel tensor.

Each signal catches a different kind of doubt. Predictive entropy flags pixels where the class probabilities are spread out rather than peaked. Margin uncertainty flags pixels where the top two classes are almost tied, which is common right at the cytoplasm nucleus border.

Test time augmentation variance captures how much the prediction wobbles when the image is flipped four ways and averaged back, a cheap read on instability. A probability gradient map, computed with a Sobel style filter on the foreground probability, lights up where the class changes fast, which is the contour. And a soft boundary hint, from local maximum minus minimum pooling, marks the edge directly. None of these alone is a reliable error detector, so the method fuses them rather than trusting any single one. This instinct, letting a model expose where it is unsure, echoes other work on segmenters that flag their own uncertain pixels.

The gated residual refiner

Given that reliability tensor and the base logits, the refiner predicts two things at once, a residual logit correction and a spatial gate. The gate is a per pixel value between zero and one that decides how much of the correction is actually applied.

Read what the gate does. Where it is near zero, the refined logits stay equal to the base logits, so the confident interior of the cell is preserved untouched. Where it is near one, the correction takes full effect. This biases the whole system toward local repair rather than wholesale replacement, and it is enforced during training by a small penalty on the total amount of gating, which discourages the refiner from meddling where it is not needed.

The training objective mixes four losses, cross entropy and multiclass Dice for overlap, focal loss to concentrate on hard pixels, and a differentiable boundary Dice loss that operates on a soft boundary map. That boundary term is the one that matters most here, because the whole complaint about the base model was boundary accuracy. The refiner adds its gate penalty on top.

Finally, a light biological post processing step enforces two facts about white blood cells that a pixel classifier does not inherently know, that the foreground should be one main connected component and that nucleus pixels should sit inside that foreground. Small stray islands are removed and holes are filled. It is a modest, sensible cleanup rather than a learned module.

What the numbers say

The evaluation uses the WbcMSBench dataset from the MedSegBench suite, with an official split of 280 training, 40 validation, and 80 held out test images at 256 by 256 resolution. The test set was touched only for final reporting. The comparison that matters is the cached base prediction against the full refined result.

Look past the Dice column. It moves by about a thousandth, which sounds like nothing, and in isolation it would be nothing. The story is in the boundary metrics. Boundary Dice rises from 0.3448 to 0.3611, and the HD95 surface error drops from 3.0354 to 2.8274. HD95 captures the large outlier contour deviations, so a lower value means the refined mask is not just slightly better on average but less likely to contain a bad excursion at the edge. A representative case in the paper makes the split vivid, where Dice improved by only 0.004 while Boundary Dice improved by 0.062 on the same image. The refinement is doing exactly and only what it claims.

Because the base is so strong, the honest way to judge the gains is per image statistics, not just the mean. The paired analysis backs the improvement up.

Dice, Boundary Dice, and HD95 all improved with statistical significance. And then there is the last row, which the author does not hide and neither should this article.

The calibration surprise, and why it matters

This is the most interesting finding in the paper, precisely because it is a negative one. A less careful write up would claim RABR-Net improves trustworthiness across the board. It does not. It improves spatial and boundary reliability while probability calibration gets worse, and saying so out loud is what separates a credible result from an oversold one.

Why would fixing the geometry hurt the confidence? Several reasons line up. The refiner is trained mostly with segmentation and boundary losses that reward spatial accuracy, not calibrated probability. The gated correction modifies logits in the uncertain regions, which can improve the final label while distorting the underlying confidence distribution. And the biological post processing changes the discrete mask without ever recalibrating the probability map behind it. The takeaway is not that the method failed. It is that spatial correctness and calibrated confidence are related but distinct objectives, and improving one does not hand you the other for free. Future refinement schemes would need explicit calibration aware training, temperature scaling, or uncertainty consistency terms to fix both at once.

Which parts pull weight

A staged ablation traces where the improvement comes from, and the answer is clear.

The reliability guided refiner is the workhorse. It lifts Boundary Dice from 0.3448 to 0.3609 and pulls HD95 down from 3.0354 to 2.8545 in one step. The biological post processing adds only a small structural consistency benefit on top. So the gain is not a post processing trick, it is the uncertainty guided residual correction doing the real work, which is the whole thesis of the paper confirmed by removing everything else.

The clinical translation gap

Good boundary numbers on a benchmark are not the same as a tool a laboratory can rely on, and this work sits some distance from that use. The most important reason is scope. The method was developed and tested on a single public dataset with one fixed split. Blood smear appearance varies enormously with staining protocol, scanner, and laboratory, and a model that improves boundaries on one benchmark has not been shown to hold up across that variability. The paper’s own robustness experiment, applied to the base model, shows how fragile microscopy models can be to exactly these shifts. Under a severe contrast change the base Dice fell to 0.8815 and Boundary Dice to 0.2588, and a strong stain shift pushed HD95 up sharply, while blur, noise, and compression were tolerated far better.

The second reason the clinic is still far off is the calibration result above. A morphology support tool that becomes less calibrated after refinement is exactly the kind of thing that needs fixing before anyone leans on its confidence to flag a suspicious cell. And the third is that the robustness test covered only the base model, not the full refinement pipeline, so how the refiner behaves under a stain shift is genuinely unknown. The paper is careful to describe RABR-Net as a segmentation support method that does not make clinical decisions, and that framing is the right one.

Honest limitations

The single dataset scope is the headline limitation. External validation on independent microscopy datasets, across staining protocols, acquisition devices, and laboratory conditions, is necessary before anyone calls this generalizable, and the author says so plainly. One benchmark with a fixed split is a proof of concept, not evidence of robustness across the messy reality of clinical labs.

Dataset and design bias run alongside it. The refiner depends entirely on the quality of the base model. When the base is confidently wrong, the uncertainty cues do not highlight the true error region, and the refiner cannot fix a mistake it cannot see. The failure cases in the paper show exactly this, where uncertainty does not align with the real segmentation error and the refiner either leaves an error unresolved or introduces a small new one. The method inherits its base model’s blind spots.

Generalization concerns are concrete rather than abstract. The biological post processing assumes one main foreground component with the nucleus inside it, which is reasonable for isolated single cells but would struggle with overlapping cells or more complex cell types that need topology aware constraints. And the interpretability the method offers, its uncertainty maps and correction overlays, helps a human inspect the output but does not replace expert validation of whether the corrected masks actually improve downstream morphology measurements. That last check, connecting boundary refinement to real measurement accuracy, has not been done yet.

Above all of this sits one plain fact. This is a preprint by a single author, and it has not been through peer review. Its results should be read as a well argued early signal, not a settled conclusion.

Why the approach matters anyway

Set the caveats aside for a moment and the paper makes a point worth keeping. It reframes what reliable segmentation even means. For years the field optimized global overlap and reported one Dice number. RABR-Net insists that a model for morphology sensitive analysis should be judged on boundary quality, calibration, robustness, uncertainty localization, and inspectable corrections, all together. Evaluating trustworthiness beyond global Dice is a useful correction whether or not this specific network is the one that reaches the clinic. The same reliability first spirit runs through recent medical segmentation work such as diagnosis conditioned cardiac segmentation and few shot medical segmentation that must generalize past its training hospital.

The design also travels. A frozen strong base, a fused uncertainty map that says where to distrust it, and a gated residual that only corrects the uncertain regions is a pattern that would fit any dense prediction task where errors concentrate at boundaries and confident interiors should be left alone. The gate makes the whole thing auditable, which is rare and valuable in a field that too often ships an opaque second network and calls it refinement.

Reference implementation in PyTorch

The code below is a runnable reconstruction of the RABR-Net idea based on the paper’s equations, the uncertainty and boundary maps, the reliability tensor, and the gated residual refiner. A compact convolutional segmenter stands in for the UNet++ EfficientNet-B4 base so the file runs without downloads. It includes entropy, margin, test time augmentation variance, probability gradient and soft boundary maps, the eleven channel reliability tensor, the gated residual refiner that outputs a correction and a spatial gate, the compound segmentation loss with a boundary Dice term, the gate regularized refiner loss, and a smoke test on dummy tensors. Swap in the real base model and dataset for actual experiments.

# rabr_net_reference.py
# Runnable reconstruction of the RABR-Net reliability guided refiner.
# Replace BaseSegmenterStub with a UNet++ EfficientNet-B4 for real runs.

import torch
import torch.nn as nn
import torch.nn.functional as F

C = 3  # classes: 0 background, 1 cytoplasm, 2 nucleus


class BaseSegmenterStub(nn.Module):
    """Stand in for the frozen UNet++ base. Returns logits (B, C, H, W)."""
    def __init__(self, ch=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, ch, 3, 1, 1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True),
            nn.Conv2d(ch, ch, 3, 1, 1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True),
            nn.Conv2d(ch, C, 1))

    def forward(self, x):
        return self.net(x)


def soft_boundary(q):
    """B(q) = maxpool(q) - minpool(q), min via negated maxpool."""
    mx = F.max_pool2d(q, 3, 1, 1)
    mn = -F.max_pool2d(-q, 3, 1, 1)
    return mx - mn


def reliability_tensor(base, x, tta_var, eps=1e-6):
    """Build the 11 channel reliability tensor from base logits."""
    p = F.softmax(base(x), dim=1)                    # (B, C, H, W)
    # predictive entropy, normalized to [0, 1]
    ent = -(p * torch.log(p + eps)).sum(1, keepdim=True) / torch.log(torch.tensor(float(C)))
    # margin uncertainty from the top two class probabilities
    top2 = p.topk(2, dim=1).values
    margin = 1.0 - (top2[:, 0:1] - top2[:, 1:2])
    # foreground probability, its Sobel gradient, and a boundary hint
    fg = p[:, 1:].sum(1, keepdim=True)
    kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]],
                      dtype=fg.dtype, device=fg.device).view(1, 1, 3, 3) / 8
    ky = kx.transpose(2, 3)
    gx, gy = F.conv2d(fg, kx, padding=1), F.conv2d(fg, ky, padding=1)
    grad = torch.sqrt(gx ** 2 + gy ** 2 + eps)
    grad = grad / (grad.amax(dim=(2, 3), keepdim=True) + eps)
    bd = soft_boundary(fg)
    R = torch.cat([x, p, ent, margin, tta_var, grad, bd], dim=1)  # 3+3+5=11
    return R, torch.log(p + eps)                      # tensor and base logits


class GatedResidualRefiner(nn.Module):
    """Predicts a residual logit correction and a spatial gate."""
    def __init__(self, in_ch=11, ch=48):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Conv2d(in_ch, ch, 3, 1, 1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True),
            nn.Conv2d(ch, ch, 3, 1, 1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True))
        self.delta_head = nn.Conv2d(ch, C, 1)          # residual logits
        self.gate_head = nn.Conv2d(ch, 1, 1)           # spatial gate

    def forward(self, R, base_logits):
        h = self.enc(R)
        delta = self.delta_head(h)
        gate = torch.sigmoid(self.gate_head(h))       # G in [0, 1]
        refined = base_logits + gate * delta          # z_r = z_b + G * dz
        return refined, gate


def dice_loss(probs, oh, eps=1.0):
    inter = (probs * oh).sum((2, 3))
    denom = probs.sum((2, 3)) + oh.sum((2, 3))
    return (1 - ((2 * inter + eps) / (denom + eps))[:, 1:]).mean()


def boundary_dice_loss(probs, oh, eps=1.0):
    bp, by = soft_boundary(probs[:, 1:]), soft_boundary(oh[:, 1:])
    inter = (bp * by).sum((2, 3))
    denom = bp.sum((2, 3)) + by.sum((2, 3))
    return (1 - (2 * inter + eps) / (denom + eps)).mean()


def compound_loss(logits, target):
    """ce 0.35, dice 0.35, focal 0.10, boundary dice 0.20."""
    probs = F.softmax(logits, dim=1)
    oh = F.one_hot(target, C).permute(0, 3, 1, 2).float()
    ce = F.cross_entropy(logits, target)
    pt = (probs * oh).sum(1).clamp(min=1e-6)
    focal = (-(1 - pt) ** 2 * torch.log(pt)).mean()
    return (0.35 * ce + 0.35 * dice_loss(probs, oh)
            + 0.10 * focal + 0.20 * boundary_dice_loss(probs, oh))


def refiner_loss(refined, gate, target, lam_g=0.005):
    return compound_loss(refined, target) + lam_g * gate.mean()


if __name__ == "__main__":
    base = BaseSegmenterStub()
    for p in base.parameters():
        p.requires_grad = False                       # base is frozen, cached in practice
    refiner = GatedResidualRefiner()
    opt = torch.optim.AdamW(refiner.parameters(), lr=3e-4, weight_decay=1e-4)
    B, H, W = 2, 256, 256
    x = torch.rand(B, 3, H, W)
    tta_var = torch.rand(B, 1, H, W)                   # cached TTA variance map
    y = torch.randint(0, C, (B, H, W))
    for step in range(3):
        R, base_logits = reliability_tensor(base, x, tta_var)
        refined, gate = refiner(R, base_logits)
        loss = refiner_loss(refined, gate, y)
        opt.zero_grad(); loss.backward(); opt.step()
        print("step", step, "loss", round(loss.item(), 4),
              "mean gate", round(gate.mean().item(), 4))

Conclusion

The core achievement of RABR-Net is to move the conversation about cell segmentation from global overlap to boundary reliability, and to back that move with a method that measurably improves the contours while leaving the confident interior alone. On the held out test set it improved Boundary Dice from 0.3448 to 0.3611 and cut the HD95 surface error from 3.0354 to 2.8274, gains that a near flat Dice score would have hidden entirely. For morphology, where the boundary is the measurement, that is the improvement that counts.

The conceptual shift worth keeping is that trustworthiness is not one number. This paper pulls apart spatial accuracy, boundary quality, calibration, robustness, and interpretability, and shows they do not move together. The clearest evidence is the calibration result, where a sharper mask came with worse confidence estimates, expected calibration error rising from 0.0079 to 0.0145 on every test image. A weaker paper would have buried that. Reporting it is what makes the rest of the claims believable, and it hands the field a specific open problem rather than a false all clear.

The design also travels beyond blood cells. A frozen strong base, a fused map of where to distrust it, and a gated residual that only edits the uncertain regions is a general recipe for any dense prediction task whose errors live at the boundary. The learned gate makes the correction inspectable, so a human can see which pixels were changed and why, which is a real answer to the black box worry that dogs second stage refinement networks.

The honest limitations keep expectations where they belong. This is a single author preprint on one dataset with a fixed split, the refinement worsens calibration, the robustness test covered only the base model, and the whole thing depends on a base segmenter whose confident mistakes it cannot see. The biological cleanup assumes tidy isolated cells. None of that erases the contribution. It frames it as a well argued starting point, and names the next steps, external validation, calibration aware training, and a check that better boundaries actually improve downstream morphology measurements.

For anyone building segmentation for a measurement task, the practical message is compact. Stop trusting global Dice as your only gauge. Measure the boundary, measure calibration, and treat them as separate goals. Let the model expose where it is unsure, and correct only there. RABR-Net shows the payoff is real at the edge, and it is candid that the confidence estimates still need work, which is exactly the kind of honesty a morphology sensitive clinical setting will demand.

Frequently asked questions

What does RABR-Net actually do?

RABR-Net refines the output of a white blood cell segmentation model. It takes the base model’s prediction along with several uncertainty and boundary maps, and a second network makes small corrections only at the uncertain cytoplasm and nucleus boundaries, controlled by a learned spatial gate. The confident interior of the cell is left unchanged.

Why not just use the Dice score to measure a segmentation model?

Global Dice is dominated by the large, easy interior of a cell, so it barely changes when the boundary is a few pixels off. But the boundary is what determines nuclear shape, cytoplasm area, and the nucleus to cytoplasm ratio, the measurements that matter for morphology. A model can score a high Dice and still draw unreliable contours, which is why RABR-Net emphasizes Boundary Dice and the HD95 contour error instead.

Did the refinement improve everything?

No, and the paper is clear about it. Boundary Dice, HD95, and global Dice all improved with statistical significance, but calibration got worse, with expected calibration error rising from 0.0079 to 0.0145 on all 80 test images. Improving the geometry of the mask did not improve the trustworthiness of the model’s confidence, which are two separate objectives.

How does RABR-Net decide where to make corrections?

It builds an eleven channel reliability tensor that combines the image, the base probabilities, and five signals of doubt, predictive entropy, margin uncertainty, test time augmentation variance, a probability gradient, and a soft boundary hint. A learned gate then reads this and allows corrections only where the base prediction looks unreliable, which in practice concentrates along cell and nuclear contours.

Is RABR-Net ready to use in a laboratory?

No. It is an unreviewed preprint tested on a single dataset with a fixed split, its calibration worsened after refinement, and its robustness was only tested on the base model. The authors describe it as a segmentation support method rather than a diagnostic system. External validation across staining protocols, scanners, and laboratories would be needed before any clinical use.

What is the reusable idea beyond blood cells?

The pattern is a frozen strong base model, a fused uncertainty map that marks where to distrust it, and a gated residual that only corrects the uncertain regions while leaving confident areas alone. That is applicable to any dense prediction task where errors concentrate at boundaries, and the learned gate makes the correction inspectable rather than a black box.