FIX(psoct): propagate no_pool to pyramid generator + FIX(generate_pyr…

…amid): do not crash if last chunk in a row only has a single voxel
lincbrain · Nov 22, 2024 · 41cbd1f · 41cbd1f
1 parent 7ecfc08
commit 41cbd1f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/linc_convert/modalities/psoct/_utils.py b/linc_convert/modalities/psoct/_utils.py
@@ -203,17 +203,25 @@ def generate_pyramid(
                 slice(i * max_load, min((i + 1) * max_load, n))
                 for i, n in zip(chunk_index, prev_shape)
             ]
+            fullshape = omz[str(level - 1)].shape
             dat = omz[str(level - 1)][tuple(slicer)]
 
             # Discard the last voxel along odd dimensions
-            crop = [0 if x == 1 else x % 2 for x in dat.shape[-ndim:]]
+            crop = [
+                0 if y == 1 else x % 2
+                for x, y in zip(dat.shape[-ndim:], fullshape)
+            ]
             # Don't crop the axis not down-sampling
             # cannot do if not no_pyramid_axis since it could be 0
             if no_pyramid_axis is not None:
                 crop[no_pyramid_axis] = 0
             slcr = [slice(-1) if x else slice(None) for x in crop]
             dat = dat[tuple([Ellipsis, *slcr])]
 
+            if any(n == 0 for n in dat.shape):
+                # last strip had a single voxel, nothing to do
+                continue
+
             patch_shape = dat.shape[-ndim:]
 
             # Reshape into patches of shape 2x2x2
@@ -234,7 +242,7 @@ def generate_pyramid(
             # -> flatten patches
             smaller_shape = [max(n // 2, 1) for n in patch_shape]
             if no_pyramid_axis is not None:
-                smaller_shape[2 * no_pyramid_axis] = patch_shape[no_pyramid_axis]
+                smaller_shape[no_pyramid_axis] = patch_shape[no_pyramid_axis]
 
             dat = dat.reshape(batch + smaller_shape + [-1])
 

diff --git a/linc_convert/modalities/psoct/multi_slice.py b/linc_convert/modalities/psoct/multi_slice.py
@@ -294,7 +294,7 @@ def convert(
 
         inp[i] = None  # no ref count -> delete array
 
-    generate_pyramid(omz, nblevels - 1, mode="mean")
+    generate_pyramid(omz, nblevels - 1, mode="mean", no_pyramid_axis=no_pool)
 
     print("")
 

diff --git a/linc_convert/modalities/psoct/single_volume.py b/linc_convert/modalities/psoct/single_volume.py
@@ -208,7 +208,7 @@ def convert(
             i * inp_chunk[2]: i * inp_chunk[2] + loaded_chunk.shape[2],
         ] = loaded_chunk
 
-    generate_pyramid(omz, nblevels - 1, mode="mean")
+    generate_pyramid(omz, nblevels - 1, mode="mean", no_pyramid_axis=no_pool)
 
     print("")