diff --git a/src/ttmask/sphere.py b/src/ttmask/sphere.py index f5aabfc..19dd82b 100644 --- a/src/ttmask/sphere.py +++ b/src/ttmask/sphere.py @@ -3,11 +3,13 @@ import einops import napari import typer +from scipy.ndimage import distance_transform_edt @cli.command(name='sphere') def sphere( boxsize: int = typer.Option(...), sphere_diameter: float = typer.Option(...), + soft_edge_size: int = typer.Option(...), ): sphere_radius = sphere_diameter / 2 @@ -29,6 +31,12 @@ def sphere( idx = distance < sphere_radius mask[idx] = 1 + distance_from_edge = distance_transform_edt(mask == 0) + boundary_pixels = (distance_from_edge <= soft_edge_size) & (distance_from_edge != 0) + normalised_distance_from_edge = (distance_from_edge[boundary_pixels] / soft_edge_size) * np.pi + + mask[boundary_pixels] = (0.5 * np.cos(normalised_distance_from_edge) + 0.5) + viewer = napari.Viewer(ndisplay=3) viewer.add_image(mask) napari.run()