ct/chunking.py

@@ -47,7 +47,7 @@ class ChunkedCT:
         output_dir,
         angle_chunks,
         vol_chunks,
-        read_batch_size=20,
+        read_batch_size=None,
         processes=None,
         verbose=None,
         use_pickle_geometries=False
@@ -62,7 +62,7 @@ class ChunkedCT:
             tiff_dir: Directory where all TIFFs are assumed to be projections. Defaults to parent of config_path.
             angle_chunks: How many chunks should the projections be chunked along the angle dimension.
             vol_chunks: Tuple of number of chunks along each direction in the order of Z, Y, X. Must be divisors of the corresponding default ig dimension length.
-            read_batch_size: Number of TIFFs per worker at a time during parallelized reading.
+            read_batch_size: Number of TIFFs per worker at a time during parallelized reading. Default 20.
             processes: Number of workers during parallelized reading. Default None chooses this to all available CPUs.
             verbose: If slices from the reconstructions should be saved along the way to see the progress visually. Default to 1.
             use_pickle_geometries: Used to pass custom geometries. It is assumed the geometries are saved as 'ag.pkl' and 'ig.pkl' inside of tiff_dir.
@@ -85,7 +85,7 @@ class ChunkedCT:
         self.angle_chunks = angle_chunks
         self.vol_chunks = vol_chunks
         
-        self.read_batch_size = read_batch_size
+        self.read_batch_size = read_batch_size if read_batch_size is not None else 20
         self.processes = processes
         
         os.makedirs(self.output_dir, exist_ok=True)
@@ -187,15 +187,16 @@ class ChunkedCT:
     def make_ig_chunk(self, chunk_coords):
         ig_chunk = self.ig.copy()
         for axis, dim_label in enumerate(['z', 'y', 'x']):
-            setattr(ig_chunk, f'voxel_num_{dim_label}', self.vol_chunk_shape[axis])
             chunk = chunk_coords[axis]
-            num_chunks = self.vol_chunks[axis]
+            voxel_num = min(self.vol_chunk_shape[axis], self.ig.shape[axis] - chunk*self.vol_chunk_shape[axis])
+            setattr(ig_chunk, f'voxel_num_{dim_label}', voxel_num)
             ig_voxel_num = getattr(self.ig, f'voxel_num_{dim_label}')
             ig_chunk_voxel_num = getattr(ig_chunk, f'voxel_num_{dim_label}')
             ig_chunk_voxel_size = getattr(ig_chunk, f'voxel_size_{dim_label}')
             
             # center is in 'world' coordinates, hence the scaling with voxel_size
-            center = (-(1 - 1/num_chunks)/2 * ig_voxel_num + chunk * ig_chunk_voxel_num) * ig_chunk_voxel_size
+            initial_shift = (-ig_voxel_num/2 + self.vol_chunk_shape[axis]/2) * ig_chunk_voxel_size
+            center = initial_shift + (chunk*self.vol_chunk_shape[axis] - (self.vol_chunk_shape[axis] - ig_chunk_voxel_num)/2) * ig_chunk_voxel_size
             setattr(ig_chunk, f'center_{dim_label}', center)
         
         return ig_chunk
@@ -207,11 +208,12 @@ class ChunkedCT:
         for proj_chunk_index in range(self.angle_chunks):
             recon_chunk = self.reconstruct_projection_chunk(ig_chunk, proj_chunk_index)
             vol_accumulated += recon_chunk.as_array()
-            if self.verbose:
+            if self.verbose > 0:
                 str_chunk = '_'.join(str(i) for i in chunk_coords)
+            if self.verbose > 1:
                 save_center_slice(recon_chunk.as_array(), self.output_dir / f'center/recon_vol_{str_chunk}_proj_{proj_chunk_index}.png')
 
-        if self.verbose:
+        if self.verbose > 0:
             save_center_slice(vol_accumulated, self.output_dir / f'center/recon_{str_chunk}_accumulated.png')
         
         recon_zarr = zarr.open_array(store=self.recon_zarr_path, mode='a')
@@ -225,7 +227,7 @@ class ChunkedCT:
             store=self.recon_zarr_path,
             mode='w',
             dtype=np.float32,
-            shape=self.ig.shape[::-1],
+            shape=self.ig.shape,
             chunks=self.vol_chunk_shape,
             compressor=None
         )
@@ -241,6 +243,7 @@ def main(
     output_dir: str='out',
     angle_chunks: int=1,
     vol_chunks: str='1,1,1',
+    read_batch_size: int=None,
     processes: str=None,
     verbose: int=None,
     use_pickle_geometries: bool=False,
@@ -262,6 +265,12 @@ def main(
         Number of processes for multiprocessing. None uses all available.
     verbose: int
     """
+    print('-' * 60)
+    print("Performing chunked CT with the configurations:")
+    for k, v in locals().items():
+        print(f'  {k}: {v}')
+    print('-' * 60 + '\n')
+    
     vol_chunks = tuple(int(d) for d in vol_chunks.split(','))
     if len(vol_chunks) != 3:
         raise ValueError
@@ -276,6 +285,7 @@ def main(
         output_dir=output_dir,
         angle_chunks=angle_chunks,
         vol_chunks=vol_chunks,
+        read_batch_size=read_batch_size,
         processes=processes,
         verbose=verbose,
         use_pickle_geometries=use_pickle_geometries,