---
title: 'TP-UNet: Temporal & Topology-Aware Segmentation'
url: https://www.emergentmind.com/topics/tp-unet
type: topic
---

# TP-UNet: Temporal & Topology-Aware Segmentation

TP-UNet refers to two distinct but influential frameworks in medical image segmentation: Temporal Prompt Guided UNet for integrating temporal priors via prompts [2411.11305], and the Topology-Preserving Segmentation Network, leveraging diffeomorphic deformation fields for topological correctness [2202.13331]. Each addresses limitations of standard UNet-based architectures by incorporating domain-specific priors—temporal or topological—into the deep learning pipeline, thereby achieving state-of-the-art results in their respective tasks.

## 1. Temporal Prompt Guided UNet

The Temporal Prompt Guided UNet (TP-UNet) [2411.11305] extends conventional UNet models for 2D medical image segmentation by integrating temporal information using textual prompts. Traditional UNet variants process each axial slice of a 3D scan independently, omitting organ-order and slice-position cues that experienced radiologists leverage intuitively. This often results in confusion among adjacent organs with similar appearances but distinct anatomical progression along the scan axis.

### Motivation and Temporal Priors

The appearance probability of abdominal organs as a function of axial slice position can be modeled by organ-specific normal distributions over the normalized slice index ($i/N \in [0,1]$): for example, the stomach commonly appears in lower-index slices, the small intestine mid-range, and the large intestine in later slices. Standard UNet architectures ignore this temporal prior; TP-UNet injects it directly.

Each slice is paired with an auto-generated prompt:
> "This is {an MRI / a CT} of the {organ} with a segmentation period of {i/N}."

The prompt encodes both the imaging modality, the target organ, and the slice's relative position, directly guiding the model’s expectations for that slice’s content.

## 2. Architecture and Multimodal Integration

TP-UNet is implemented as a multimodal extension of the classic 2D UNet:

- **Backbone**: Four-block encoder with max-pooling, and a decoder with upsampling and skip connections.
- **Text/Image Embedding**: Deepest encoder output $F_m \in \mathbb{R}^{B \times C \times H \times W}$ is paired with a text prompt embedding $F_t \in \mathbb{R}^{B \times L \times D}$, where textual features are extracted with CLIP+LoRA or Electra+SFT encoders.

Two core mechanisms are used between encoder and decoder:

### Semantic Alignment via Contrastive Learning

A bidirectional InfoNCE loss aligns text and image embedding spaces, drawing matched pairs of image and prompt representations together, and separating unmatched pairs. Given image feature $F_{m_i}$ and its corresponding prompt $F_{t_i}$, the losses are

\[
\ell_i^{(F_m \to F_t)} = -\log \frac{\exp(\langle F_{m_i},F_{t_i}\rangle/\tau)}{\sum_{k=1}^N \exp(\langle F_{m_i},F_{t_k}\rangle/\tau)}
\]
\[
\ell_i^{(F_t \to F_m)} = -\log \frac{\exp(\langle F_{t_i},F_{m_i}\rangle/\tau)}{\sum_{k=1}^N \exp(\langle F_{t_i},F_{m_k}\rangle/\tau)}
\]
with total contrastive loss:

\[
\mathcal{L}_{\text{contra}} = \frac{1}{N} \sum_{i=1}^N \left[\lambda \ell_i^{(F_m \to F_t)} + (1-\lambda)\ell_i^{(F_t \to F_m)}\right]
\]

### Cross-Attention Fusion

Aligned image and prompt embeddings $[F_m'; F_t']$ are concatenated and fused via pixel-wise cross-attention, computed as:

\[
\mathcal{F} = \operatorname{softmax}\left(\frac{[F_m';F_t']W^Q \left([F_m';F_t']W^K \right)^\top}{\sqrt{d_k}}\right)\left([F_m';F_t']W^V\right)
\]

The resulting feature map is merged by a $1 \times 1$ convolution and is injected into the decoder via the skip-connection for downstream mask prediction.

## 3. Optimization, Experimental Datasets, and Results

### Loss Functions

Segmentation supervision is achieved by combining binary cross-entropy and Tversky loss:
\[
\mathcal{L}_{\text{seg}} = \mathcal{L}_{\text{BCE}} + \mathcal{L}_{\text{Tversky}}
\]
with optional Dice loss reporting.

### Datasets and Setup

- **UW-Madison GI Tract (MRI):**
  - 26,746 training, 3,820 validation, 7,642 test 2D slices, multi-organ (stomach, small, large intestine)
  - 7:1:2 data split
- **LITS 2017 (CT):**
  - 13,000 3D scans, 10,967 2D slices used for liver segmentation
  - 7:1:2 split

Preprocessing involves CoarseDropout, HorizontalFlip, and ShiftScaleRotate. The optimizer is AdamW with weight decay $1\times10^{-6}$ and initial learning rate $3\times10^{-5}$, cosine-annealing scheduling, and temperature $\tau \in [25, 96.875]$ for contrastive alignment.

### Quantitative Outcomes

| Dataset                | Backbone         | Dice Score | Jaccard | Improvement over UNet |
|------------------------|------------------|------------|---------|----------------------|
| UW-Madison (MRI)       | Electra          | 0.9266     | 0.8943  | +4.44% mDice         |
| LITS (CT)              | Not specified    | 0.9125     | 0.8780  | +6.08% mDice         |

TP-UNet sets state-of-the-art performance, outperforming Swin-UNet by +1.3% Dice on UW-Madison and previous LITS SOTA by +9.21% [2411.11305]. Qualitative analysis demonstrates sharper boundaries and fewer false positives, especially for slices where organ appearance is ambiguous.

### Ablation Studies

Removal of temporal prompts reduces mDice by 2.1% (UW-Madison), and eliminating prompts entirely causes a 5.36% drop (LITS). Semantic alignment and cross-attention fusion contribute additional performance. The dominant contributor is the temporal prompt.

## 4. Topology-Preserving Segmentation Network (TPSN)

The Topology-Preserving Segmentation Network (TPSN, also referred to as TP-UNet in some sources) [2202.13331] addresses the problem of topology errors in medical segmentation. TPSN outputs a diffeomorphic deformation field that warps a simple binary template mask (with the desired connectivity) onto the anatomy in the image, guaranteeing that resulting segmentations possess the correct topological structure.

### Mathematical Formulation

Let $I: \Omega \to \mathbb{R}$ be the image ($d$-dimensional), and $M_{temp}: \Omega \to \{0,1\}$ a template mask with correct topology. A UNet (2D or 3D) is adapted to take both $I$ and $M_{temp}$ as input and outputs a deformation field $f: \Omega \to \mathbb{R}^d$ describing a spatial map $\phi$.

- **Deformation**: $\phi(x) = x + u(x)$ or directly as $\phi(x) = f(x)$
- **Prediction**: $M_{pred}(x) = M_{temp}(\phi(x))$

A spatial transformer layer executes mask warping using bilinear or trilinear sampling.

### Loss and Topology Regularization

To guarantee topology preservation, the Jacobian determinant of $\phi$ is constrained to remain positive via an $\varepsilon$-ReLU penalty:

\[
L_{Jac}(f) = \|\mathrm{ReLU}_\varepsilon ( -\det \nabla \phi(x) ) \|_1 = \int_\Omega | \mathrm{ReLU}_\varepsilon( - \det \nabla \phi(x) ) | dx
\]

The total loss combines Dice loss, Jacobian penalty, and optional Laplacian smoothness:

\[
L = \lambda_{Dice} L_{Dice}(M_{pred}, M_{label}) + \lambda_{Jac} L_{Jac}(f) + \lambda_{Lap} L_{Lap}(f)
\]
with reported weights $\lambda_{Dice}=1.0$, $\lambda_{Jac}=1.0$, $\lambda_{Lap}=0.1$ [2202.13331].

### Multi-Scale Cascade

A three-level coarse-to-fine cascade (ml-TPSN) improves boundary detail. The machine learning pipeline trains and infers at low, medium, and full resolution, with each TPSN module starting from the previous upsampled mask.

## 5. Experimental Evaluation and Comparative Analysis

### Datasets and Protocol

- **Ham10000 (2D dermoscopy):** 10,015 images, split 9,000/1,015, resized to $128\times128$
- **KiTS21 (3D CT kidneys):** 300 volumes, 210/70 split, resampled to $64\times128\times128$

### Results

| Dataset          | UNet         | UNet+CCA     | TEDSNet      | TPSN         | ml-TPSN       |
|------------------|-------------|-------------|--------------|--------------|---------------|
| Ham10000 (2D)    | 93.50±0.29% | 93.63±0.29% | 89.91±0.27%  | 93.77±0.37%  | 94.42±0.39%   |
| KiTS21 (3D)      | 92.76±0.38% | 92.81±0.38% | 87.67±0.29%  | 92.73±0.48%  | 93.17±0.41%   |

TPSN consistently outperforms baseline UNet and UNet+CCA (connected-component analysis), especially in cases with partial data loss or low contrast, yielding a segmentation that is, by construction, topologically correct and free of spurious disconnected components.

### Inference and Post-processing

Given a test image and template mask, TPSN produces the final mask in a single forward pass, and no further post-processing is required to ensure topology correctness.

## 6. Distinctions and Nomenclature

The term "TP-UNet" is ambiguous, referring to both Temporal Prompt Guided UNet [2411.11305] and Topology-Preserving Segmentation Network [2202.13331] in contemporary literature. Each framework leverages different forms of prior:

- **Temporal Prompt Guidance**: Uses slice-wise textual prompts encoding anatomical order.
- **Topology Preservation**: Incorporates topological priors via template-based diffeomorphic deformation and Jacobian regularization.

A plausible implication is that future works may combine both temporal and topological priors within unified segmentation models if anatomical constraints and scan protocols permit.

## 7. Prospective Extensions

For Temporal Prompt Guided UNet, proposed future directions include employing 3D prompt sequences, enabling uncertainty-aware prompt construction, and integrating clinical report information to improve volumetric spatial–temporal consistency [2411.11305].

For TPSN, plausible extensions involve templates with more complex topology, data-driven template adaptation, or fusion with semantic priors beyond topology, aiming to generalize topology preservation across multi-object and variable-connectivity segmentation tasks [2202.13331].

Source: https://www.emergentmind.com/topics/tp-unet