Generate binary masks from IMOD model files (and saves them as MRC files).

mod2mask.py

import typer
import logging
from pathlib import Path
import imodmodel
import numpy as np
import matplotlib.path as mplPath
import mrcfile
from scipy.interpolate import interp1d

app = typer.Typer(add_completion=False)


def setup_logger(verbosity: int):
    logging_levels = {
        0: logging.WARNING,
        1: logging.INFO,
        2: logging.DEBUG
    }
    logging.basicConfig(level=logging_levels.get(verbosity, logging.DEBUG),
                        format='%(asctime)s - %(levelname)s - %(message)s')
    return logging.getLogger(__name__)


def readmod(mod, contour_id: int, orientation1, orientation2, logger):
    model = imodmodel.read(mod)
    obj = model[model['object_id'] == contour_id][[orientation1, orientation2]].to_numpy()
    logger.debug(f"Read contour {contour_id} from {mod}. Shape: {obj.shape}")
    return obj


def create_3d_mask(obj_xy, obj_xz, shape, logger):
    z_min, z_max = int(np.min(obj_xz[:, 1])), int(np.max(obj_xz[:, 1]))
    z_range = np.arange(z_min, z_max + 1)

    mask = np.zeros(shape, dtype=np.uint8)
    _, unique_indices = np.unique(obj_xz[:, 1], return_index=True)
    obj_xz_unique = obj_xz[unique_indices]

    xz_width = np.max(obj_xz_unique[:, 0]) - np.min(obj_xz_unique[:, 0])
    xz_scale = (obj_xz_unique[:, 0] - np.min(obj_xz_unique[:, 0])) / xz_width
    scale_func = interp1d(obj_xz_unique[:, 1], xz_scale, kind='linear', bounds_error=False, fill_value='extrapolate')

    xy_center = np.mean(obj_xy, axis=0)

    x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
    xy_points = np.column_stack((x.ravel(), y.ravel()))

    scales = scale_func(z_range)

    for i, z in enumerate(z_range):
        if z < 0 or z >= shape[2]:
            continue
        logger.debug(f"Processing slice z={z}")
        scale = scales[i]
        scaled_xy = (obj_xy - xy_center) * scale + xy_center
        scaled_polygon = mplPath.Path(scaled_xy)

        mask[:, :, z] = scaled_polygon.contains_points(xy_points).reshape(shape[0], shape[1])
        logger.debug(f"Slice {z} sum: {np.sum(mask[:, :, z])}")

    for x in range(shape[0]):
        for y in range(shape[1]):
            z_profile = mask[x, y, :]
            z_indices = np.where(z_profile)[0]
            if len(z_indices) > 1:
                mask[x, y, z_indices[0]:z_indices[-1] + 1] = 1

    return mask


def get_volume_info(volume_path: Path, logger):
    mrc_files = list(volume_path.glob("*.mrc"))
    if not mrc_files:
        raise ValueError(f"No MRC files found in {volume_path}")

    first_mrc = mrc_files[0]
    logger.info(f"Using {first_mrc} for volume information")

    with mrcfile.open(first_mrc) as mrc:
        shape = mrc.data.shape
        voxel_size = mrc.voxel_size.x

    logger.info(f"Volume shape (Z,Y,X): {shape}")
    logger.info(f"Voxel size: {voxel_size}")

    return shape[::-1], voxel_size  # Return shape as (X,Y,Z)


@app.command(no_args_is_help=True)
def generate_mask(
        input_path: Path = typer.Option(..., "--i", help="Path to the directory containing .mod files"),
        output_path: Path = typer.Option("./masks/", "--o", help="Path to save the output mask files"),
        volume_path: Path = typer.Option(..., "--v", help="Path to the directory containing volume data (.mrc files)"),
        verbosity: int = typer.Option(1, help="Verbosity level (0: WARNING, 1: INFO, 2: DEBUG)")
):
    """
    Script Description:\n
    This script generates binary masks from IMOD model files (.mod) and saves them as MRC files. \n
    It uses the dimensions and voxel size from existing volume data to ensure compatibility.\n
    The idea is very strongly inspired by https://www.ncbi.nlm.nih.gov/pmc/articles/PMC9596874/\n
    Workflow:\n
    1. Reads volume information (shape and voxel size) from the first MRC file in the specified volume directory.\n
    2. Processes each .mod file in the input directory:\n
       a. Reads XY and XZ contours from the .mod file.\n
       b. Creates a 3D binary mask based on these contours.\n
       c. Saves the resulting mask as an MRC file in the output directory.\n
    \n
    Key Assumptions:\n
    1. The first MRC file in the volume directory is representative of all volumes in terms of shape and voxel size.\n
    2. Each .mod file contains exactly two objects: one for the XY contour (object_id 0) and one for the XZ contour (object_id 1).\n
    3. The XY contour represents the shape of the mask in each Z-slice.\n
    4. The XZ contour represents how the shape scales across different Z-slices.\n
    \n
    Potential Failure Cases:\n
    1. No MRC files in the specified volume directory.\n
    2. Inconsistent volume sizes or voxel sizes among MRC files (only the first file is used for reference).\n
    3. .mod files with missing or incorrectly numbered objects.\n
    4. XY or XZ contours that extend beyond the volume boundaries.\n
    5. Extremely complex or self-intersecting contours that may cause unexpected behavior in mask generation.\n
    6. Insufficient disk space to save output mask files.\n
    7. Lack of read/write permissions in the input, volume, or output directories.\n
    \n
    The script uses logging to provide detailed information about its progress and any issues encountered during execution. Adjust the verbosity level for more or less detailed output.
    """
    logger = setup_logger(verbosity)

    volume_shape_xyz, voxel_size = get_volume_info(volume_path, logger)
    logger.info(f"Using volume shape (X,Y,Z): {volume_shape_xyz}, voxel size: {voxel_size}")

    modfiles = list(input_path.glob("*.mod"))
    logger.info(f"Found {len(modfiles)} mod files")

    if output_path.exists() and not output_path.is_dir():
        output_path.mkdir(parents=True, exist_ok=True)
        logger.info(f"Output path {output_path} is not a directory. Creating directory.")

    elif not output_path.exists():
        logger.info(f"Creating output directory {output_path}")
        output_path.mkdir(parents=True, exist_ok=True)

    for modfile in modfiles:
        logger.info(f"Processing {modfile}")
        obj_xy = readmod(modfile, 0, 'x', 'y', logger)
        obj_xz = readmod(modfile, 1, 'x', 'z', logger)

        logger.debug("XY contour:")
        logger.debug(obj_xy)
        logger.debug("XZ contour:")
        logger.debug(obj_xz)

        mask = create_3d_mask(obj_xy, obj_xz, volume_shape_xyz, logger)

        mask_permuted = np.transpose(mask, (2, 1, 0))

        logger.info(f"Final mask shape: {mask_permuted.shape}, sum: {np.sum(mask_permuted)}")

        output_file = output_path / f"{modfile.stem}.mrc"
        mrcfile.write(output_file, mask_permuted.astype(np.uint8), overwrite=True, voxel_size=voxel_size)

        logger.info(f"Generated binary mask for {modfile.stem}")

    logger.info("Script execution completed successfully.")

if __name__ == "__main__":
    app()