Getting started¶

The best way to use the segmenteverygrain package is to run the Segment_every_grain.ipynb notebook.

The notebook goes through the steps of loading the models, running the segmentation, interactively updating the result, and saving the grain data and the mask. The text below summarizes the steps that you need to take to run the segmentation.

Loading the models¶

To load the U-Net model, you can use the ‘load_model’ function from Keras. The U-Net model is saved in the ‘seg_model_smooth_labels.keras’ file.

import segmenteverygrain as seg
from keras.saving import load_model
model = load_model("seg_model_smooth_labels.keras", custom_objects={'weighted_crossentropy': seg.weighted_crossentropy})

Note

As of v0.4.0, the U-Net model outputs raw logits (pre-softmax). Softmax is applied automatically at inference time in predict_image_tile. Models trained with v0.3.0 and earlier (e.g. seg_model.keras) are incompatible with this version — use seg_model_smooth_labels.keras or retrain your custom models.

This assumes that you are using Keras 3 and ‘seg_model_smooth_labels.keras’ was saved using Keras 3. Older models created with a segmenteverygrain version that was based on Keras 2 do not work with the latest version of the package.

The SAM 2.1 model can be downloaded from this link. You can also download it programmatically:

import os
import urllib.request

if not os.path.exists("./models/sam2.1_hiera_large.pt"):
    url = "https://huggingface.co/facebook/sam2.1-hiera-large/resolve/main/sam2.1_hiera_large.pt"
    urllib.request.urlretrieve(url, "./models/sam2.1_hiera_large.pt")

Running the segmentation¶

The recommended way of running the segmentation is the predict_large_image function, which is what the Segment_every_grain.ipynb notebook uses by default. Despite its name, it works well on images of any size: it splits the image into overlapping patches (a single patch for small images), runs the U-Net prediction and the SAM segmentation on each patch, and blends the results. It is described in the Large image processing section below.

The functions in this section and the next two expose the individual steps of the pipeline — running them one by one is useful for understanding the workflow and for QC-ing the intermediate U-Net output, especially when working with a new image type.

To run the U-Net segmentation on an image and label the grains in the U-Net output:

import numpy as np
import matplotlib.pyplot as plt
from keras.utils import load_img

# Load your image
fname = "path/to/your/image.jpg"
image = np.array(load_img(fname))

# Run U-Net prediction
image_pred = seg.predict_image(image, model, I=256)
labels, coords = seg.label_grains(image, image_pred, dbs_max_dist=20.0)

The input image should not be much larger than ~2000x3000 pixels, in part to avoid long running times; it is supposed to be a numpy array with 3 channels (RGB). Grains should be well defined in the image and not too small (e.g., only a few pixels in size).

Quality control of U-Net prediction¶

The U-Net prediction should be QC-d before running the SAM segmentation:

plt.figure(figsize=(15,10))
plt.imshow(image_pred)
plt.scatter(np.array(coords)[:,0], np.array(coords)[:,1], c='k')
plt.xticks([])
plt.yticks([])

The black dots in the figure represent the SAM prompts that will be used for grain segmentation. If the U-Net segmentation is of low quality, the base model can be (and should be) finetuned using the steps outlined below.

SAM segmentation¶

Here is an example showing how to run the SAM segmentation on an image, using the outputs from the U-Net model:

import torch
from sam2.build_sam import build_sam2
# Auto-detect device: CUDA for NVIDIA, MPS for Apple Silicon, CPU as fallback
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
sam = build_sam2("configs/sam2.1/sam2.1_hiera_l.yaml", "sam2.1_hiera_large.pt", device=device) # load the SAM 2.1 model
all_grains, labels, mask_all, grain_data, fig, ax = seg.sam_segmentation(sam, image, image_pred, coords, labels, min_area=400.0, plot_image=True, remove_edge_grains=False, remove_large_objects=False)

The all_grains list contains shapely polygons of the grains detected in the image. labels is an image that contains the labels of the grains. grain_data is a pandas dataframe with a number of grain parameters.

Interactive editing of results¶

After the initial segmentation, you can interactively edit the results using the GrainPlot class from the interactions module. This provides a modern, interactive interface for deleting, merging, and adding grains.

First, convert the polygons to Grain objects and set up the SAM predictor:

import segmenteverygrain.interactions as si
from sam2.sam2_image_predictor import SAM2ImagePredictor
from tqdm import tqdm

# Convert polygons to Grain objects
grains = si.polygons_to_grains(all_grains, image=image)
for g in tqdm(grains, desc='Measuring detected grains'):
    g.measure()

# Set up SAM 2.1 predictor for adding new grains
predictor = SAM2ImagePredictor(sam)
predictor.set_image(image)

Then create the interactive plot:

plot = si.GrainPlot(
    grains,
    image=image,
    predictor=predictor,
    blit=True,                      # Use blitting for faster rendering
    color_palette='tab20b',         # Matplotlib colormap for grain colors
    figsize=(12, 8),                # Figure size in inches
    scale_m=500*1e-6,               # Length of scale bar in meters (for unit conversion)
    px_per_m=1.0,                   # Pixels per meter (will be updated if scale bar is drawn)
)
plot.activate()

Interactive controls (also shown in the figure title bar):

Mouse controls:

Left-click on existing grain: Select/unselect the grain
Left-click in grain-free area: Place foreground prompt for instant grain creation (auto-create)
Alt + Left-click: Place foreground prompt for multi-prompt grain creation (hold Alt, click multiple times, release Alt to create)
Alt + Right-click: Place background prompt for multi-prompt creation
Shift + Left-drag: Draw a scale bar line (red line) for unit conversion
Middle-click or Shift + Left-click on grain: Show grain measurement info

Keyboard controls:

d or Delete: Delete selected (highlighted) grains
m: Merge selected grains (must be touching)
z: Undo (delete the most recently created grain)
Ctrl (hold): Temporarily hide all grain masks
Esc: Remove all prompts and unselect all grains
c: Create grain from existing prompts (alternative to auto-create)

After editing, retrieve the updated grains and deactivate the interactive features:

grains = plot.get_grains()  # Get the edited list of grains
plot.deactivate()           # Turn off interactive features

# Optionally draw the major and minor axes on each grain
plot.draw_axes()

Scale bar and unit conversion¶

To convert measurements from pixels to real-world units, you can either:

Specify the scale when creating the plot using the px_per_m parameter (pixels per meter).
Draw a scale bar interactively: Hold Shift and drag the mouse to draw a line on a known reference object. The scale_m parameter specifies the real-world length of this reference object in meters.

After drawing a scale bar, the pixel-to-meter conversion is automatically updated:

# Retrieve the scale after drawing a scale bar
px_per_m = plot.px_per_m

Grain size analysis¶

Generate a summary dataframe and histogram of grain measurements:

# Get summary dataframe with all grain measurements
summary = si.get_summary(grains, px_per_m=plot.px_per_m)

# Create histogram of major and minor axis lengths
hist = seg.plot_histogram_of_axis_lengths(
    summary['major_axis_length'] * 1000,  # Convert to mm
    summary['minor_axis_length'] * 1000,
    binsize=0.25
)

Area-weighted grain size distributions¶

Counting grains in an image produces a number-weighted distribution: every grain contributes equally, no matter how small. Field and laboratory methods such as sieving, point counting, and Wolman pebble counts instead sample grains roughly in proportion to their size, so their distributions are coarser than a per-grain count of the same sediment. To make image-based results comparable with these methods, the grain size distribution can be weighted by grain area.

As of v0.5.0, area weighting uses explicit weights: each grain’s area is used directly as its weight when building histograms, cumulative distributions, and percentiles. Earlier versions approximated this by replicating each grain size in proportion to its area; because the replication count was truncated to an integer, grains smaller than half the mean area were dropped entirely, which biased the statistics coarse. The old get_area_weighted_distribution function is therefore deprecated (it still works, but emits a DeprecationWarning), and area-weighted results computed with v0.5.0 or later will differ somewhat from earlier ones — the fine tail of the distribution is now fully represented.

To plot an area-weighted histogram, pass the grain areas to plot_histogram_of_axis_lengths:

# Area-weighted histogram (axis lengths in mm, areas in mm^2)
fig, ax = seg.plot_histogram_of_axis_lengths(
    summary['major_axis_length'] * 1000,  # m to mm
    summary['minor_axis_length'] * 1000,
    area=summary['area'] * 1e6,           # m^2 to mm^2
    binsize=0.25,
)

When area is provided, the y axis is labeled ‘area-weighted count’. The weights are normalized so that they sum to the number of grains; this keeps the bar heights on a count-like scale and makes them independent of the units the areas are given in. Omit area for an ordinary count-based histogram.

The get_histogram and save_histogram functions of the interactions module are area-weighted by default (as long as the summary dataframe has an ‘area’ column); pass area_weighted=False for a count-based histogram.

For statistics computed directly from the data, without binning, two functions accept an optional weights argument:

weighted_ecdf(values, weights=None) returns the sorted values together with the cumulative and exceedance fractions of the total weight, i.e. a weighted empirical cumulative distribution function.
weighted_percentile(values, percentiles, weights=None) computes weighted percentiles, such as the D50 or D84 of an area-weighted distribution. It uses the midpoint convention, so with equal weights the results are very close to numpy.percentile.

major_mm = summary['major_axis_length'] * 1000  # m to mm
areas = summary['area'] * 1e6                   # m^2 to mm^2

# Area-weighted cumulative distribution
sizes, cdf, exceedance = seg.weighted_ecdf(major_mm, weights=areas)
plt.figure()
plt.plot(sizes, cdf)
plt.xscale('log')
plt.xlabel('major axis length (mm)')
plt.ylabel('cumulative fraction (area-weighted)')

# Area-weighted percentiles
d16, d50, d84 = seg.weighted_percentile(major_mm, [16, 50, 84], weights=areas)

With weights=None, both functions fall back to ordinary count-based statistics, so the same code can be used for weighted and unweighted analyses.

Saving results¶

The interactions module provides convenient functions for saving all results:

out_fn = "./examples/output/my_image"  # Base filename (without extension)

# Save grain shapes as GeoJSON
si.save_grains(out_fn + '_grains.geojson', grains)

# Save the plot with grain overlays
plot.savefig(out_fn + '_grains.jpg')

# Save grain measurements as CSV
summary = si.save_summary(out_fn + '_summary.csv', grains, px_per_m=plot.px_per_m)

# Save histogram as image (area-weighted by default;
# pass area_weighted=False for a count-based histogram)
si.save_histogram(out_fn + '_summary.jpg', summary=summary)

# Save binary mask for training (0-1 values)
si.save_mask(out_fn + '_mask.png', grains, image, scale=False)

# Save human-readable mask (0-255 values)
si.save_mask(out_fn + '_mask2.jpg', grains, image, scale=True)

Large image processing¶

The predict_large_image function is the recommended default for running the segmentation, whatever the image size: it splits the image into overlapping patches (small images are processed as a single patch), runs the U-Net prediction and the SAM segmentation on each patch, and blends the results:

from PIL import Image
import segmenteverygrain.interactions as si

Image.MAX_IMAGE_PIXELS = None  # needed for very large images

fname = "./examples/my_large_image.jpg"
image = si.load_image(fname)

all_grains, image_pred, all_coords = seg.predict_large_image(
    fname, unet, sam,
    use_sam=True,          # use SAM for grain boundary detection
    min_area=400.0,
    patch_size=2000,
    overlap=600,           # the larger the grains, the larger this overlap needs to be
                           # (~ the size of the largest grains)
    dbs_max_dist=100.0,    # reduce this if there are not enough prompts and
                           # some grains are not detected
    remove_edge_grains=False  # keep grains on outer edges of the full image
)

You can also run the segmentation using the U-Net only, without SAM refinement. This is significantly faster and avoids the need to load the SAM model:

all_grains, image_pred, all_coords = seg.predict_large_image(
    fname, unet,
    patch_size=2000,
    overlap=200,
    use_sam=False,
    dilation=3,        # expand grain labels to recover boundary pixels
    min_area=400.0,
)

When use_sam=False, grain labeling is performed once on the full blended U-Net prediction (rather than per patch), so grains that straddle patch boundaries are handled correctly. The dilation parameter expands each grain label by the specified number of pixels, compensating for the boundary channel that the U-Net predicts between grains.

Interrupting and resuming long runs¶

Segmenting a very large image can take hours, and before v0.5.0 a crash, an interrupted kernel, or a machine going to sleep meant starting over. As of v0.5.0, passing a checkpoint_dir to predict_large_image makes the run resumable:

all_grains, image_pred, all_coords = seg.predict_large_image(
    fname, unet, sam,
    min_area=400.0,
    patch_size=2000,
    overlap=600,
    remove_edge_grains=False,
    checkpoint_dir="./my_large_image_checkpoints",
    verbose=False,
)

With checkpoint_dir set, the grains of each patch are written to GeoJSON (in global image coordinates) as soon as the patch is processed, the blended U-Net prediction is stored on disk as a numpy.memmap instead of a large in-memory array (which also lowers the RAM footprint), and a manifest keeps track of the completed patches. If the run is interrupted for any reason, re-running the same call with the same checkpoint_dir — and the same parameters — skips the finished patches and continues where it stopped.

A few things to keep in mind:

The returned image_pred is a numpy.memmap backed by a file in the checkpoint directory; copy it with np.array(image_pred) if you need it to outlive the directory.
Once the run has finished and the results are saved, the checkpoint directory can be deleted.
For images of about 100 megapixels or more, predict_large_image emits a warning if checkpoint_dir is not set, as a reminder that an interruption would otherwise lose all progress.

The verbose=False setting used above is also new in v0.5.0: it replaces the per-patch progress bars with a single concise line per patch. This is recommended for large images, where the volume of progress-bar output can make notebook frontends sluggish or appear to hang.

Just like before, the all_grains list contains shapely polygons of the grains detected in the image. The image containing the grain labels can be generated like this:

labels = seg.rasterize_grains(all_grains, image)

To interactively edit the results, convert to Grain objects and use GrainPlot:

grains = si.polygons_to_grains(all_grains, image=image)
predictor.set_image(image)

plot = si.GrainPlot(
    grains,
    image=image,
    predictor=predictor,
    color_palette='tab20b',
)
plot.activate()

See the Segment_every_grain.ipynb notebook for a complete example of how the models can be loaded and used for segmenting an image and QC-ing the result. The notebook goes through all the steps described above in an interactive format.

Grain extraction and clustering¶

The grain_utils module provides functions for extracting individual grain images and clustering them for classification tasks.

Extracting individual grains¶

Extract standardized, square images of individual grains suitable for machine learning:

from segmenteverygrain import extract_all_grains, extract_grain_image

# Extract all grains as standardized 224x224 images
grain_images, grain_masks, grain_preds = extract_all_grains(
    all_grains, image, image_pred, target_size=224
)

# Or extract a single grain
grain_img, grain_mask, grain_pred, orig_size = extract_grain_image(
    all_grains[0], image, image_pred, target_size=224, pad=10
)

Feature extraction¶

Extract deep learning features using pre-trained CNNs for clustering or classification:

from segmenteverygrain import extract_vgg16_features, extract_color_features

# Extract VGG16 features (4096-dimensional)
features, model = extract_vgg16_features(grain_images, model_name='VGG16')

# Or extract color-based features
color_features = extract_color_features(grain_images, color_space='hsv')

Clustering grains¶

Cluster grains based on their features:

from segmenteverygrain import cluster_grains, create_clustered_grain_montage

# Cluster using K-means with PCA dimensionality reduction
labels, reduced_features, pca, clusterer = cluster_grains(
    features, n_clusters=10, n_components=25
)

# Create a visual montage of clustered grains
montage, cluster_info = create_clustered_grain_montage(
    labels, grain_images, grid_cols=20, draw_boundaries=True
)

Interactive grain selection and labeling¶

Use the ClusterMontageSelector for quality control (removing bad grains or clusters):

from segmenteverygrain import ClusterMontageSelector

selector = ClusterMontageSelector(labels, grain_images, all_grains)
selector.activate()  # Interactive mode

# After selection, get filtered results
filtered_grains = selector.get_filtered_grains()

Use the ClusterMontageLabeler for labeling grains with custom categories:

from segmenteverygrain import ClusterMontageLabeler

labeler = ClusterMontageLabeler(
    labels, grain_images, all_grains,
    label_names=['quartz', 'feldspar', 'lithic', 'other']
)
labeler.activate()  # Interactive mode

# Export labels
labeler.export_labels('grain_labels.csv')
labeler.save_labeled_images('labeled_grains/')

Visualizing classified grains¶

Plot grains colored by their classification:

from segmenteverygrain import plot_classified_grains

fig, ax = plot_classified_grains(
    image, all_grains, classifications,
    class_colors={'quartz': 'blue', 'feldspar': 'red', 'other': 'green'}
)

Hardware requirements¶

For training a new U-Net model or fine tuning the existing one, GPU access is necessary. The easiest way of getting access to a powerful GPU is Google Colab. In inference mode, a moderately powerful computer with at least 16 GB of memory should be enough. That said, larger CPU speeds and more memory will significantly reduce inference time.

Finetuning the U-Net model¶

The last section of the Segment_every_grain.ipynb notebook shows how to finetune the U-Net model. The first step is to create patches (usually 256x256 pixels in size) from the images and the corresponding masks that you want to use for training.

image_dir, mask_dir = seg.patchify_training_data(input_dir, patch_dir)

The input_dir should contain the images and masks that you want to use for training. These files should have ‘image’ and ‘mask’ in their filenames, for example, ‘sample1_image.png’ and ‘sample1_mask.png’. An example image and the corresponding mask are available in the repository.

The mask is an 8-bit image and should contain only three numbers: 0, 1, and 2. 0 is the background, 1 is the grain, and 2 is the grain boundary. Usually the mask is generated using the segmenteverygrain workflow, that is, by running the U-Net segmentation first, the SAM segmentation second, and then cleaning up the result. That said, when the U-Net ouputs are of low quality, it might be a good idea to generate the masks directly with SAM. Once you have a good mask, you can save it using cv2.imwrite (see also the example notebook):

cv2.imwrite('sample1_mask.png', mask)

The patch_dir is the directory where the patches will be saved. A folder named ‘Patches’ will be created in this directory, and the patches will be saved in subfolders named ‘images’ and ‘labels’.

Next, training, validation, and test datasets are created from the patches:

train_dataset, val_dataset, test_dataset = seg.create_train_val_test_data(image_dir, mask_dir, augmentation=True)

Now we are ready to load the existing model weights and to train the model:

model = seg.create_and_train_model(train_dataset, val_dataset, test_dataset, model_file='seg_model.keras', epochs=100)

You can also evaluate the model separately on any dataset to get per-class IoU scores:

results = seg.evaluate_model(model, test_dataset)

This prints test loss, accuracy, mean IoU, and per-class IoU (background, grain, boundary), and returns them as a dictionary.

If you are happy with the finetuned model, you will want to save it:

model.save('seg_model_finetuned.keras')

If you want to use this new model to make predictions, you will need to load it with the custom loss function:

model = load_model("seg_model_finetuned.keras", custom_objects={'weighted_crossentropy': seg.weighted_crossentropy})

Training the U-Net model from scratch¶

If you want to train a U-Net model from scratch, you can use the Train_Unet_model.ipynb notebook, which mostly consists of the code snippets below.

import segmenteverygrain as seg
model = seg.Unet() # create model
model.compile(optimizer=Adam(), loss=seg.weighted_crossentropy, metrics=["accuracy"])

If you place the training images and masks in the ‘images’ directory (the filenames are supposed to terminate with ‘_image.png’ and ‘_mask.png’), you can create the training dataset like this:

input_dir = "../images/"
patch_dir = "../patches/"
image_dir, mask_dir = seg.patchify_training_data(input_dir, patch_dir)
image_dir = '../patches/Patches/images'
mask_dir = '../patches/Patches/labels'
train_dataset, val_dataset, test_dataset = seg.create_train_val_test_data(image_dir, mask_dir, augmentation=True)

Then you can train and test the model:

model = seg.create_and_train_model(train_dataset, val_dataset, test_dataset, epochs=200)

The model can be saved using the Keras save method:

model.save('seg_model.keras')

The U-Net model in the GitHub repository was trained using 66 images and the corresponding masks of a variety of grains, split into 61,013 patches of 256x256 pixels. 48 of these image-mask pairs are available at this Zenodo repository: https://zenodo.org/records/15786086. The model was trained for 200 epochs with a batch size of 32, using the Adam optimizer and a weighted cross-entropy loss function with label smoothing (0.1). The model is available in the repository as ‘seg_model_smooth_labels.keras’.