UniPro: Unified Multi-Modal Medical Image Segmentation from 2D to 3D in One Framework

Medical image segmentation has long been troubled by two disconnected axes: segmentation paradigm and data dimensionality. Semantic segmentation, in-context segmentation, and interactive segmentation are typically developed as three separate technical tracks, each further subdivided into specialized versions for 2D images or 3D volumetric data. This fragmentation may appear orderly in academic papers, but it fits awkwardly into real clinical workflows. UniPro, a model proposed in a recent arXiv preprint, attempts to stitch these scattered capabilities back together within a single unified framework.

The Real Pain Points of Clinical Segmentation

The starting point for a clinical segmentation task can vary enormously. Sometimes it's automatically initialized by a semantic prediction, sometimes guided by a reference image, and sometimes driven entirely by a clinician's manual clicks. Regardless of the starting point, fine-grained corrections are almost always performed in 2D views — consistent with how radiologists review images slice by slice.

The problem is that for volumetric scans like CT and MRI, a 2D edit a physician makes on one slice must propagate coherently across all remaining slices in the volume. If every slice requires manual correction one by one, the annotation cost becomes prohibitively high. Existing methods tend to excel at either 2D or 3D, or support only one interaction paradigm, making it difficult to cover this flexible and variable real-world workflow.

Core Insight: Propagation as In-Context Segmentation

UniPro's most critical theoretical breakthrough lies in the authors' identification of a previously overlooked equivalence: volumetric propagation and in-context segmentation fundamentally share the same "reference-conditioned prediction" mechanism.

The only difference between the two is where the reference image-mask pairs come from. In in-context segmentation, reference samples come from other cases; in volumetric propagation, reference samples come from already-segmented neighboring slices within the same volume. Since the underlying prediction mechanism is identical, it's entirely natural to handle both tasks with the same model. This perspective unifies two seemingly unrelated problems under the same mathematical framework.

In-context segmentation is a paradigm recently migrated from the "in-context learning" ideas of NLP into the vision domain: at inference time, the model requires no fine-tuning — it simply receives a few reference exemplars (image-mask pairs) and generalizes the same segmentation logic to new query images. This differs fundamentally from traditional supervised training, where conditions are encoded into model weights rather than dynamically provided. In medical imaging, this paradigm is especially suited for few-shot organ or rare lesion scenarios: annotate a small batch of reference slices, and the model can infer boundaries across the remaining unannotated regions. UniPro's key finding is that volumetric propagation is mathematically isomorphic to this mechanism: already-segmented neighboring slices naturally play the role of "reference exemplars," while the slice to be predicted plays the role of the "query image." This equivalence not only saves model parameters but also enables mutual transfer learning between the two tasks — the rich paired data from neighboring slices also reinforces cross-case in-context segmentation capability.

Four Modes Within a Single Slice-Based Framework

Building on this insight, UniPro supports four segmentation modes simultaneously within a unified, slice-centric framework:

Semantic Segmentation

Using class priors as conditional input, the model can predict directly based on category information for target organs or lesions.

In-Context Segmentation

Leveraging reference exemplars — image-mask pairs from other cases — to guide segmentation of the current image. Well-suited for few-shot or novel task scenarios.

Interactive Segmentation

Accepting user clicks as conditions, allowing clinicians to quickly specify or refine target regions through simple interactions.

Propagation Segmentation

Using neighboring-slice predictions as conditions to coherently extend 2D segmentation along the volumetric axis into 3D.

All four modes share the same network architecture; the only difference is the mode-specific conditional input fed to the model. This design eliminates the redundancy of training separate models for each task and enables seamless switching between paradigms.

Bidirectional and 3D Supervision for More Reliable Propagation

While propagation is an elegant mechanism, slice-by-slice propagation can accumulate errors, causing segmentation results to gradually drift in the deeper portions of a volume. To improve propagation reliability, UniPro introduces bidirectional supervision and 3D supervision on top of the traditional per-slice loss.

Bidirectional supervision means propagation is constrained in both directions rather than advancing unidirectionally, reducing single-direction error accumulation. 3D supervision applies regularization to slice-by-slice propagation from the perspective of the entire volume, ensuring spatial consistency across slices in three-dimensional space. Together, these two designs jointly constrain slice-level propagation behavior, producing more coherent and stable final 3D segmentation.

Error accumulation in slice-by-slice propagation is a long-standing problem in computer vision and medical imaging, fundamentally analogous to the "drift" phenomenon in video object tracking: each prediction step depends on the previous output, and slight deviations can amplify linearly or even super-linearly with the number of slices. Bidirectional propagation is a common approach to this problem: the model propagates simultaneously from "start to end" and "end to start," with the two predictions meeting midway to mutually correct each other, preventing errors from accumulating indefinitely in a single direction. 3D supervision provides global consistency constraints through joint optimization of voxel-level loss across the entire volume — even if a local prediction on an intermediate slice deviates, the global loss function pulls it back toward the correct direction. The combination forms a "local-global co-constraint" training strategy that is a key engineering factor behind UniPro's propagation accuracy improvements over slice-by-slice methods.

Experimental Performance and Practical Value

The authors conducted extensive experiments across multiple imaging modalities and different anatomical structures. Results show that UniPro achieves strong performance across all segmentation settings, demonstrating that a single model is fully capable of handling multi-paradigm, cross-dimensional segmentation tasks without sacrificing accuracy for generality.

Perhaps more practically significant is the annotation efficiency improvement it delivers. UniPro supports efficient complete 3D segmentation starting from sparse 2D initialization, significantly reducing the workload of slice-by-slice correction for clinicians. For medical image annotation — an extremely costly process — this "minimal interaction, full propagation" workflow has the potential to meaningfully reduce human effort in both clinical and research settings.

The high cost of medical image annotation stems from multiple factors: high professional barriers requiring licensed radiologists or clinical experts to review each case; workload proportional to slice count, with an abdominal CT routinely containing hundreds of slices; and heterogeneity in annotation standards across different anatomical structures, device parameters, and imaging protocols. Industry estimates suggest that annotating a high-quality medical image segmentation dataset often costs hundreds to thousands of dollars per case. UniPro's "sparse 2D initialization + propagation completion" workflow can compress a physician's active operations from slice-by-slice interaction to careful annotation of a small number of key slices, with the model automatically completing the rest for rapid expert review and correction. This not only reduces absolute annotation hours but also shifts the physician's cognitive burden from repetitive slice-by-slice operations toward higher-value quality control work — aligning well with the practical needs of today's medical AI deployment for "human-machine collaborative" annotation pipelines.

Conclusion

UniPro's value lies not only in its across-the-board metric improvements, but also in the new framework it offers for thinking about medical image segmentation. By identifying the intrinsic unity between propagation and in-context segmentation, it integrates a previously fragmented technical stack into a single model — accommodating the diverse starting points of clinical workflows while solving the core challenge of coherently propagating 2D edits into 3D volumes. For researchers and medical AI teams looking to reduce annotation burden and build general-purpose segmentation systems, this work is well worth close attention.