feat: more seg channels

Author: seankmartin

examples/load_tiff_seg.py

@@ -1,5 +1,6 @@
 # %% Import packages
 
+from copy import copy
 import itertools
 from pathlib import Path
 from bfio import BioReader
@@ -60,12 +61,11 @@
 
 num_channels = br.shape[-1]
 data_type = "uint16"
-chunk_size = [256, 256, 128]
 volume_size = [br.shape[1], br.shape[0], br.shape[2]]  # XYZ
 
 # %% Setup somewhere to hold progress
+CHANNEL = 0
 progress_dir = OUTPUT_PATH / "progress"
-progress_dir.mkdir(exist_ok=True)
 
 # %% Functions for moving data
 shape = np.array([br.shape[1], br.shape[0], br.shape[2]])
@@ -79,7 +79,13 @@ def chunked_reader(x_i, y_i, z_i):
     z_start, z_end = z_i * chunk_shape[2], min((z_i + 1) * chunk_shape[2], shape[2])
 
     # Read the chunk from the BioReader
-    chunk = br.read(X=(x_start, x_end), Y=(y_start, y_end), Z=(z_start, z_end))
+    chunk = br.read(
+        X=(x_start, x_end), Y=(y_start, y_end), Z=(z_start, z_end), C=(CHANNEL,)
+    )
+    chunk = np.atleast_3d(chunk)
+    # Remove the last flattened dimension
+    if (len(chunk.shape) > 3) and (chunk.shape[-1] == 1):
+        chunk = chunk[:, :, :, 0]
 
     # Return the chunk
     return chunk.swapaxes(0, 1)
@@ -99,21 +105,22 @@ def process(args):
     )
     if f_name.exists() and not OVERWRITE:
         return
-    print("Working on", f_name)
     rawdata = chunked_reader(x_i, y_i, z_i)
     # TEMP
-    print(np.unique(rawdata))
     print(rawdata.shape)
     # Create the segmentation mask
-    dimensions = [start, end]
-    seg_chunk = create_segmentation_chunk(rawdata, dimensions, convert_non_zero_to=None)
+    # The writer expects the data to be in ZYX order so need to swap the axes
+    start_zyx = [start[2], start[1], start[0]]
+    end_zyx = [end[2], end[1], end[0]]
+    dimensions = [start_zyx, end_zyx]
+    seg_chunk = create_segmentation_chunk(rawdata, dimensions, convert_non_zero_to=1)
     seg_chunk.write_to_directory(OUTPUT_PATH / "data")
     touch(f_name)
 
 
 # %% Try with a single chunk to see if it works
-x_i, y_i, z_i = 0, 0, 0
-process((x_i, y_i, z_i))
+# x_i, y_i, z_i = 0, 0, 0
+# process((x_i, y_i, z_i))
 
 # %% Loop over all the chunks
 coords = itertools.product(
@@ -126,34 +133,43 @@ def process(args):
 reversed_coords.reverse()
 
 # %% Create the metadata
+int_size = [int(size_x), int(size_y), int(size_z)]
+real_size = [size_x, size_y, size_z]
 metadata = {
     "@type": "neuroglancer_multiscale_volume",
     "data_type": "uint32",
     "num_channels": 1,
     "scales": [
         {
-            "chunk_sizes": [chunk_size],
+            "chunk_sizes": [[int(c) for c in chunk_shape]],
             "encoding": "compressed_segmentation",
-            "copmressed_segmentation_block_size": [8, 8, 8],
-            "resolution": [size_x, size_y, size_z],
+            "compressed_segmentation_block_size": [8, 8, 8],
+            "resolution": int_size,
             "key": "data",
             "size": volume_size,
         }
     ],
     "mesh": "mesh",
     "type": "segmentation",
 }
-write_metadata(
-    metadata,
-    OUTPUT_PATH,
-    overwrite=OVERWRITE,
-)
-
-# %% Now create the mesh
-mesh_shape = np.array(volume_size)
-generate_multiresolution_mesh_from_segmentation(OUTPUT_PATH, "mesh", 3, mesh_shape)
-
 
 # %% Move the data across with a single worker
-for coord in reversed_coords:
-    process(coord)
+original_path = copy(OUTPUT_PATH)
+for i in range(0, 1):
+    OUTPUT_PATH = original_path.with_stem(f"{original_path.stem}_ch{i}")
+    OUTPUT_PATH.mkdir(exist_ok=True, parents=True)
+    print("Writing to ", OUTPUT_PATH)
+    # Make the progress directory
+    progress_dir = OUTPUT_PATH / "progress"
+    progress_dir.mkdir(exist_ok=True)
+    CHANNEL = i
+    write_metadata(metadata, OUTPUT_PATH)
+    for coord in reversed_coords:
+        process(coord)
+        print("COMPLETED", i)
+    mesh_shape = np.array(volume_size)
+    generate_multiresolution_mesh_from_segmentation(
+        OUTPUT_PATH, "mesh", 2, mesh_shape, fill_missing=True
+    )
+
+# %% For running all