💻 computer science

SAMamba3D: adapting Segment Anything for generalizable 3D segmentation of multiphase pore-scale images

The paper introduces SAMamba3D, a parameter-efficient framework that adapts the frozen Segment Anything Model with Mamba-based volumetric context modeling to achieve generalizable, high-performance 3D segmentation of multiphase pore-scale rock images across diverse geological and scanning conditions without requiring extensive retraining.

Original authors: Rui Zhang, Xianzhi Song, Linqi Zhu, Branko Bijeljic, Gensheng Li, Martin J. Blunt

Published 2026-05-05

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Rui Zhang, Xianzhi Song, Linqi Zhu, Branko Bijeljic, Gensheng Li, Martin J. Blunt

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, incredibly detailed 3D puzzle made of rock. Inside this rock, there are tiny tunnels (pores) filled with different fluids like oil, water, or gas. Scientists use special X-ray cameras to take pictures of these rocks to understand how fluids move underground, which is crucial for things like storing carbon dioxide or finding oil.

The problem is that these X-ray pictures are just shades of gray. To do any science with them, a computer needs to "color-code" the image: turning the rock green, the water blue, and the oil red. This process is called segmentation.

The Old Way: The "Custom Tailor" Problem

Until now, creating these colored maps was like hiring a custom tailor for every single outfit. If you changed the type of rock, the type of fluid, or even the camera used to take the picture, the old computer program would get confused. It would need to be completely retrained from scratch for every new situation. It was slow, expensive, and often made mistakes in the tiny, tricky spots where the fluids touch the rock.

The New Solution: SAMamba3D

The authors of this paper created a new tool called SAMamba3D. Think of it as a "universal translator" for these rock images.

Here is how it works, using a simple analogy:

The Expert Eye (SAM): The system starts with a pre-trained "brain" called SAM (Segment Anything Model). Imagine SAM as a world-class artist who has seen millions of 2D drawings and knows exactly how to draw a perfect line around an object. However, SAM only knows how to look at flat, 2D pictures.
The 3D Context (Mamba): The rock images are 3D, and the fluids wrap around each other in complex ways. To help SAM understand the 3D shape, the researchers added a second brain called Mamba. Think of Mamba as a structural engineer who understands how buildings (or in this case, rock pores) hold together in three dimensions.
The Teamwork: Instead of letting the artist (SAM) and the engineer (Mamba) work separately, SAMamba3D makes them talk to each other constantly.
- SAM says: "I see a sharp edge here!"
- Mamba says: "I see that this edge connects to a tunnel over there, so it must be part of the water layer."
- Together, they decide exactly where the water ends and the oil begins, even in the tiniest, most confusing spots.

Why It's a Big Deal

The paper claims that this new team-up is a game-changer for three main reasons:

It Doesn't Need Re-training: Usually, if you show a computer a new type of rock, you have to teach it all over again. SAMamba3D is like a smart assistant that has learned the principles of rock and fluid shapes. You can show it a completely new rock type, a new fluid (like hydrogen instead of oil), or a new camera setting, and it just works without needing a new lesson.
It's Fast and Light: Because it uses a "frozen" (pre-trained) artist and only adds a small, efficient helper (Mamba), it is much faster and requires less computer power than the old, heavy methods. It's like upgrading from a massive, fuel-guzzling truck to a sleek, high-speed electric car that gets the same job done.
It Gets the Physics Right: The most important claim is that the results aren't just "pretty pictures." The way the computer colors the fluids matches real-world physics. For example, in a water-loving rock, the water naturally clings to the rock grains. SAMamba3D correctly identifies these thin, clinging layers of water, whereas older methods often missed them or broke them apart. This means scientists can trust the numbers they calculate (like how much oil is actually trapped) without having to manually fix the computer's mistakes.

The Bottom Line

The paper demonstrates that by combining a powerful, pre-trained 2D image expert with a smart 3D context builder, they created a system that can look at a wide variety of complex underground rock images and accurately separate the rock from the fluids. It does this without needing to be retrained for every single new experiment, saving time and providing more reliable data for understanding how fluids move underground.

Technical Summary: SAMamba3D

Problem Statement

Reliable segmentation of multiphase pore-scale X-ray micro-computed tomography (micro-CT) images is essential for quantifying fluid saturation, connectivity, and interfacial geometry in porous media. However, current 3D segmentation methods are typically dataset-specific, requiring retraining or extensive fine-tuning whenever rock type, fluid configuration, scanner hardware, or acquisition conditions change. While foundational models like the Segment Anything Model (SAM) offer strong 2D boundary priors, they are not directly applicable to 3D pore-scale data due to the monochromatic nature of the images, low contrast, reconstruction artifacts, and the critical need for volumetric continuity. Existing approaches often fail to preserve connected pore structures, thin wetting layers, and interfacial topology when applied across diverse domains.

Methodology

Dataset and Preprocessing

The study utilizes a curated multi-source dataset comprising sandstone, carbonate, and glass bead pack samples with varying fluid systems (oil–brine, scCO2–brine, H2–brine), wettability states, and scanner resolutions. To address domain shifts without information leakage, the authors employ a preprocessing pipeline involving:

Denoising: Non-local means (NLM) filtering to suppress noise while preserving boundaries.
Percentile-based Intensity Alignment: Aligning source volumes to a fixed reference range derived solely from training data.
Global Standardization: Standardizing volumes using global mean and standard deviation computed from the training split.

SAMamba3D Architecture

SAMamba3D is a parameter-efficient framework designed to adapt a largely frozen SAM encoder to 3D volumetric data by coupling it with a Mamba-based branch for global structural modeling. The architecture consists of five modules:

3D-Adapted SAM Encoder: Extends the 2D SAM backbone to 3D using fine-grained $4 \times 4 \times 4$ patch embeddings (to capture thin wetting layers), Low-Rank Adaptation (LoRA) for parameter-efficient tuning, and anisotropy-aware positional encoding.
Hierarchical Mamba Branch: A state-space model branch that extracts multi-scale volumetric context ( $s_0$ to $s_3$ ) and generates global descriptors and spatial importance maps.
Adaptive Cross-Scale Fusion: Enables bidirectional interaction between the SAM and Mamba branches. Early layers use one-way injection (Mamba to SAM), while deeper layers allow bidirectional exchange. A gating mechanism selectively processes high-importance regions (interfaces) with full attention while using lightweight pathways for low-information regions.
Mamba Controller: Transforms global Mamba features into modulation parameters (FiLM) to condition the decoder.
FiLM-Modulated Multi-Scale Decoder: Progressively fuses features from the SAM branch, Mamba skip connections, and global conditioning signals to produce the final segmentation.

Loss Function and Training

To address class imbalance and uncertainty near phase boundaries, the authors propose a boundary-aware composite loss comprising:

Weighted Dice Loss: Handles class imbalance.
Tversky Loss: Penalizes missed small structures more heavily than spurious voxels.
Confidence-Weighted Focal Loss: Focuses optimization on difficult voxels while down-weighting uncertain boundary labels based on distance to the nearest manual annotation boundary.

Training follows a two-stage progressive unfreezing strategy:

Stage A (Warm-up): Only 3D-specific components (Mamba branch, adapters, decoder) are trained; the SAM backbone remains frozen.
Stage B (Adaptation): SAM layers (LayerNorm, LoRA, adapters) are unfrozen to allow gradual adaptation to the volumetric domain.

Key Contributions

Parameter-Efficient Adaptation: Introduction of a method to adapt the frozen SAM encoder to 3D multiphase pore-scale micro-CT segmentation using LoRA and lightweight 3D modules, avoiding full retraining.
Volumetric Context Fusion: Development of a hybrid architecture that combines local boundary-aware representations (SAM) with broader structural reasoning (Mamba) through bidirectional cross-scale interaction.
Generalization and Physical Fidelity: Evaluation of the method not only on segmentation accuracy (Dice/IoU) but also on the preservation of downstream physical descriptors (porosity, saturation, Euler number, contact angle, curvature) across unseen datasets.

Results

Accuracy vs. Efficiency: SAMamba3D achieved a mean Dice score of 0.94 on unseen test datasets (Bentheimer sandstone and Estaillades carbonate), outperforming baselines including U-Net, nnU-Net, UNETR, SwinUNETR, and U-Mamba. Notably, it reduced computational cost by nearly 40x compared to nnU-Net (16 GFLOPs vs. 622 GFLOPs) while improving accuracy by 15%.
Physical Descriptor Preservation: On water-wet Bentheimer sandstone, the model corrected topological errors in the baseline segmentation, yielding a brine Euler number of –1384 (indicating a connected network) compared to +1931 for the baseline (indicating fragmentation). This aligns with the expected physics of water-wet systems.
Cross-Domain Generalization: Without retraining or fine-tuning, the model successfully segmented diverse unseen datasets, including:
- Different Rock Types: From homogeneous sandstones to heterogeneous carbonates with micro-porosity.
- Different Wettability: Water-wet, mixed-wet, and oil-wet conditions, correctly capturing fluid morphologies (e.g., brine in corners for water-wet, oil in corners for oil-wet).
- Different Fluids: Successfully generalized to scCO2–brine and H2–brine systems, which have distinct grayscale statistics and low contrast compared to the training oil–brine data.
Interface Refinement: The model demonstrated superior ability to resolve thin wetting layers and complex interfacial geometries, resulting in curvature and contact angle distributions more consistent with mixed-wet physical expectations than baseline methods.

Significance and Claims

The paper claims that multiphase pore-scale micro-CT segmentation can be treated as a reusable foundational model adaptation problem rather than a dataset-specific training task. The authors assert that SAMamba3D provides a viable route toward scalable and physically reliable segmentation that:

Reduces Dependency on Retraining: Eliminates the need for case-specific retraining when rock type, fluid, or scanner conditions change.
Preserves Physical Meaning: Generates segmentations that maintain physically meaningful descriptors (connectivity, interface morphology) essential for pore-network extraction and flow simulation.
Offers Computational Efficiency: Delivers state-of-the-art accuracy with significantly lower inference costs compared to fully retrained 3D baselines.

The authors acknowledge limitations, noting that performance may degrade in underrepresented lithologies (e.g., organic-rich shales) and that the current framework does not encode explicit physical constraints (e.g., Young–Laplace consistency) or handle time-resolved 4D imaging. However, they conclude that this approach substantially reduces the dependence of pore-scale analysis on manual correction and repeated model redevelopment.