create segmentation masks

pick2d.py

import napari
import mrcfile

import numpy as np

from pathlib import Path
from magicgui import magicgui
from napari.types import LabelsData, PointsData, ImageData
from scipy.ndimage import distance_transform_edt

viewer = napari.Viewer(ndisplay=2)


@magicgui(call_button='Open')
def micrographs_in_folder(mic: Path):
    with mrcfile.open(mic) as mrc:
        img_data = mrc.data
    plane_layer = viewer.add_image(img_data)
    points_layer = viewer.add_points([], ndim=2)
    return plane_layer

@magicgui(call_button='Segment It!')
def points_to_semantic_segmentation(Points: PointsData, img: ImageData, Radius: float) -> LabelsData:
    labels_data = np.ones_like(img).astype(int)
    indices = tuple(np.round(Points).astype(int).T)  # Shape 2,3,3
    labels_data[indices] = 0
    edt = distance_transform_edt(labels_data)
    labels_data = edt < Radius
    return labels_data


@magicgui(call_button='Save Segmentation')
def save_mask(basename: str, segmentation: LabelsData):
    output_fname = f'{basename}_mask.mrc'
    mrcfile.new(output_fname, data=segmentation)


viewer.window.add_dock_widget(micrographs_in_folder)
viewer.window.add_dock_widget(points_to_semantic_segmentation)
viewer.window.add_dock_widget(save_mask)
napari.run()