diff --git a/src/ttmask/add_padding.py b/src/ttmask/add_padding.py new file mode 100644 index 0000000..38435f2 --- /dev/null +++ b/src/ttmask/add_padding.py @@ -0,0 +1,10 @@ +import numpy as np +from scipy.ndimage import distance_transform_edt + + +def add_padding(mask: np.ndarray, width: float) -> np.ndarray: + distance_from_edge = distance_transform_edt(mask == 0) + boundary_pixels = (distance_from_edge <= width) & (distance_from_edge != 0) + output = np.copy(mask) + output[boundary_pixels] = 1 + return output diff --git a/src/ttmask/map2mask.py b/src/ttmask/map2mask.py index 5c1010e..90bec28 100644 --- a/src/ttmask/map2mask.py +++ b/src/ttmask/map2mask.py @@ -5,6 +5,7 @@ from ._cli import cli from .soft_edge import add_soft_edge +from .add_padding import add_padding @cli.command(name='map2mask') def map2mask( @@ -14,7 +15,7 @@ def map2mask( output_mask: Path = typer.Option(Path("mask.mrc")), pixel_size: float = typer.Option(...), soft_edge_width: int = typer.Option(0), - + padding_width: int = typer.Option(0), ): with mrcfile.open(input_map) as mrc: map_data = np.array(mrc.data) @@ -25,6 +26,7 @@ def map2mask( map_data[above_threshold] = 1 map_data[below_threshold] = 0 - mask = add_soft_edge(map_data, soft_edge_width) + padded_mask = add_padding(map_data, padding_width) + mask = add_soft_edge(padded_mask, soft_edge_width) - mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True) + mrcfile.write(output_mask, mask, voxel_size=pixel_size, overwrite=True) \ No newline at end of file