diff --git a/src/ttmask/__init__.py b/src/ttmask/__init__.py index e1e99b2..e464458 100644 --- a/src/ttmask/__init__.py +++ b/src/ttmask/__init__.py @@ -16,4 +16,5 @@ from .cube import cube from .cone import cone from .ellipsoid import ellipsoid +from .map2mask import map2mask diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py new file mode 100644 index 0000000..5c1010e --- /dev/null +++ b/src/ttmask/map2mask.py @@ -0,0 +1,30 @@ +import mrcfile +import numpy as np +import typer +from pathlib import Path + +from ._cli import cli +from .soft_edge import add_soft_edge + +@cli.command(name='map2mask') +def map2mask( + + input_map: Path = typer.Option(Path("map.mrc")), + binarization_threshold: float = typer.Option(...), + output_mask: Path = typer.Option(Path("mask.mrc")), + pixel_size: float = typer.Option(...), + soft_edge_width: int = typer.Option(0), + +): + with mrcfile.open(input_map) as mrc: + map_data = np.array(mrc.data) + + above_threshold = map_data >= binarization_threshold + below_threshold = map_data < binarization_threshold + + map_data[above_threshold] = 1 + map_data[below_threshold] = 0 + + mask = add_soft_edge(map_data, soft_edge_width) + + mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True)