Skip to content

Commit

Permalink
add operator save-swc; upgrade the load-zarr to use zarr v4 protocol;…
Browse files Browse the repository at this point in the history
… show skeletons in neuroglancer but it is not working yet
  • Loading branch information
xiuliren committed Dec 8, 2023
1 parent fcf5b01 commit f26c0b8
Show file tree
Hide file tree
Showing 5 changed files with 174 additions and 41 deletions.
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,7 @@ After installation, You can simply type `chunkflow` and it will list all the ope
| save-precomputed| Save chunk to local/cloud storage volume |
| save-tif | Save chunk as TIFF file |
| save-synapses | Save synapses as a HDF5 file. |
| save-swc | Save skeletons as a SWC file. |

## Affiliation
This package is developed at Princeton University and Flatiron Institute.
Expand Down
82 changes: 66 additions & 16 deletions chunkflow/flow/flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ def adjust_bbox(tasks, corner_offset: tuple):
@main.command('skip-task-by-file')
@click.option('--prefix', '-p', required=True, type=str,
help='the pre part of result file path')
@click.option('--suffix', '-s', required=True, type=str,
@click.option('--suffix', '-s', default='', type=str,
help='the post part of result file path. Normally include file extention.')
@click.option('--mode', '-m',
type=click.Choice(['missing', 'empty', 'exist']), default='exist',
Expand Down Expand Up @@ -270,9 +270,9 @@ def skip_task_by_blocks_in_volume(tasks, volume_path: str, mip: int, use_https:


@main.command('mark-complete')
@click.option('--prefix', '-p', type=str, default=None,
@click.option('--prefix', '-p', type=str, required=True,
help='pre-path of a file.')
@click.option('--suffix', '-s', type=str, default=None,
@click.option('--suffix', '-s', type=str, default='',
help='suffix of the flag file.')
@operator
def mark_complete(tasks, prefix: str, suffix: str):
Expand All @@ -298,7 +298,7 @@ def mark_complete(tasks, prefix: str, suffix: str):
type=str, default=DEFAULT_CHUNK_NAME, help='input chunk name')
@click.option('--prefix', '-p', type=str, default=None,
help = 'pre-path of a file. we would like to keep a trace that this task was executed.')
@click.option('--suffix', '-s', type=str, default=None,
@click.option('--suffix', '-s', type=str, default='',
help='post-path of a file. normally include the extention of result file.')
@click.option('--adjust-size', '-a', type=click.INT, default=None,
help='change the bounding box of chunk if it do not match with final result file name.')
Expand All @@ -312,7 +312,7 @@ def skip_all_zero(tasks, input_chunk_name: str, prefix: str, suffix: str, adjust
chunk = task[input_chunk_name]
if not np.any(chunk):
print('all zero chunk, skip this task')
if prefix is not None or suffix is not None:
if prefix is not None:
if chunk_bbox:
bbox = chunk.bbox.clone()
else:
Expand Down Expand Up @@ -1261,15 +1261,18 @@ def load_precomputed(tasks, name: str, volume_path: str, mip: int,
help='voxel offset for easier usage. This will be translated to physical coordinate using voxel size.')
@click.option('--voxel-size', '-s', type=click.FLOAT, nargs=3, default=None,
help='voxel size for translating voxel offset to physical coordinate offset. We do not set default 1x1x1 in case it is not set with voxel_offset together.')
@click.option('--output-name', '-o', type-str, default='skel', help='output name of oid2skel. Note that it is a dict to map object ID to skeleton.')
@click.option('--output-name', '-o', type=str, default=DEFAULT_SKELETON_NAME, \
help='output name of oid2skel. Note that it is a dict to map object ID to skeleton.')
@operator
def load_skeleton(tasks, fname: str, offset: Tuple, voxel_offset: Tuple, voxel_size: Tuple, output_name: str):
if offset is None and voxel_offset is not None:
assert voxel_size is not None
voxel_offset = Cartesian.from_collection(voxel_offset)
voxel_size = Cartesian.from_collection(voxel_size)
offset = voxel_offset * voxel_size
if not isinstance(offset, Cartesian):
if offset is None:
offset = Cartesian(0, 0, 0)
elif not isinstance(offset, Cartesian):
offset = Cartesian.from_collection(offset)

for task in tasks:
Expand Down Expand Up @@ -1299,34 +1302,81 @@ def load_skeleton(tasks, fname: str, offset: Tuple, voxel_offset: Tuple, voxel_s
for oid, skels in oid2skel.items():
skel = Skeleton.simple_merge(skels).consolidate()
if offset is not None:
skel.vertices[:] += np.asarray((offset)[::-1]
skel.vertices[:] += np.asarray(offset)[::-1]
oid2skel[oid] = skel
task[output_name] = oid2skel
yield task


@main.command('save-swc')
@click.option('--input-name', '-i', type=str, default=DEFAULT_SKELETON_NAME, \
help='input name in task')
@click.option('--output-prefix', '-p', type=str, default='./', \
help='prefix of output file name.')
@operator
def save_swc(tasks, input_name: str, output_prefix: str):
for task in tasks:
if task is not None:
if os.path.isdir(output_prefix) and not output_prefix.endswith('/'):
output_prefix += '/'

skels = task[input_name]
for sid, skel in skels.items():
fname = f'{output_prefix}{sid}.swc'
with open(fname, 'w') as file:
file.write(skel.to_swc())
yield task

@main.command('load-npy')
@click.option('--file-name', '-f',
type=click.Path(exists=True, file_okay=True, dir_okay=False, readable=True, resolve_path=True), required=True, help='file name path')
@click.option('--voxel-offset', '-t', type=int, nargs=3, default=(0,0,0),
help='voxel offset')
@click.option('--voxel-size', '-s', type=click.FLOAT, nargs=3,
default=(1,1,1), help='voxel size')
@click.option('--output-name', '-o', type=str, default=DEFAULT_CHUNK_NAME,
help='output name in the task dictionary.')
@operator
def load_npy(tasks, file_name: str, voxel_offset: Tuple, voxel_size: Tuple,
output_name: str):
for task in tasks:
if task is not None:
assert file_name.endswith('.npy')
arr = np.load(file_name)
voxel_offset = Cartesian.from_collection(voxel_offset)
voxel_size = Cartesian.from_collection(voxel_size)
chunk = Chunk(arr, voxel_offset=voxel_offset, voxel_size=voxel_size)
task[output_name] = chunk
yield task


@main.command('load-zarr')
@click.option('--store', '-f', type=str, required=True,
help='Zarr store path')
help='Zarr store path: n5, local, precomputed')
@click.option('--path', '-p', type=str, default = None,
help = 'Zarr path in the store')
@click.option('--chunk-start', '-s', type=click.INT, nargs=3, default=None,
help='voxel offset or start')
@click.option('--chunk-size', '-z', type=click.INT, nargs=3, default=None,
help='chunk size')
@click.option('--voxel-size', '-v', type=click.FLOAT, nargs=3, default=None)
@click.option('--backend', '-b', type=str, default='NestedDirectoryStore',
@click.option('--driver', '-b', type=click.Choice(['n5', 'local']), default='NestedDirectoryStore',
help='the storage backend.')
@click.option('--output-chunk-name', '-o', type=str, default=DEFAULT_CHUNK_NAME,
help='output chunk name.')
@operator
def load_zarr(tasks, store: str, path: str, chunk_start: tuple, voxel_size: tuple,
chunk_size: tuple, backend: str, output_chunk_name: str):
chunk_size: tuple, driver: str, output_chunk_name: str):
"""Load Zarr arrays."""
if backend == 'NestedDirectoryStore':
if driver in ['NestedDirectoryStore', 'local'] :
store = zarr.NestedDirectoryStore(store)
z = zarr.open(store, mode='r', path=path)
attrs = z.attrs.asdict()
elif driver in ['n5', 'N5FSStore']:
store = zarr.N5FSStore(store, anon=True)
else:
raise ValueError(f'invalide store type, we only support local and n5, but got {store}')

zdata = zarr.convenience.open(store, mode='r', path=path)
attrs = zdata.attrs.asdict()

# Note that this is the physical
if 'offset' in attrs:
Expand All @@ -1350,7 +1400,7 @@ def load_zarr(tasks, store: str, path: str, chunk_start: tuple, voxel_size: tupl
for task in tasks:
if task is not None:
if chunk_size is None and chunk_start is None and 'bbox' not in task:
arr = z[:]
arr = zdata[:]
voxel_offset = volume_offset
else:
if chunk_start is not None and chunk_size is not None:
Expand All @@ -1364,7 +1414,7 @@ def load_zarr(tasks, store: str, path: str, chunk_start: tuple, voxel_size: tupl

arr_start = bbox.start - volume_offset
arr_bbox = BoundingBox.from_delta(arr_start, bbox.shape)
arr = z[arr_bbox.slices]
arr = zdata[arr_bbox.slices]
voxel_offset = chunk_start
chunk = Chunk(arr, voxel_offset=voxel_offset, voxel_size=voxel_size)
task[output_chunk_name] = chunk
Expand Down
Loading

0 comments on commit f26c0b8

Please sign in to comment.