painting patches of a large labels layer using tensorstore

tensorstore_labels_patches.py

# %% 
# The idea behind this is to enable painting into a patch of a large
# labels layer that is loaded multiscale into napari
# By using tensorstore we can paint into a view of the array
# in memory and then commit the painting to disc async.
# This ends up being quite responsive and performant.
# %%
import numpy as np
import napari
import tensorstore as ts
import zarr
# %%
# Zarr setup
array_size = (20000, 20000)
zarr_path = r'test_zarr_array.zarr'
# %%
# write the zarr as a starting point
z = zarr.zeros(array_size, chunks=(512, 512), dtype='uint16')
zarr.save(zarr_path, z)

# %%
# setup tensorstore reading the zarr
spec = {
  'driver': 'zarr',
  'kvstore': {
    'driver': 'file', 
    'path': zarr_path,
  },
}

ts_array = ts.open(spec).result()
# %%
viewer = napari.Viewer()
# %%
# add the layer as multiscale -- this is view only basically
# napari Labels layers don't support painting in multiscale
# you can't even use the picker to sample the color!
viewer.add_labels([ts_array, ts_array[::2,::2], ts_array[::4,::4]], name="labels")
# %%
# add a writable crop, using transactions
# tried without transactions, but it was glitchy, maybe due to 2 arrays open?
# by using blending = 'opaque' we can see the live data
txn = ts.Transaction()
viewer.add_labels(ts_array.with_transaction(txn)[0:1000, 0:1000], name='paint', blending='opaque')
viewer.layers['paint'].bounding_box.visible = True
# do some painting in the box in the gui
# %%
# commit the painting, close the painting layer, update the data
txn.commit_async()
viewer.layers.remove(viewer.layers['paint'])
viewer.layers['labels'].refresh()
# %%
# make a new painting patch
# tensorstore returns a view with the view coordinates
# napari expects the corner to be at 0, 0.
# to fix this we use TS to translate the view to 0
# and then translate the napari layer to the correct location
txn = ts.Transaction()
viewer.add_labels(ts_array.with_transaction(txn)[ts.d[:][1000:2000, 1000:2000].translate_to[0]], name='paint', blending="opaque")
viewer.layers['paint'].translate = [1000, 1000]
viewer.layers['paint'].bounding_box.visible = True

# do some painting
# %%
# commit the painting, update the data, and close the painting layer
txn.commit_async()
viewer.layers.remove(viewer.layers['paint'])
viewer.layers['labels'].refresh()
# %%