save work for projection chunking

ct_chunking.py

+import zarr
+import numpy as np
+
+from cil.io import ZEISSDataReader, NikonDataReader
+from cil.framework import AcquisitionData
+from cil.processors import TransmissionAbsorptionConverter
+from cil.recon import FDK
+
+def create_reader(file_name, roi=None):
+    if file_name.endswith('txrm'):
+        DataReader = ZEISSDataReader
+    elif file_name.endswith('xtekct'):
+        DataReader = NikonDataReader
+    else:
+        raise ValueError("Unrecognizable CT metadata file. File extension should either be '.txrm' or '.xtekct'")
+    
+    if roi is None:
+        return DataReader(file_name=file_name)
+    else:
+        return DataReader(file_name=file_name, roi=roi)
+
+def write_projection_chunks(config_path, zarr_path, chunk_size, proj_shape):
+    """
+    # (ignore for now, meant for custom reader) input_path: Directory containing .tiff projection data.
+    out_path: Directory to write zarr file.
+    """
+    nchunks = (proj_shape[0]-1)//chunk_size + 1
+    compressors = zarr.codecs.Blosc(cname='lz4')
+    zarr_array = zarr.open(
+        store=zarr_path,
+        mode='w',
+        dtype=np.uint16,
+        chunks=(chunk_size, *proj_shape[1:]),
+        compressors=compressors)
+    
+    for i in range(nchunks):
+        start = i*chunk_size
+        end = min((i+1)*chunk_size, proj_shape[0])
+        reader = create_reader(config_path, roi={'angle': (start, end, 1)})
+        chunk = reader.read()
+
+        expected_shape = (end-start, *proj_shape[1:])
+        if chunk.shape != expected_shape:
+            raise ValueError(f'Chunk shape mismatch: expected {expected_shape}, got {chunk.shape}')
+
+        zarr_array[start:end,:,:] = chunk
+
+def read_projection_chunk(zarr_path, chunk_size, chunk_index):
+    zarr_array = zarr.open_array(store=zarr_path, mode='r')
+    proj_shape = zarr_array.shape
+    start = chunk_index*chunk_size
+    end = min((chunk_index+1)*chunk_size, proj_shape[0])
+    chunk = zarr_array[start:end,:,:]
+    return chunk
+
+def recon_projection_chunk(ag, ig, zarr_path, chunk_size, chunk_index):
+    """
+    ag: AcquisitionGeometry with chunked projection.
+    ig: ImageGeometry for the reconstruction volume.
+    """
+    start = chunk_index*chunk_size
+    end = min((chunk_index+1)*chunk_size, ag.shape[0])
+    ag_chunk = ag.copy().set_angles(ag.angles[start:end])
+    chunk = read_projection_chunk(zarr_path, chunk_size, chunk_index)
+    chunk = chunk.astype(np.float32)
+    data_chunk = AcquisitionData(array=chunk, deep_copy=False, geometry=ag_chunk)
+    data_chunk = TransmissionAbsorptionConverter()(data_chunk)
+    data_chunk.reorder(order='tigre')
+    recon_chunk = FDK(data_chunk, image_geometry=ig).run(verbose=1)
+    return recon_chunk
+
+def chunk_volumes():
+    pass
+
+def main():
+    path = '/dtu/3d-imaging-center/projects/2024_QIM_platform/analysis/data/valnut/valnut_2014-03-21_643_28/tomo-A/valnut_tomo-A.txrm'
+    chunk_size = 600
+    zarr_path = 'ct_chunk.zarr'
+
+    ag = create_reader(path).get_geometry()
+    ig = ag.get_ImageGeometry()
+    proj_shape = ag.shape
+
+    # get number of projections. then use CIL reader to read in subset of tiff files for each chunk and have a function to write to disk.
+    # potentially replace this with a custom reader
+    write_projection_chunks(config_path=path, zarr_path=zarr_path, chunk_size=chunk_size, proj_shape=proj_shape)
+    
+    nchunks = (proj_shape[0]-1)//chunk_size + 1
+    recon_accumulated = np.zeros(ig.shape, dtype=np.float32)
+    for chunk_index in range(nchunks):
+        recon_accumulated += recon_projection_chunk(ag, ig, zarr_path, chunk_size=chunk_size, chunk_index=chunk_index)
+