shared_memory_hacking.py

# To use me:
# Run `(venv1) > python shared_memory_hacking.py --server`
# this will print the name of a shared memory buffer like `psm_5bee30a5` and open a napari window
# copy that name of the shared memorry buffer into another terminal with:
# `(venv2) > python shared_memory_hacking.py --client --shared_memory_name psm_5bee30a5`
# you should see another napari window pop up
# paint into the labels layer in either viewer - changes in one will be reflected in the other
# you need to refresh the visual to see the changes - for example hide/show the labels layer

import argparse
from multiprocessing.managers import SharedMemoryManager
from multiprocessing.shared_memory import SharedMemory

import numpy as np
from skimage import data
from skimage.color import rgb2gray
from skimage.segmentation import slic

import napari


def start_server():
    with SharedMemoryManager(
        address=('localhost', 8080), authkey=b'secret'
    ) as smm:
        print('started server', smm)
        smm.register('hello')
        array_shape = (512, 512)
        nbytes = np.prod(array_shape) * np.dtype(np.uint8).itemsize
        shared_data = smm.SharedMemory(size=nbytes)
        print(shared_data)
        shared_array = np.ndarray(
            array_shape, dtype=np.uint8, buffer=shared_data.buf
        )
        astro = data.astronaut()
        labels = slic(astro, channel_axis=-1, compactness=20) + 1
        shared_array[:] = labels.astype(np.uint8)

        viewer = napari.Viewer()
        astro = data.astronaut()
        viewer.add_image(rgb2gray(astro), name='astronaut', rgb=False)
        viewer.add_labels(shared_array, name='segmentation')
        napari.run()
        shared_data.close()


def start_client(shared_memory_name):
    smm = SharedMemoryManager(address=('localhost', 8080), authkey=b'secret')
    smm.connect()
    print('connected to server', smm)
    shared_data = SharedMemory(name=shared_memory_name)
    print(shared_data)
    shared_array = np.ndarray(
        (512, 512), dtype=np.uint8, buffer=shared_data.buf
    )
    viewer = napari.Viewer()
    viewer.add_labels(shared_array, name='segmentation')
    napari.run()
    shared_data.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--server', action='store_true', help='Start the shared memory server'
    )
    parser.add_argument(
        '--client', action='store_true', help='Start the shared memory client'
    )
    parser.add_argument(
        '--shared_memory_name',
        type=str,
        help='Name of the shared memory to connect to (client only)',
    )

    args = parser.parse_args()
    if args.server:
        start_server()
    elif args.client:
        start_client(args.shared_memory_name)