diff --git a/src/correct_images/correct_images.py b/src/correct_images/correct_images.py index 5912b16b..9de52308 100644 --- a/src/correct_images/correct_images.py +++ b/src/correct_images/correct_images.py @@ -89,6 +89,19 @@ def main(args=None): default="", help="Expected suffix for correct_images configuration and output folders.", ) + subparser_parse.add_argument( + "--memmap-location", + dest="memmap_location", + default=".", + help="Location for memmap files. Defaults to current directory.", + ) + subparser_parse.add_argument( + "--memmap-size-gb", + dest="memmap_size_gb", + default=50, + type=int, + help="Size of individual memmap files in GB. Defaults to 50.", + ) subparser_parse.set_defaults(func=call_parse) # subparser process @@ -244,7 +257,12 @@ def call_parse(args): "parse", args.force, args.suffix, camera, correct_config=None ) # call new list-compatible implementation of parse() - corrector.parse(path_list, correct_config_list) + corrector.parse( + path_list, + correct_config_list, + args.memmap_location, + args.memmap_size_gb, + ) corrector.cleanup() Console.info( diff --git a/src/correct_images/corrector.py b/src/correct_images/corrector.py index 0ee25730..f7abbc41 100644 --- a/src/correct_images/corrector.py +++ b/src/correct_images/corrector.py @@ -276,7 +276,11 @@ def load_configuration(self, correct_config=None): else: self.loader.set_loader("default") - def parse(self, path_list, correct_config_list): + def parse( + self, path_list, correct_config_list, memmap_location=".", memmap_size_gb=50.0 + ): + self.memmap_location = memmap_location + self.memmap_size_gb = memmap_size_gb # both path_list and correct_config_list are assumed to be valid + equivalent for i in range(len(path_list)): # for each dive path = path_list[i] @@ -669,7 +673,9 @@ def get_altitude_and_depth_maps(self): return elif self.distance_metric == "depth_map": # Load depth maps - path_depth = self.path_processed/"3d_reconstruction/depth_maps/gp_smoothed" + path_depth = ( + self.path_processed / "3d_reconstruction/depth_maps/gp_smoothed" + ) if not path_depth.exists(): Console.quit("Depth maps folder", path_depth, "does not exist.") images_to_drop = [] @@ -747,8 +753,7 @@ def generate_attenuation_correction_parameters(self): * 4.0 / (1024.0**3) ) - max_bin_size_gb = 50.0 - max_bin_size = int(max_bin_size_gb / image_size_gb) + max_bin_size = int(self.memmap_size_gb / image_size_gb) self.bin_band = 0.1 hist_bins = np.arange( @@ -831,7 +836,7 @@ def generate_attenuation_correction_parameters(self): images_map, distances_map, max_bin_size, - max_bin_size_gb, + self.memmap_size_gb, distance_vector, ) for idx_bin in range(hist_bins.size - 1) @@ -1143,6 +1148,7 @@ def compute_distance_bin( memmap_filename, memmap_handle = create_memmap( bin_images, dimensions, + self.memmap_location, loader=self.loader, ) self.memmaps_to_remove.append(memmap_filename) @@ -1153,6 +1159,7 @@ def compute_distance_bin( memmap_filename, memmap_handle = create_memmap( bin_images, dimensions, + self.memmap_location, loader=self.loader, ) self.memmaps_to_remove.append(memmap_filename) @@ -1163,6 +1170,7 @@ def compute_distance_bin( memmap_filename, memmap_handle = create_memmap( bin_images, dimensions, + self.memmap_location, loader=self.loader, ) self.memmaps_to_remove.append(memmap_filename) diff --git a/src/correct_images/tools/memmap.py b/src/correct_images/tools/memmap.py index 48e72d65..73ccf557 100644 --- a/src/correct_images/tools/memmap.py +++ b/src/correct_images/tools/memmap.py @@ -8,6 +8,7 @@ import uuid from datetime import datetime +from pathlib import Path import cv2 import numpy as np @@ -17,7 +18,7 @@ from ..loaders import default -def create_memmap_name() -> str: +def create_memmap_name(memmap_location) -> str: filename_map = ( "memmap_" + datetime.now().strftime("%Y%m%d_%H%M%S_") @@ -27,11 +28,15 @@ def create_memmap_name() -> str: + str(uuid.uuid4()) + ".map" ) - return filename_map + memmap_location = Path(memmap_location) + if not memmap_location.exists(): + memmap_location.mkdir(parents=True) + filename_map = memmap_location / filename_map + return str(filename_map) -def create_memmap(image_list, dimensions, loader=default.loader): - filename_map = create_memmap_name() +def create_memmap(image_list, dimensions, memmap_location, loader=default.loader): + filename_map = create_memmap_name(memmap_location) Console.info("Creating memmap at", filename_map) # If only 1 channel, do not create a 3D array if dimensions[-1] == 1: @@ -55,15 +60,15 @@ def create_memmap(image_list, dimensions, loader=default.loader): return filename_map, image_memmap -def open_memmap(shape, dtype): - filename_map = create_memmap_name() +def open_memmap(shape, dtype, memmap_location): + filename_map = create_memmap_name(memmap_location) Console.info("Creating memmap (open_memmap) at", filename_map) image_memmap = np.memmap(filename=filename_map, mode="w+", shape=shape, dtype=dtype) return filename_map, image_memmap -def convert_to_memmap(array, loader=default.loader): - filename_map = create_memmap_name() +def convert_to_memmap(array, memmap_location, loader=default.loader): + filename_map = create_memmap_name(memmap_location) Console.info("Creating memmap (convert_to_memmap) at", filename_map) image_memmap = np.memmap( filename=filename_map, mode="w+", shape=array.shape, dtype=array.dtype