Fix single file backprojection output

- TiffWriter cannot be pickled (and arguments were wrong anyway).
- Attempting to write single file as it comes in is very slow due to order requirements, so switched (back) to wait to finish and rewrite.

Author: den-sq

python/ouroboros/helpers/files.py

@@ -7,7 +7,7 @@
 import numpy as np
 from numpy.typing import ArrayLike
 from pathlib import Path
-from tifffile import imread, TiffWriter, memmap, TiffFile
+from tifffile import imread, TiffWriter, TiffFile
 import time
 
 from .shapes import DataShape
@@ -139,7 +139,8 @@ def np_convert(dtype: np.dtype, source: ArrayLike, normalize=True):
         return source.astype(dtype)
 
 
-def generate_tiff_write(write_func, compression, micron_resolution, backprojection_offset):
+def generate_tiff_write(write_func: callable, compression: str | None, micron_resolution: np.ndarray[float],
+               backprojection_offset: np.ndarray, **kwargs):
     # Volume cache resolution is in voxel size, but .tiff XY resolution is in voxels per unit, so we invert.
     resolution = [1.0 / voxel_size for voxel_size in micron_resolution[:2] * 0.0001]
     resolutionunit = "CENTIMETER"
@@ -158,7 +159,8 @@ def generate_tiff_write(write_func, compression, micron_resolution, backprojecti
                    metadata=metadata,
                    resolution=resolution,
                    resolutionunit=resolutionunit,
-                   software="ouroboros")
+                   software="ouroboros",
+                   **kwargs)
 
 
 def write_small_intermediate(file_path: os.PathLike, *series):

python/ouroboros/helpers/memory_usage.py

@@ -40,7 +40,7 @@ def calculate_gigabytes_from_dimensions(shape: tuple[int], dtype: np.dtype) -> f
     num_elements = np.prod(shape)
 
     # Calculate the total number of bytes
-    num_bytes = np.multiply(num_elements, dtype_size, dtype=np.uint64)
+    num_bytes = np.multiply(num_elements, dtype_size, dtype=np.uint64, casting='unsafe')
 
     return num_bytes / GIGABYTE
 

python/ouroboros/pipeline/backproject_pipeline.py

@@ -135,13 +135,14 @@ def _process(self, input_data: any) -> tuple[any, None] | tuple[None, any]:
 
         if config.make_single_file:
             is_big_tiff = calculate_gigabytes_from_dimensions(np.prod(write_shape), np.uint16) > 4     # Check Dtype
-            single_tiff = tifffile.TiffWriter(folder_path.with_suffix(".tiff"), abigtiff=is_big_tiff)
+        else:
+            is_big_tiff = calculate_gigabytes_from_dimensions(np.prod(write_shape[1:]), np.uint16) > 4     # Check Dtype
 
         bp_offset = pipeline_input.backprojection_offset if config.backproject_min_bounding_box else None
-        tif_write = generate_tiff_write(single_tiff.write if config.make_single_file else tifffile.imwrite,
-                                        config.backprojection_compression,
-                                        volume_cache.get_resolution_um(),
-                                        bp_offset)
+        tif_write = partial(generate_tiff_write,
+                            compression=config.backprojection_compression,
+                            micron_resolution=volume_cache.get_resolution_um(),
+                            backprojection_offset=bp_offset)
 
         # Process each bounding box in parallel, writing the results to the backprojected volume
         try:
@@ -198,17 +199,12 @@ def note_written(write_future):
 
                     if np.any(writeable == 1):
                         write = np.flatnonzero(writeable == 1)
-                        if config.make_single_file:
-                            write = write[write == (np.indices(write.shape) + len(writeable > 1))]
-
-                        # Will need to multiprocess
+                        # Single File needs to be in order
                         for index in write:
-                            path_args = [] if config.make_single_file else [folder_path.joinpath(f"{index:05}.tiff")]
-
                             write_futures.append(write_executor.submit(
                                 write_conv_vol,
-                                tif_write, i_path.joinpath(f"i_{index:05}"),
-                                ImgSlice(*write_shape[1:]), np.uint16, *path_args
+                                tif_write(tifffile.imwrite), i_path.joinpath(f"i_{index:05}"),
+                                ImgSlice(*write_shape[1:]), np.uint16, folder_path.joinpath(f"{index:05}.tif")
                             ))
                             write_futures[-1].add_done_callback(note_written)
 
@@ -227,7 +223,9 @@ def note_written(write_future):
         start = time.perf_counter()
 
         if config.make_single_file:
-            shutil.rmtree(folder_path)
+            writer = tif_write(tifffile.TiffWriter(folder_path.with_suffix(".tiff"), bigtiff=is_big_tiff).write)
+            for fname in get_sorted_tif_files(folder_path):
+                writer(tifffile.imread(folder_path.joinpath(fname)))
 
         # Rescale the backprojected volume to the output mip level
         if pipeline_input.slice_options.output_mip_level != config.output_mip_level:
@@ -265,6 +263,9 @@ def note_written(write_future):
         pipeline_input.backprojected_folder_path = folder_path
 
         self.add_timing("export", time.perf_counter() - start)
+        
+        if config.make_single_file:        
+            shutil.rmtree(folder_path)
 
         return None