fixed further index problem with 5D data array

suite2p/io/h5.py

@@ -58,10 +58,9 @@ def h5py_to_binary(ops):
             for key in keys:
                 hdims = f[key].ndim
                 # keep track of the plane identity of the first frame (channel identity is assumed always 0)
-                ncp = nplanes * nchannels
+                ncp = nplanes * nchannels  # number of frame per 1 timestamp
                 nbatch = ncp * math.ceil(ops1[0]["batch_size"] / ncp)
-                nframes_all = f[key].shape[
-                    0] if hdims == 3 else f[key].shape[0] * f[key].shape[1]
+                nframes_all = np.prod(f[key].shape[:-2]) # last two dim= row x col
                 nbatch = min(nbatch, nframes_all)
                 nfunc = ops["functional_chan"] - 1 if nchannels > 1 else 0
                 # loop over all tiffs
@@ -78,13 +77,19 @@ def h5py_to_binary(ops):
                             1)
                         if irange.size == 0:
                             break
-                        im = f[key][irange, ...]
+                        if hdims == 4:
+                            im = f[key][irange, ...]
+                        elif hdims == 5:
+                            im = f[key][:, irange, :, :, :]
                         if im.ndim == 5 and im.shape[0] == nchannels:
                             im = im.transpose((1, 0, 2, 3, 4))
                         # flatten to frames x pixels x pixels
+                        # for 5D data with 2 channels, 2 planes, this results
+                        # [fr0 ch0 pl0], [fr0 ch0 pl1], [fr0 ch1 pl0],
+                        # [fr0 ch1 pl1] [fr1 ch0 pl0], ...
                         im = np.reshape(im, (-1, im.shape[-2], im.shape[-1]))
                     nframes = im.shape[0]
-                    if type(im[0, 0, 0]) == np.uint16:
+                    if im.dtype == np.uint16:
                         im = im / 2
                     for j in range(0, nplanes):
                         if iall == 0:
@@ -94,7 +99,7 @@ def h5py_to_binary(ops):
                                 ops1[j]["meanImg_chan2"] = np.zeros(
                                     (im.shape[1], im.shape[2]), np.float32)
                             ops1[j]["nframes"] = 0
-                        i0 = nchannels * ((j) % nplanes)
+                        i0 = j
                         im2write = im[np.arange(int(i0) +
                                                 nfunc, nframes, ncp), :, :].astype(
                                                     np.int16)