Merge pull request #5 from tdincer/nd2support

Nd2support

Author: sidhulyalkar

docs/inputs.rst

@@ -148,6 +148,14 @@ Scanbox binary files (*.sbx) work out of the box if you set ``ops['input_format'
 When recording in bidirectional mode some columns might have every other line saturated; to trim these during loading set ``ops['sbx_ndeadcols']``. Set this option to ``-1`` to let suite2p compute the number of columns automatically, a positive integer to specify the number of columns to trim.
 Joao Couto (@jcouto) wrote the binary sbx parser.
 
+
+Nikon nd2 files
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Suite2p reads nd2 files using the nd2 package and returns a numpy array representing the data with a minimum of two dimensions (Height, Width). The data can also have additional dimensions for Time, Depth, and Channel. If any dimensions are missing, Suite2p adds them in the order of Time, Depth, Channel, Height, and Width, resulting in a 5-dimensional array. To use Suite2p with nd2 files, simply set ``ops['input_format'] = "nd2".``
+
+
+
 BinaryRWFile
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

setup.py

@@ -14,7 +14,8 @@
         'h5py',
         'sbxreader',
         'scikit-learn',
-        'cellpose']
+        'cellpose',
+        'nd2']
 
 gui_deps = [
         "pyqt5",

suite2p/io/nd2.py

@@ -8,7 +8,7 @@
 
 
 def nd2_to_binary(ops):
-    """  finds nd2 files and writes them to binaries
+    """finds nd2 files and writes them to binaries
 
     Parameters
     ----------
@@ -39,72 +39,83 @@ def nd2_to_binary(ops):
     for file_name in fs:
         # open nd2
         nd2_file = nd2.ND2File(file_name)
-        im = nd2_file.asarray()
+        nd2_dims = {k: i for i, k in enumerate(nd2_file.sizes)}
 
-        # expand dimensions to have [Time (T), Depth (Z), Channel (C), Height (Y), Width (X)].
-        if 'T' not in nd2_file.sizes:
-            im = np.expand_dims(im, 0)
-        if 'C' not in nd2_file.sizes:
-            im = np.expand_dims(im, -3)
-        if 'Z' not in nd2_file.sizes:
-            im = np.expand_dims(im, 1)
+        valid_dimensions = "TZCYX"
+        assert set(nd2_dims) <= set(
+            valid_dimensions
+        ), f"Unknown dimensions {set(nd2_dims)-set(valid_dimensions)} in file {file_name}."
 
-        nplanes = nd2_file.sizes['Z'] if 'Z' in nd2_file.sizes else 1
-        nchannels = nd2_file.sizes['C'] if 'C' in nd2_file.sizes else 1
-        nframes = nd2_file.sizes['T'] if 'T' in nd2_file.sizes else 1
+        # Sort the dimensions in the order of TZCYX, skipping the missing ones.
+        im = nd2_file.asarray().transpose(
+            [nd2_dims[x] for x in valid_dimensions if x in nd2_dims]
+        )
 
-        iblocks = np.arange(0, nframes, ops1[0]['batch_size'])
+        # Expand array to include the missing dimensions.
+        for i, dim in enumerate("TZC"):
+            if dim not in nd2_dims:
+                im = np.expand_dims(im, i)
+
+        nplanes = nd2_file.sizes["Z"] if "Z" in nd2_file.sizes else 1
+        nchannels = nd2_file.sizes["C"] if "C" in nd2_file.sizes else 1
+        nframes = nd2_file.sizes["T"] if "T" in nd2_file.sizes else 1
+
+        iblocks = np.arange(0, nframes, ops1[0]["batch_size"])
         if iblocks[-1] < nframes:
             iblocks = np.append(iblocks, nframes)
 
         if nchannels > 1:
-            nfunc = ops1[0]['functional_chan'] - 1
+            nfunc = ops1[0]["functional_chan"] - 1
         else:
             nfunc = 0
 
-        if im.dtype.type == np.uint16:
-            im = (im // 2).astype(np.int16)
-        elif im.dtype.type == np.int32:
-            im = (im // 2).astype(np.int16)
-        elif im.dtype.type != np.int16:
-            im = im.astype(np.int16)
+        assert im.max() < 32768 and im.min() >= -32768, "image data is out of range"
+        im = im.astype(np.int16)
 
         # loop over all frames
         for ichunk, onset in enumerate(iblocks[:-1]):
             offset = iblocks[ichunk + 1]
             im_p = np.array(im[onset:offset, :, :, :, :])
-            im2mean = im_p.mean(axis = 0).astype(np.float32) / len(iblocks)
+            im2mean = im_p.mean(axis=0).astype(np.float32) / len(iblocks)
             for ichan in range(nchannels):
                 nframes = im_p.shape[0]
                 im2write = im_p[:, :, ichan, :, :]
                 for j in range(0, nplanes):
                     if iall == 0:
-                        ops1[j]['meanImg'] = np.zeros((im_p.shape[3], im_p.shape[4]), np.float32)
-                        if nchannels>1:
-                            ops1[j]['meanImg_chan2'] = np.zeros((im_p.shape[3], im_p.shape[4]), np.float32)
-                        ops1[j]['nframes'] = 0
+                        ops1[j]["meanImg"] = np.zeros(
+                            (im_p.shape[3], im_p.shape[4]), np.float32
+                        )
+                        if nchannels > 1:
+                            ops1[j]["meanImg_chan2"] = np.zeros(
+                                (im_p.shape[3], im_p.shape[4]), np.float32
+                            )
+                        ops1[j]["nframes"] = 0
                     if ichan == nfunc:
-                        ops1[j]['meanImg'] += np.squeeze(im2mean[j, ichan, :, :])
-                        reg_file[j].write(bytearray(im2write[:, j, :, :].astype('int16')))
+                        ops1[j]["meanImg"] += np.squeeze(im2mean[j, ichan, :, :])
+                        reg_file[j].write(
+                            bytearray(im2write[:, j, :, :].astype("int16"))
+                        )
                     else:
-                        ops1[j]['meanImg_chan2'] += np.squeeze(im2mean[j, ichan, :, :])
-                        reg_file_chan2[j].write(bytearray(im2write[:, j, :, :].astype('int16')))
-                        
-                    ops1[j]['nframes'] += im2write.shape[0]
+                        ops1[j]["meanImg_chan2"] += np.squeeze(im2mean[j, ichan, :, :])
+                        reg_file_chan2[j].write(
+                            bytearray(im2write[:, j, :, :].astype("int16"))
+                        )
+
+                    ops1[j]["nframes"] += im2write.shape[0]
             ik += nframes
             iall += nframes
-        
+
         nd2_file.close()
 
     # write ops files
-    do_registration = ops1[0]['do_registration']
+    do_registration = ops1[0]["do_registration"]
     for ops in ops1:
-        ops['Ly'] = im.shape[3]
-        ops['Lx'] = im.shape[4]
+        ops["Ly"] = im.shape[3]
+        ops["Lx"] = im.shape[4]
         if not do_registration:
-            ops['yrange'] = np.array([0, ops['Ly']])
-            ops['xrange'] = np.array([0, ops['Lx']])
-        np.save(ops['ops_path'], ops)
+            ops["yrange"] = np.array([0, ops["Ly"]])
+            ops["xrange"] = np.array([0, ops["Lx"]])
+        np.save(ops["ops_path"], ops)
     # close all binary files and write ops files
     for j in range(0, nplanes):
         reg_file[j].close()