AI Tech Inspire

Why U-Nets Miss Cosmic Filaments—and Practical Fixes for Ultra-Thin Segmentation

When segmentation models meet the universe’s hairline structures, things get weird. A recent project tried to map 1‑pixel cosmic filaments from galaxy distributions and watched a familiar pattern unfold: U-Net predictions collapsed into fuzzy blobs around clusters while the real filaments stayed invisible. If you’ve ever chased roads, vessels, cracks, or nerves, you’ve probably seen the same failure mode. At AI Tech Inspire, we spotted this case as a perfect lens on a broader question: why do otherwise strong vision models struggle with ultra-thin topology—and what actually works?

Quick facts from the experiment

  • Input: 256×256 multi-channel per cluster: raw galaxy points, smoothed density, gradient magnitude, and a radial distance map.
  • Target: 1-pixel-wide filament skeleton generated with DisPerSE (a topological filament finder).
  • Dataset size: ~1900 samples, consistent and aligned; masks sit on density ridges.
  • Observed failure: predictions degenerate into fuzzy blobs or circular structures around cluster centers.
  • Metrics: Dice 0.08–0.12, Dilated Dice 0.18–0.23, IoU ≈0.00–0.06.
  • Attempts so far: U-Net, Dice/BCE/Tversky/Focal Tversky losses, 5-channel input, heavy augmentation, oversampling positives, LR schedules, longer training, and thick→thin mask variants.
  • Open questions: Are U-Nets a poor fit for super-thin topology? Should targets be thicker or replaced by distance maps? Would other models or pretraining help?

Why ultra-thin skeletons break vanilla segmentation

Skinny structures behave like adversarial examples for pixel-wise segmentation. Three issues dominate:

  • Extreme class imbalance: a 1‑px skeleton may be <1% of the image. Losses like BCE and plain Dice are dominated by easy negatives, encouraging smooth “blobs” that reduce false positives but erase the centerline.
  • Label aliasing: 1‑px ground truth is discretization-sensitive—subpixel shifts during augmentation or down/up-sampling move negatives to positives with a 1‑px error. The model “learns” that edges are unreliable and regresses to fat, safe shapes.
  • Topology blindness: standard losses don’t penalize broken connectivity. A prediction that misses a continuous filament but paints a circular halo can still look statistically “ok” to the loss.

Key takeaway: If your target is a 1‑pixel skeleton, your loss and architecture must explicitly value connectivity and centerlines, not just pixel accuracy.

Targets that play to your model’s strengths

Instead of forcing a network to guess a single-pixel curve, teach it richer, more stable signals and collapse to a skeleton later:

  • Distance transform (DT) regression: Predict a continuous signed distance or truncated distance map to the skeleton. Train with L1/L2 or a combo with BCE on a thick mask. Then apply ridge detection + skeletonization at inference. This reduces aliasing and emphasizes geometry.
  • Soft-thick labels: Convolve the skeleton with a small Gaussian to make a soft tube; use soft Dice/Tversky. The network learns a centerline likelihood field, not binary spikes.
  • Curriculum: thick → thin: Start with 3–5 px tubes; progressively anneal to thinner targets. Keep a multi-task head that always sees a stable thick mask even in late training.
  • Orientation fields: Predict both filament probability and local tangent orientation (2D unit vector). Later, track along high-probability, consistent orientations to recover clean filaments.

If you’re prototyping in PyTorch or TensorFlow, it’s straightforward to stack heads: p(y), d(x), theta(x), with a composite loss like loss = 0.4*clDice + 0.3*Tversky + 0.3*L1(DT).

Architecture pivots worth trying

Classic U-Net is a good baseline, but hairline structures often respond better to edge- and high-resolution-biased backbones:

  • Edge detectors: HED/RCF or modern boundary detectors with deep supervision specialize in thin, connected contours. Use them to produce a filament probability map and post-skeletonize.
  • High-resolution features: HRNet retains high-res representations throughout, preserving precise localization better than heavy downsampling paths.
  • Attention U-Net / U-Net++: Enhanced skip connections and attention gates help preserve small structures across scales.

Whatever you choose, enable deep supervision on intermediate scales and constrain strides so the effective receptive field doesn’t erase the 1‑px signal. If GPU memory allows (hello, CUDA), keep input at native 256×256 without aggressive downsampling.

Losses that reward topology and centerlines

  • clDice: A topology-preserving loss for tubular structures that aligns centerlines between prediction and ground truth. Often paired with Dice/Tversky for stability.
  • Boundary loss / Hausdorff loss: Distance-aware objectives that penalize boundary deviations more than interior area.
  • Focal Tversky: Good under high imbalance; tune alpha/beta so false negatives hurt more than false positives.
  • Weighted BCE: Explicit per-class weights (or online hard example mining) to keep gradients flowing on the tiny positive class.

A simple, effective recipe many teams report for vessels/roads:

# pseudo-loss blend for ultra-thin targets
loss = 0.5 * clDice(pred, gt_skeleton) \
     + 0.3 * Tversky(pred, gt_thick, alpha=0.7, beta=0.3) \
     + 0.2 * L1(pred_dt, gt_distance)

Post-processing: where the skeleton actually emerges

Even with better targets, the cleanest centerlines typically appear after post-processing:

  • Non-maximum suppression along gradient/orientation to thin ridges.
  • Morphological thinning or skeletonize on a high-confidence tube.
  • Graph refinement: Turn the probability map into a graph, prune spurs, and connect gaps with shortest paths respecting orientation and density ridges.

Given the DisPerSE pedigree of the labels, an alternative is to predict a ridgeness field and hand it to the same topological pipeline (or a simplified variant) to produce filaments. That hybrid keeps the physics/topology prior intact while letting the network denoise and enrich the signal.

Evaluation: measure what you actually care about

Pixel IoU near zero is common with 1‑px targets. Consider:

  • Centerline Dice (skeleton Dice): Dice after skeletonizing predictions at a matched scale.
  • Symmetric centerline distance: Average distance between predicted and true centerlines.
  • Dilated Dice curves: Dice as you dilate the GT from 1→5 px; report AUC rather than a single radius.

These metrics reflect whether you captured the filament’s path and connectivity, not whether every pixel aligned perfectly.

Data and training hygiene that matter more than usual

  • Augment gently: Subpixel translations, small rotations, and elastic deforms can break a 1‑px label. Favor intensity/contrast, mild geometric jitter, and multi-scale crops over heavy warps.
  • Positive-biased sampling: Use crops/patches centered on filaments to keep the batch informative.
  • Deep supervision at multiple scales keeps gradients alive for thin features.
  • Longer schedules with cyclical LR can help, but watch for the model collapsing to background; add a connectivity-aware loss early.

If you’re iterating quickly in notebooks, a tiny ritual helps: press Shift+Enter, visualize both the thick mask and predicted DT, and only then inspect the skeleton. It’s remarkable how often the centerline is “there,” just not yet thinned.

Why this matters beyond cosmology

Engineers working on medical imaging (vessels, nerves), geospatial mapping (roads, rivers), manufacturing (cracks, fibers), and materials science (grain boundaries) face the same topology problem. The consistent signal across domains: 1‑px targets are not model-friendly. Teach the model a field (probability, distance, orientation), then derive the skeleton with post-processing.

For those who want off-the-shelf starting points, consider edge/sketch models on Hugging Face, adapt them to your channels, and fine-tune with a topology-aware loss. Whether you run on PyTorch or TensorFlow, the recipe transfers.

A compact game plan to try next

  • Switch targets to a 3–5 px soft tube + DT map; keep the 1‑px skeleton only for clDice supervision.
  • Adopt an edge-centric model (e.g., HED) or high-res backbone (HRNet) with deep supervision.
  • Blend clDice + Tversky + L1(DT); weight positives in BCE if you keep it.
  • Post-process with NMS + skeletonize + graph pruning; evaluate with centerline metrics.
  • Optionally, predict ridgeness and run a topo-aware extractor rather than direct skeletons.

“Predict fields, not lines. Then extract the lines you care about.”

If your current runs plateau around Dice 0.10, this shift often unlocks an order-of-magnitude jump in centerline metrics—and more importantly, produces filaments that look and behave like the structures domain experts expect.

Recommended Resources

As an Amazon Associate, I earn from qualifying purchases.