AI Built for Cells Now Measures River Gravels With Record Precision

Measuring the size and shape of sediment grains sounds like one of the most straightforward tasks in geoscience, yet it has stubbornly resisted automation for decades. Field crews still kneel on gravel bars counting pebbles with templates, and laboratories still pour sediment through nested sieves, because the digital alternatives have repeatedly fallen short. Now a team led by David Mair of the University of Bern reports a major step forward. In a study published in Earth Surface Dynamics, the researchers introduce ImageGrains 2.0, an open-source framework that borrows a cutting-edge artificial intelligence model from biomedical research and retrains it to detect, outline and measure individual grains in photographs and X-ray scans of sediment. The result, they show, outperforms every comparable method they tested, and it does so across a strikingly wide range of image types.

The stakes are higher than they might appear. Grain morphometry underpins a remarkable breadth of Earth science. The sizes and shapes of clasts record how sediments are produced on hillslopes, how rivers hydraulically sort and transport their loads, and how deposits in the stratigraphic archive were laid down. Grain size data have even been used to reconstruct transport conditions on Mars, where the shapes of ancient pebbles hint at how far water once flowed. But the traditional ways of collecting these data are laborious and biased. Manual Wolman counts and line sampling yield relatively few observations and depend heavily on the observer. Point-cloud methods struggle with grain occlusion and require geometric assumptions. Earlier image-based tools either need field calibration or systematic manual correction, and they tend to over- or underestimate grain sizes in predictable but troublesome ways.

Machine learning promised to change this, and to some extent it has. Texture-based neural networks can predict summary statistics of a grain size distribution directly from image texture, while segmentation-based approaches delineate every grain individually, opening the door to datasets of tens of thousands of measurements. The catch, as Mair and colleagues emphasize, is that these segmentation models have behaved like narrow specialists. Trained on homogeneous, task-specific datasets, they lose accuracy the moment they encounter a new camera, a new lighting regime or a new sediment type. Each new setting has effectively demanded a new model and a freshly curated training set, which has blunted much of the promised efficiency.

The Bern team’s solution was to reach for a different class of model altogether. In recent years, transformer-based foundation models such as Meta’s Segment Anything Model, or SAM, trained on eleven million images and more than a billion annotation masks, have shown an extraordinary ability to generalize to unseen data. But SAM in its basic form is inefficient at densely segmenting large numbers of similar objects, which is precisely what a gravel bar photograph demands. The breakthrough came with Cellpose-SAM, a next-generation model developed for segmenting cells in biomedical imagery. It grafts SAM’s transformer image encoder onto the flow-based segmentation machinery of the original Cellpose framework, which tracks predicted vector flows to separate touching objects. The hybrid keeps SAM’s generalization while excelling at dense, prompt-free segmentation of many instances of one object class, and its architecture has grown to more than 304 million trainable parameters, roughly fifty times larger than the U-Net backbone it replaces.

Retraining such a model for geoscience required data, and lots of carefully annotated ones. The researchers assembled the ImageGrains 2.0 dataset, combining 81 previously labelled tiles with 162 new ones for a total of 243 image tiles containing more than 29,000 manually annotated grain masks. The imagery spans fluvial gravels photographed from uncrewed aerial vehicles and handheld cameras in Switzerland, Spain, Peru, New Zealand and Namibia; painted and unpainted pebbles used to track grain mobility; angular proglacial deposits; near-vertical outcrops of lithified conglomerates; X-ray computed tomography slices of glacial till; and micro-CT images of bioclastic marine sand. Crucially, the team deliberately included awkward content, such as vegetation, sieves, hands, scale objects, water bodies and harsh shadows, to test whether the model could distinguish genuine grains from everything else in a messy field photograph.

The performance gains are substantial. Benchmarked against the manual annotations using average precision at an intersection-over-union threshold of 0.5, the fine-tuned Cellpose-SAM model achieved a median score of 0.72 across all image tiles, roughly 20 percent higher than the second-best method, a Cellpose 2 model trained on the same data. The advantage held in both training and unseen test splits, and it proved largely independent of how many images each sediment type contributed to training. The model scored highest of all on X-ray CT tiles, which made up only about a tenth of the dataset. Even more telling was the generalization test: when the researchers withheld two entire subsets from training, the model still matched or beat specialist models that had been fine-tuned on those very images, evidence that the transfer-learning strategy preserves the broad visual competence of its SAM heritage.

Accurate masks translate directly into accurate measurements, and here the numbers are striking. The model recovered 93 percent of ground-truth grains, and mean differences in grain size were around 3 percent or less for both the long and short axes of ellipse fits to each grain. The resulting grain size distributions were statistically indistinguishable from the ground truth in 88 percent of cases for the a-axis and 82 percent for the b-axis, compared with only 54 and 57 percent for the second-best model. Shape metrics fared equally well, with average deviations below 2 percent for roundness, roughness, elongation and orientation. The team also identified a practical rule of thumb: perfect statistical agreement between predicted and true grain size distributions was only achieved when segmentation scores exceeded roughly 0.68 average precision, a threshold that varied with how well sorted the sediment was.

The framework extends into three dimensions as well. Applying the model to a stack of 400 X-ray CT slices from a drill core of glacio-fluvial sediment in southern Germany, the researchers segmented 4,647 well-defined coarse grains in full volumetric detail, something 2D photographs can never provide. Because CT scans capture complete grains rather than partially occluded surfaces, they sidestep the geometric extrapolations that plague LiDAR and photogrammetry-based approaches. The authors note that rigorous 3D benchmarking still awaits a manually annotated volumetric ground truth, whose creation has so far been impractical, but the demonstration suggests CT-based sediment analysis, from clast fabric to depositional history, is now within reach of automated workflows.

Accessibility was a design priority. ImageGrains 2.0 ships as an installable Python package with an interactive graphical user interface for annotating images and correcting model predictions, so researchers need not write code at all. Training runs in under ninety minutes on a high-end GPU, and although transformer models are demanding, the authors show that fine-tuning remains feasible on a mid-range desktop graphics card, with inference requiring as little as 3 gigabytes of memory. Because the model builds on already-trained foundation weights, its marginal energy footprint is modest compared with training such architectures from scratch. Adapting the model to a genuinely new sediment type or imaging style can require as few as seven additional annotated tiles, a dramatic lowering of the barrier that once made every new site a modelling project of its own.

The deeper message of the study concerns where progress will come from next. The authors argue that future gains in grain measurement will depend less on new architectures than on larger, more variable, publicly shared annotated datasets, ideally with annotations from multiple experts to establish consensus-level benchmarks. They explicitly frame their work as a step toward a foundation model for granular particles in geoscientific imagery, one that could serve any discipline needing dense segmentation and morphometric analysis, from sedimentology to planetary science. With the dataset, model weights and code all openly available, the Bern team has effectively invited the community to build on their groundwork, turning what was once a tedious field chore into a problem that a well-fed neural network can now handle in minutes.

Subject of Research: Deep-learning-based automated segmentation and morphometric measurement of sediment grains in geoscientific imagery

Article Title: ImageGrains 2.0: Improved precision and generalization for grain segmentation

Article References: Mair, D., Witz, G., Do Prado, A., Garefalakis, P., Wild, A., Ville, F., Schuster, B., Horn, M., Österle, J., Fabbri, S. C., Litty, C., Achleitner, S., Leistner, S., Hiller, C., & Schlunegger, F. (2026). ImageGrains 2.0: Improved precision and generalization for grain segmentation. Earth Surface Dynamics, 14(4), 527-551. https://doi.org/10.5194/esurf-14-527-2026

Image Credits: AI Generated

DOI: 10.5194/esurf-14-527-2026

Keywords: ImageGrains 2.0, grain segmentation, Cellpose-SAM, deep learning, sedimentology, grain size distribution, X-ray computed tomography, transfer learning, Segment Anything Model, fluvial sediment, geomorphology, open-source software