From a6240dbb09c5722b770f731c5bb51b6b6af640fb Mon Sep 17 00:00:00 2001 From: Weaver Goldman Date: Wed, 7 Aug 2024 10:35:27 -0400 Subject: [PATCH] Add support for supersampling backprojection to match mip levels. --- python/ouroboros/common/logging.py | 11 +- python/ouroboros/common/server_handlers.py | 10 +- python/ouroboros/helpers/options.py | 2 +- .../pipeline/backproject_pipeline.py | 211 +++++++++++++++++- src/renderer/src/interfaces/options.tsx | 1 + 5 files changed, 219 insertions(+), 16 deletions(-) diff --git a/python/ouroboros/common/logging.py b/python/ouroboros/common/logging.py index d1ca2bd..53e6f48 100644 --- a/python/ouroboros/common/logging.py +++ b/python/ouroboros/common/logging.py @@ -2,12 +2,5 @@ server_logger = logging.getLogger("uvicorn") -DEFAULT_LOGGER = server_logger - - -def get_logger(): - return DEFAULT_LOGGER - - -def get_server_logger(): - return server_logger +# Primary export, import this from other modules +logger = server_logger diff --git a/python/ouroboros/common/server_handlers.py b/python/ouroboros/common/server_handlers.py index 29af096..28e2f90 100644 --- a/python/ouroboros/common/server_handlers.py +++ b/python/ouroboros/common/server_handlers.py @@ -6,7 +6,7 @@ save_output_for_backproject_docker, save_output_for_slice_docker, ) -from ouroboros.common.logging import get_logger +from ouroboros.common.logging import logger from ouroboros.common.pipelines import backproject_pipeline, slice_pipeline from ouroboros.common.server_types import BackProjectTask, SliceTask, Task @@ -28,8 +28,8 @@ def handle_slice_core(task: SliceTask, slice_options): return error # Log the pipeline statistics - get_logger().info("Slice Pipeline Statistics:") - get_logger().info(pipeline.get_step_statistics()) + logger.info("Slice Pipeline Statistics:") + logger.info(pipeline.get_step_statistics()) def handle_slice(task: SliceTask): @@ -88,8 +88,8 @@ def handle_backproject_core(task: BackProjectTask, options): return error # Log the pipeline statistics - get_logger().info("Backproject Pipeline Statistics:") - get_logger().info(pipeline.get_step_statistics()) + logger.info("Backproject Pipeline Statistics:") + logger.info(pipeline.get_step_statistics()) def handle_backproject(task: BackProjectTask): diff --git a/python/ouroboros/helpers/options.py b/python/ouroboros/helpers/options.py index 55386a8..3600ded 100644 --- a/python/ouroboros/helpers/options.py +++ b/python/ouroboros/helpers/options.py @@ -12,6 +12,7 @@ class CommonOptions(BaseModel): flush_cache: bool = False # Whether to flush the cache after processing make_single_file: bool = True # Whether to save the output to a single file max_ram_gb: int = 0 # Maximum amount of RAM to use in GB (0 means no limit) + output_mip_level: int = 0 # MIP level for the output image layer @model_with_json @@ -33,7 +34,6 @@ class SliceOptions(CommonOptions): False # Whether to connect the start and end of the given annotation points ) annotation_mip_level: int = 0 # MIP level for the annotation layer - output_mip_level: int = 0 # MIP level for the output image layer @field_serializer("bounding_box_params") def serialize_bounding_box_params(self, value: BoundingBoxParams): diff --git a/python/ouroboros/pipeline/backproject_pipeline.py b/python/ouroboros/pipeline/backproject_pipeline.py index ae3df4d..a305958 100644 --- a/python/ouroboros/pipeline/backproject_pipeline.py +++ b/python/ouroboros/pipeline/backproject_pipeline.py @@ -1,3 +1,5 @@ +import scipy +import scipy.ndimage from ouroboros.helpers.memory_usage import GIGABYTE, calculate_gigabytes_from_dimensions from ouroboros.helpers.slice import ( detect_color_channels, @@ -5,7 +7,7 @@ make_volume_binary, write_slices_to_volume, ) -from ouroboros.helpers.volume_cache import VolumeCache +from ouroboros.helpers.volume_cache import VolumeCache, get_mip_volume_sizes from ouroboros.helpers.bounding_boxes import BoundingBox from .pipeline import PipelineStep from ouroboros.helpers.options import BackprojectOptions @@ -346,6 +348,37 @@ def _process(self, input_data: any) -> tuple[any, None] | tuple[None, any]: except BaseException as e: return f"Error creating single tif file: {e}" + # Rescale the backprojected volume to the output mip level + if pipeline_input.slice_options.output_mip_level != config.output_mip_level: + output_name = f"{folder_path}-temp" + + error = rescale_mip_volume( + pipeline_input.source_url, + pipeline_input.slice_options.output_mip_level, + config.output_mip_level, + single_path=( + None if config.make_single_file is False else folder_path + ".tif" + ), + folder_path=(folder_path if config.make_single_file is False else None), + output_name=output_name, + compression=config.backprojection_compression, + ) + + if error is not None: + return error + + # Remove the original backprojected volume + if config.make_single_file: + os.remove(folder_path + ".tif") + else: + shutil.rmtree(folder_path) + + # Rename the rescaled volume + if config.make_single_file: + os.rename(output_name + ".tif", folder_path + ".tif") + else: + os.rename(output_name, folder_path) + # Update the pipeline input with the output file path pipeline_input.backprojected_folder_path = folder_path @@ -507,3 +540,179 @@ def create_volume_chunks( chunks_and_boxes.append((bounding_box, chunk_boxes)) return chunks_and_boxes + + +def rescale_mip_volume( + source_url, + current_mip, + target_mip, + single_path=None, + folder_path=None, + output_name="out", + compression=None, +) -> str | None: + """ + Rescale the volume to the mip level. + + Parameters + ---------- + source_url : str + The URL of the volume. + current_mip : int + The current mip level of the volume. + target_mip : int + The target mip level of the volume. + single_path : str + The path to the single tif file. + folder_path : str + The path to the folder containing the tif files. + compression : str, optional + The compression to use for the resulting tif file. + The default is None. + + Returns + ------- + str | None + Error message if an error occurred. + """ + + if single_path is None and folder_path is None: + return "Either single_path or folder_path must be provided." + + if target_mip == current_mip: + return None + + if single_path is not None: + return rescale_single_tif( + source_url, + current_mip, + target_mip, + single_path, + compression=compression, + file_name=output_name + ".tif", + ) + + return rescale_folder_tif( + source_url, + current_mip, + target_mip, + folder_path, + compression=compression, + folder_name=output_name, + ) + + +def rescale_single_tif( + source_url, + current_mip, + target_mip, + single_path, + file_name="out.tif", + compression=None, +): + with tifffile.TiffFile(single_path) as tif: + tif_shape = (len(tif.pages),) + tif.pages[0].shape + + scaling_factors, _ = calculate_scaling_factors( + source_url, current_mip, target_mip, tif_shape + ) + + with tifffile.TiffWriter(file_name) as output_volume: + for i in range(tif_shape[0]): + tif_layer = tif.pages[i].asarray() + + # Give the tif a new 0th dimension so that it can be resized along all axes + tif_layer = np.expand_dims(tif_layer, axis=0) + + layers = scipy.ndimage.zoom(tif_layer, scaling_factors, order=3) + + size = layers.shape[0] + + # Save the layers to the tif file + for j in range(size): + output_volume.write( + layers[j], + contiguous=compression is None or compression == "none", + compression=compression, + software="ouroboros", + ) + + return None + + +def rescale_folder_tif( + source_url, + current_mip, + target_mip, + folder_path, + folder_name="out", + compression=None, +): + # Create output folder if it doesn't exist + output_folder = folder_name + os.makedirs(output_folder, exist_ok=True) + + tifs = get_sorted_tif_files(folder_path) + + if len(tifs) == 0: + return "No tif files found in the folder." + + # Determine the shape of the tif stack + new_shape = (len(tifs), *tifffile.imread(join_path(folder_path, tifs[0])).shape) + + scaling_factors, resolution_factors = calculate_scaling_factors( + source_url, current_mip, target_mip, new_shape + ) + + num_digits = len(str(len(tifs))) + + first_index = int(tifs[0].split(".")[0]) + + output_index = int(first_index * resolution_factors[0]) + + # Resize the volume + for i in range(len(tifs)): + tif = tifffile.imread(join_path(folder_path, tifs[i])) + + # Give the tif a new 0th dimension so that it can be resized along all axes + tif = np.expand_dims(tif, axis=0) + + layers = scipy.ndimage.zoom(tif, scaling_factors, order=3) + + size = layers.shape[0] + + # Write the layers to new tif files + for j in range(size): + tifffile.imwrite( + join_path(output_folder, f"{str(output_index).zfill(num_digits)}.tif"), + layers[j], + contiguous=True if compression is None else False, + compression=compression, + ) + output_index += 1 + + return None + + +def calculate_scaling_factors(source_url, current_mip, target_mip, tif_shape): + # Determine the current and target resolutions + mip_sizes = get_mip_volume_sizes(source_url) + + current_resolution = mip_sizes[current_mip] + target_resolution = mip_sizes[target_mip] + + # Determine the scaling factor for each axis as a tuple + resolution_factors = tuple( + max(target_resolution[i] / current_resolution[i], 1) + for i in range(len(target_resolution)) + ) + + has_color_channels = len(tif_shape) == 4 + num_channels = tif_shape[-1] if has_color_channels else 1 + + # Determine the scaling factor for each axis as a tuple + scaling_factors = resolution_factors + ( + (num_channels,) if has_color_channels else () + ) + + return scaling_factors, resolution_factors diff --git a/src/renderer/src/interfaces/options.tsx b/src/renderer/src/interfaces/options.tsx index a768f36..4096531 100644 --- a/src/renderer/src/interfaces/options.tsx +++ b/src/renderer/src/interfaces/options.tsx @@ -196,6 +196,7 @@ export class BackprojectOptionsFile extends CompoundEntry { new Entry('config_path', 'Slice Configuration File', '', 'filePath'), new Entry('output_file_folder', 'Output File Folder', './', 'filePath'), new Entry('output_file_name', 'Output File Name', 'sample', 'string'), + new Entry('output_mip_level', 'Output MIP Level', 0, 'number'), new Entry('backprojection_compression', 'Backprojection Compression', 'zlib', 'string'), new Entry('make_single_file', 'Output Single File', false, 'boolean'), new Entry('backproject_min_bounding_box', 'Output Min Bounding Box', true, 'boolean'),