Combine two binary segmentations into one quantized segmentation with values 1 and 2.

seg_combinator.py

# Take two binary segmentations in mrc format and combine them, with the first segmentation having label 1 and the second segmentation having label 2.

import mrcfile
import numpy as np
import argparse

def combine_segmentations(seg1, seg2, clear_overlaps=True):
    """
    Combine two binary segmentations into one, with the first segmentation having label 1 and the second segmentation having label 2.

    Assumes segmentations are non-overlapped.
    
    clear_overlaps: If true, sets values over 2 to 0 to clear overlaps.
    """
    # Ensure both segmentations are binary
    seg1 = (seg1 > 0).astype(np.uint8)
    seg2 = (seg2 > 0).astype(np.uint8)

    # Combine the segmentations
    combined_segmentation = seg1 + 2 * seg2
    
    if clear_overlaps:
        combined_segmentation[combined_segmentation>2] = 0


    return combined_segmentation

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Combine two binary segmentations in mrc format.")
    parser.add_argument("-s1","--segmentation1", type=str, help="Path to the first segmentation file (mrc format).", required=True)
    parser.add_argument("-s2","--segmentation2", type=str, help="Path to the second segmentation file (mrc format).", required=True)
    parser.add_argument("-c","--clear_overlaps",type=bool,help="Remove segments of overlap to simplify magic-wanding", required=False, default=True)
    parser.add_argument("-o", "--output", type=str, help="Path to the output combined segmentation file (mrc format).", default="combined_segmentation.mrc")

    args = parser.parse_args()
    voxel_size = 1.0  # Default voxel size
    # Read the first segmentation
    with mrcfile.open(args.segmentation1, mode='r') as mrc:
        seg1 = mrc.data
        voxel_size = mrc.voxel_size


    # Read the second segmentation
    with mrcfile.open(args.segmentation2, mode='r') as mrc:
        seg2 = mrc.data

    # Combine the segmentations
    combined_segmentation = combine_segmentations(seg1, seg2, clear_overlaps=args.clear_overlaps)

    # Write the combined segmentation to a new mrc file
    with mrcfile.new(args.output, overwrite=True) as mrc:
        mrc.set_data(combined_segmentation)
        mrc.voxel_size = voxel_size  # Set voxel size if needed

    print(f"Combined segmentation saved to {args.output}")

Should just need pip install mrcfile numpy to run normally.

Usage: python seg_combinator.py -s1 SEG1.mrc -s2 SEG2.mrc