From 5fbeceef78395bbfe808f43579949d680d64877b Mon Sep 17 00:00:00 2001 From: luiz Date: Thu, 22 Aug 2024 14:37:14 +0200 Subject: [PATCH] readm from nwb --- voluseg/_steps/step0.py | 83 ++++++++++++++++++++++++----------------- voluseg/_steps/step1.py | 53 +++++++++++++++++--------- voluseg/_tools/nwb.py | 50 +++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 53 deletions(-) create mode 100644 voluseg/_tools/nwb.py diff --git a/voluseg/_steps/step0.py b/voluseg/_steps/step0.py index f8e02f0..f81ea30 100644 --- a/voluseg/_steps/step0.py +++ b/voluseg/_steps/step0.py @@ -7,6 +7,7 @@ from voluseg._tools.get_volume_name import get_volume_name from voluseg._tools.parameter_dictionary import parameter_dictionary from voluseg._tools.evenly_parallelize import evenly_parallelize +from voluseg._tools.nwb import open_nwbfile_local def process_parameters(parameters0: dict) -> None: @@ -142,46 +143,58 @@ def process_parameters(parameters0: dict) -> None: volume_fullnames_input = [] volume_names = [] - for dir_input_h, dir_prefix_h in zip(input_dirs, prefix_dirs): - # get volume extension, volume names and number of segmentation timepoints - file_names = [i.split(".", 1) for i in os.listdir(dir_input_h) if "." in i] - file_exts, counts = np.unique(list(zip(*file_names))[1], return_counts=True) - ext = "." + file_exts[np.argmax(counts)] - volume_names_input_h = sorted([i for i, j in file_names if "." + j == ext]) - volume_fullnames_input_h = [ - os.path.join(dir_input_h, i) for i in volume_names_input_h - ] + if parameters.get("nwb_input_local_path", None): + acquisition_name = parameters.get("nwb_input_acquisition_name", None) + if not acquisition_name: + raise Exception("nwb_input_acquisition_name is required to read nwb file") + volume_fullnames_input = [parameters.get("nwb_input_local_path")] + volume_names = [parameters.get("nwb_input_local_path").split("/")[-1].split(".nwb")[0]] + with open_nwbfile_local(file_path=parameters["nwb_input_local_path"]) as nwbfile: + lt = nwbfile.acquisition[acquisition_name].data.shape[0] + if parameters["timepoints"]: + lt = min(lt, parameters["timepoints"]) + ext = ".nwb" + else: + for dir_input_h, dir_prefix_h in zip(input_dirs, prefix_dirs): + # get volume extension, volume names and number of segmentation timepoints + file_names = [i.split(".", 1) for i in os.listdir(dir_input_h) if "." in i] + file_exts, counts = np.unique(list(zip(*file_names))[1], return_counts=True) + ext = "." + file_exts[np.argmax(counts)] + volume_names_input_h = sorted([i for i, j in file_names if "." + j == ext]) + volume_fullnames_input_h = [ + os.path.join(dir_input_h, i) for i in volume_names_input_h + ] - # adjust parameters for packed planes data - if parameters["planes_packed"]: - parameters["res_z"] = parameters["diam_cell"] + # adjust parameters for packed planes data + if parameters["planes_packed"]: + parameters["res_z"] = parameters["diam_cell"] - def get_plane_names(tuple_fullname_volume_input): - fullname_volume_input = tuple_fullname_volume_input[1] - lp = len(load_volume(fullname_volume_input + ext)) - return [ - get_volume_name(fullname_volume_input, dir_prefix_h, pi) - for pi in range(lp) - ] + def get_plane_names(tuple_fullname_volume_input): + fullname_volume_input = tuple_fullname_volume_input[1] + lp = len(load_volume(fullname_volume_input + ext)) + return [ + get_volume_name(fullname_volume_input, dir_prefix_h, pi) + for pi in range(lp) + ] - volume_names_h = ( - evenly_parallelize(volume_fullnames_input_h) - .map(get_plane_names) - .collect() - ) - volume_names_h = [pi for ni in volume_names_h for pi in ni] - else: - volume_names_h = [ - get_volume_name(i, dir_prefix_h) for i in volume_names_input_h - ] + volume_names_h = ( + evenly_parallelize(volume_fullnames_input_h) + .map(get_plane_names) + .collect() + ) + volume_names_h = [pi for ni in volume_names_h for pi in ni] + else: + volume_names_h = [ + get_volume_name(i, dir_prefix_h) for i in volume_names_input_h + ] - # grow volume-name lists - volume_fullnames_input += volume_fullnames_input_h - volume_names += volume_names_h + # grow volume-name lists + volume_fullnames_input += volume_fullnames_input_h + volume_names += volume_names_h - volume_fullnames_input = np.array(volume_fullnames_input) - volume_names = np.array(volume_names) - lt = len(volume_names) + volume_fullnames_input = np.array(volume_fullnames_input) + volume_names = np.array(volume_names) + lt = len(volume_names) # check timepoints parameters["type_timepoints"] = parameters["type_timepoints"].lower() diff --git a/voluseg/_steps/step1.py b/voluseg/_steps/step1.py index 8af9580..b713236 100755 --- a/voluseg/_steps/step1.py +++ b/voluseg/_steps/step1.py @@ -7,6 +7,7 @@ from voluseg._tools.get_volume_name import get_volume_name from voluseg._tools.constants import ori, ali, nii, hdf from voluseg._tools.evenly_parallelize import evenly_parallelize +from voluseg._tools.nwb import open_nwbfile_local, get_nwbfile_volume def process_volumes(parameters: dict) -> None: @@ -26,7 +27,11 @@ def process_volumes(parameters: dict) -> None: p = SimpleNamespace(**parameters) - volume_fullname_inputRDD = evenly_parallelize(p.volume_fullnames_input) + if parameters.get("ext") == ".nwb": + volume_fullname_inputRDD = evenly_parallelize(range(p.lt)) + else: + volume_fullname_inputRDD = evenly_parallelize(p.volume_fullnames_input) + for color_i in range(p.n_colors): fullname_volmean = os.path.join(p.dir_output, "volume%d" % (color_i)) if os.path.isfile(fullname_volmean + hdf): @@ -114,25 +119,37 @@ def make_output_volume(name_volume, volume): # end make output volume # get full name of input volume, input data and list of planes - fullname_volume_input = tuple_fullname_volume_input[1] - volume = load_volume(fullname_volume_input + p.ext) - if len(p.input_dirs) == 1: - dir_prefix = None - else: - dir_prefix = os.path.basename(os.path.split(fullname_volume_input)[0]) - - # process output volumes - if p.planes_packed: - for pi, volume_pi in enumerate(volume): - name_volume_pi = get_volume_name( - fullname_volume_input, - dir_prefix, - pi, + if parameters.get("ext") == ".nwb": + with open_nwbfile_local(file_path=p.nwb_input_local_path) as nwbfile: + volume = get_nwbfile_volume( + nwbfile=nwbfile, + acquisition_name=p.nwb_input_acquisition_name, + volume_index=tuple_fullname_volume_input[1], ) - make_output_volume(name_volume_pi, volume_pi) + make_output_volume( + name_volume=p.volume_names[0] + f"_{tuple_fullname_volume_input[1]}", + volume=volume, + ) else: - name_volume = get_volume_name(fullname_volume_input, dir_prefix) - make_output_volume(name_volume, volume) + fullname_volume_input = tuple_fullname_volume_input[1] + volume = load_volume(fullname_volume_input + p.ext) + if len(p.input_dirs) == 1: + dir_prefix = None + else: + dir_prefix = os.path.basename(os.path.split(fullname_volume_input)[0]) + + # process output volumes + if p.planes_packed: + for pi, volume_pi in enumerate(volume): + name_volume_pi = get_volume_name( + fullname_volume_input, + dir_prefix, + pi, + ) + make_output_volume(name_volume_pi, volume_pi) + else: + name_volume = get_volume_name(fullname_volume_input, dir_prefix) + make_output_volume(name_volume, volume) # end initial_processing diff --git a/voluseg/_tools/nwb.py b/voluseg/_tools/nwb.py new file mode 100644 index 0000000..4e2144a --- /dev/null +++ b/voluseg/_tools/nwb.py @@ -0,0 +1,50 @@ +import pynwb +from contextlib import contextmanager + + +@contextmanager +def open_nwbfile_local(file_path: str): + """ + Context manager to open and close a local NWB file. + + Parameters + ---------- + file_path : str + Path to NWB file. + + Yields + ------ + pynwb.NWBFile + The opened NWB file. + """ + io = pynwb.NWBHDF5IO(file_path, 'r') + try: + nwbfile = io.read() + yield nwbfile + finally: + io.close() + + +def get_nwbfile_volume( + nwbfile: pynwb.NWBFile, + acquisition_name: str, + volume_index: int, +) -> pynwb.NWBFile: + """ + Get 3D volume from NWB file. + + Parameters + ---------- + nwbfile : pynwb.NWBFile + NWB file. + acquisition_name : str + Acquisition name. + volume_index : int + Volume index. + + Returns + ------- + pynwb.NWBFile + NWB file. + """ + return nwbfile.acquisition[acquisition_name].data[volume_index] \ No newline at end of file