readm from nwb

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()

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
 

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]