basic 3D picker

basic_3d_point_picker.py

from pathlib import Path

import mrcfile
import napari
import numpy as np
from magicgui import magicgui

OUTPUT_DIRECTORY = 'picking'
Path(OUTPUT_DIRECTORY).mkdir(exist_ok=True, parents=True)

viewer = napari.Viewer(ndisplay=3)
viewer.text_overlay.visible = True

points_layer = viewer.add_points([], ndim=3, face_color='cornflowerblue')
image_layer = viewer.add_image(np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]),
                               metadata={'filename': 'test.csv'})

image_layer.depiction = 'plane'
image_layer.plane.thickness = 10
image_layer.rendering = 'average'


def point_in_bounding_box(point: np.ndarray, bounding_box: np.ndarray) -> bool:
    """Determine whether an nD point is inside an nD bounding box.
    Parameters
    ----------
    point : np.ndarray
        (n,) array containing nD point coordinates to check.
    bounding_box : np.ndarray
        (2, n) array containing the min and max of the nD bounding box.
        As returned by `Layer._extent_data`.
    """
    if np.all(point > bounding_box[0]) and np.all(point < bounding_box[1]):
        return True
    return False


def add_point_on_plane(
        viewer,
        event,
        points_layer: napari.layers.Points = points_layer,
        plane_layer: napari.layers.Image = image_layer,
        append: bool = True,
):
    # Early exit if not alt-clicked
    if 'Alt' not in event.modifiers:
        return

    # Early exit if volume_layer isn't visible
    if not plane_layer.visible:
        return

    # Calculate intersection of click with plane through data in data coordinates
    intersection = plane_layer.plane.intersect_with_line(
        line_position=viewer.cursor.position, line_direction=viewer.cursor._view_direction
    )

    # Check if click was on plane by checking if intersection occurs within
    # data bounding box. If not, exit early.
    if not point_in_bounding_box(intersection, plane_layer.extent.data):
        return

    if append:
        points_layer.add(intersection)
    else:
        points_layer.data = intersection
        # points_layer.add(intersection)


viewer.mouse_drag_callbacks.append(add_point_on_plane)


def reset_contrast():
    image_layer.reset_contrast_limits_range()
    image_layer.reset_contrast_limits()


@magicgui(auto_call=True)
def load_tomogram(file: Path):
    data = mrcfile.open(file).data
    image_layer.data = data
    image_layer.metadata['filename'] = file.stem
    viewer.text_overlay.text = file.stem
    points_layer.data = []
    reset_contrast()


viewer.window.add_dock_widget(load_tomogram)


@magicgui(call_button='save coordinates')
def save_coordinates():
    xyz = points_layer.data[:, ::-1]
    filename = Path(OUTPUT_DIRECTORY) / f"{image_layer.metadata['filename']}.csv"
    np.savetxt(filename, xyz, delimiter=', ', fmt='%.2f')


viewer.window.add_dock_widget(save_coordinates)


@viewer.bind_key('c')
def center_on_image_layer(*args):
    viewer.camera.center = np.array(image_layer.extent.world).mean(axis=0)


napari.run()