Merge pull request #1050 from MouseLand/newio

Newio

setup.py

@@ -25,7 +25,8 @@
     "paramiko",
     "nd2",
     "sbxreader",
-    "h5py"
+    "h5py",
+    "opencv-python-headless"
 ]
 
 nwb_deps = [

suite2p/detection/detect.py

@@ -125,6 +125,8 @@ def detection_wrapper(f_reg, mov=None, yrange=None, xrange=None, ops=default_ops
     n_frames, Ly, Lx = f_reg.shape
     yrange = ops.get("yrange", [0, Ly]) if yrange is None else yrange
     xrange = ops.get("xrange", [0, Lx]) if xrange is None else xrange
+    ops["yrange"] = yrange
+    ops["xrange"] = xrange
 
     if mov is None:
         bin_size = int(
@@ -138,6 +140,10 @@ def detection_wrapper(f_reg, mov=None, yrange=None, xrange=None, ops=default_ops
         elif mov.shape[2] != xrange[-1] - xrange[0]:
             raise ValueError("mov.shape[2] is not same size as xrange")
 
+    if "meanImg" not in ops:
+        ops["meanImg"] = mov.mean(axis=0)
+        ops["max_proj"] = mov.max(axis=0)
+
     if ops.get("inverted_activity", False):
         mov -= mov.min()
         mov *= -1

suite2p/gui/reggui.py

@@ -13,6 +13,7 @@
 from scipy.ndimage import gaussian_filter1d
 from natsort import natsorted
 from tifffile import imread
+import json
 
 from . import masks, views, graphics, traces, classgui, utils
 from .. import registration
@@ -354,8 +355,7 @@ def plot_trace(self):
         self.p2.setLimits(xMin=0, xMax=self.nframes)
 
     def open(self):
-        filename = QFileDialog.getOpenFileName(self, "Open single-plane ops.npy file",
-                                               filter="ops*.npy")
+        filename = QFileDialog.getOpenFileName(self, "Open single-plane ops.npy file or single-plane ops.json file")
         # load ops in same folder
         if filename:
             print(filename[0])
@@ -430,7 +430,19 @@ def openCombined(self, save_folder):
 
     def openFile(self, filename, fromgui):
         try:
-            ops = np.load(filename, allow_pickle=True).item()
+            ext = os.path.splitext(filename)[1]
+            if ext == ".npy":
+                ops = np.load(filename, allow_pickle=True).item()
+                dirname = os.path.dirname(filename)
+            elif ext == ".json":
+                with open(filename, "r") as f:
+                    ops = json.load(f)
+                ops["Ly"] = ops["Lys"] if isinstance(ops["Lys"], int) else ops["Lys"][0]
+                ops["Lx"] = ops["Lxs"] if isinstance(ops["Lxs"], int) else ops["Lxs"][0]
+                dirname = os.path.join(os.path.dirname(filename), "suite2p/plane0/")
+                ops["reg_file"] = os.path.join(dirname, "data.bin")
+                nbytesread = np.int64(2 * ops["Ly"] * ops["Lx"])
+                ops["nframes"] = os.path.getsize(ops["reg_file"]) // nbytesread
             self.LY = ops["Ly"]
             self.LX = ops["Lx"]
             self.Ly = [ops["Ly"]]
@@ -442,7 +454,7 @@ def openFile(self, filename, fromgui):
                 self.reg_loc = [ops["reg_file"]]
             else:
                 self.reg_loc = [
-                    os.path.abspath(os.path.join(os.path.dirname(filename), "data.bin"))
+                    os.path.abspath(os.path.join(dirname, "data.bin"))
                 ]
             self.reg_file = [open(self.reg_loc[-1], "rb")]
             self.wraw = False

suite2p/gui/rungui.py

@@ -276,7 +276,7 @@ def create_buttons(self):
         self.inputformat = QComboBox()
         [
             self.inputformat.addItem(f)
-            for f in ["tif", "bruker", "sbx", "h5", "mesoscan", "haus", "nd2"]
+            for f in ["tif", "binary", "bruker", "sbx", "h5", "movie", "nd2", "mesoscan", "haus"]
         ]
         self.inputformat.currentTextChanged.connect(self.parse_inputformat)
         self.layout.addWidget(self.inputformat, 2, 0, 1, 1)

suite2p/io/__init__.py

@@ -5,6 +5,7 @@
 from .nwb import save_nwb, read_nwb, nwb_to_binary
 from .save import combined, compute_dydx, save_mat
 from .sbx import sbx_to_binary
+from .movie import movie_to_binary
 from .tiff import mesoscan_to_binary, ome_to_binary, tiff_to_binary, generate_tiff_filename, save_tiff
 from .nd2 import nd2_to_binary
 from .binary import BinaryFile, BinaryFileCombined

suite2p/io/movie.py

@@ -0,0 +1,210 @@
+try:
+    import cv2
+    HAS_CV2 = True 
+except:
+    HAS_CV2 = False 
+
+import numpy as np
+import time
+from typing import Optional, Tuple, Sequence
+from .utils import find_files_open_binaries, init_ops
+
+class VideoReader:
+    """ Uses cv2 to read video files """
+    def __init__(self, filenames: list):
+        """ Uses cv2 to open video files and obtain their details for reading
+
+        Parameters
+        ------------
+        filenames : int
+            list of video files
+        """
+        cumframes = [0]
+        containers = []
+        Ly = []
+        Lx = []
+        for f in filenames:  # for each video in the list
+            cap = cv2.VideoCapture(f)
+            containers.append(cap)
+            Lx.append(int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)))
+            Ly.append(int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+            cumframes.append(cumframes[-1] + int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))
+        cumframes = np.array(cumframes).astype(int)
+        Ly = np.array(Ly)
+        Lx = np.array(Lx)
+        if (Ly==Ly[0]).sum() < len(Ly) or (Lx==Lx[0]).sum() < len(Lx):
+            raise ValueError("videos are not all the same size in y and x")
+        else:
+            Ly, Lx = Ly[0], Lx[0]
+
+        self.filenames = filenames
+        self.cumframes = cumframes 
+        self.n_frames = cumframes[-1]
+        self.Ly = Ly
+        self.Lx = Lx
+        self.containers = containers
+        self.fs = containers[0].get(cv2.CAP_PROP_FPS)
+
+    def close(self) -> None:
+        """
+        Closes the video files
+        """
+        for i in range(len(self.containers)):  # for each video in the list
+            cap = self.containers[i]
+            cap.release()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    @property
+    def shape(self) -> Tuple[int, int, int]:
+        """
+        The dimensions of the data in the file
+
+        Returns
+        -------
+        n_frames: int
+            The number of frames
+        Ly: int
+            The height of each frame
+        Lx: int
+            The width of each frame
+        """
+        return self.n_frames, self.Ly, self.Lx
+
+    def get_frames(self, cframes):
+        """
+        read frames "cframes" from videos
+
+        Parameters
+        ------------
+        cframes : np.array
+            start and stop of frames to read, or consecutive list of frames to read
+        """
+        cframes = np.maximum(0, np.minimum(self.n_frames - 1, cframes))
+        cframes = np.arange(cframes[0], cframes[-1] + 1).astype(int)
+        # find which video the frames exist in (ivids is length of cframes)
+        ivids = (cframes[np.newaxis, :] >= self.cumframes[1:, np.newaxis]).sum(axis=0)
+        nk = 0
+        im = np.zeros((len(cframes), self.Ly, self.Lx), "uint8")
+        for n in np.unique(ivids):  # for each video in cumframes
+            cfr = cframes[ivids == n]
+            start = cfr[0] - self.cumframes[n]
+            end = cfr[-1] - self.cumframes[n] + 1
+            nt0 = end - start
+            capture = self.containers[n]
+            if int(capture.get(cv2.CAP_PROP_POS_FRAMES)) != start:
+                capture.set(cv2.CAP_PROP_POS_FRAMES, start)
+            fc = 0
+            ret = True
+            while fc < nt0 and ret:
+                ret, frame = capture.read()
+                if ret:
+                    im[nk + fc] = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+                else:
+                    print("img load failed, replacing with prev..")
+                    im[nk + fc] = im[nk + fc - 1]
+                fc += 1
+            nk += nt0
+        return im
+
+def movie_to_binary(ops):
+    """  finds movie files and writes them to binaries
+
+    Parameters
+    ----------
+    ops : dictionary
+        "nplanes", "data_path", "save_path", "save_folder", "fast_disk",
+        "nchannels", "keep_movie_raw", "look_one_level_down" (optional: "subfolders")
+
+    Returns
+    -------
+        ops : dictionary of first plane
+            "Ly", "Lx", ops["reg_file"] or ops["raw_file"] is created binary
+
+    """
+    if not HAS_CV2:
+        raise ImportError("cv2 is required for this file type, please 'pip install opencv-python-headless'")
+
+    ops1 = init_ops(ops)
+
+    nplanes = ops1[0]["nplanes"]
+    nchannels = ops1[0]["nchannels"]
+
+    # open all binary files for writing
+    ops1, filenames, reg_file, reg_file_chan2 = find_files_open_binaries(ops1)
+
+    ik = 0
+    for j in range(ops["nplanes"]):
+        ops1[j]["nframes_per_folder"] = np.zeros(len(filenames), np.int32)
+
+
+    ncp = nplanes * nchannels
+    nbatch = ncp * int(np.ceil(ops1[0]["batch_size"] / ncp))
+    print(filenames)
+    t0 = time.time()
+    with VideoReader(filenames=filenames) as vr:
+        if ops1[0]["fs"]<=0:
+            for ops in ops1:
+                ops["fs"] = vr.fs
+
+        nframes_all = vr.cumframes[-1]
+        nbatch = min(nbatch, nframes_all)
+        nfunc = ops["functional_chan"] - 1 if nchannels > 1 else 0
+        # loop over all video frames
+        ik = 0
+        while 1:
+            irange = np.arange(ik, min(ik + nbatch, nframes_all), 1)
+            if irange.size == 0:
+                break
+            im = vr.get_frames(irange).astype("int16")
+            nframes = im.shape[0]
+            for j in range(0, nplanes):
+                if ik == 0:
+                    ops1[j]["meanImg"] = np.zeros((im.shape[1], im.shape[2]),
+                                                    np.float32)
+                    if nchannels > 1:
+                        ops1[j]["meanImg_chan2"] = np.zeros(
+                            (im.shape[1], im.shape[2]), np.float32)
+                    ops1[j]["nframes"] = 0
+                i0 = nchannels * ((j) % nplanes)
+                im2write = im[np.arange(int(i0) +
+                                        nfunc, nframes, ncp), :, :].astype(
+                                            np.int16)
+                reg_file[j].write(bytearray(im2write))
+                ops1[j]["meanImg"] += im2write.astype(np.float32).sum(axis=0)
+                if nchannels > 1:
+                    im2write = im[np.arange(int(i0) + 1 -
+                                            nfunc, nframes, ncp), :, :].astype(
+                                                np.int16)
+                    reg_file_chan2[j].write(bytearray(im2write))
+                    ops1[j]["meanImg_chan2"] += im2write.astype(
+                        np.float32).sum(axis=0)
+                ops1[j]["nframes"] += im2write.shape[0]
+                #ops1[j]["nframes_per_folder"][ih5] += im2write.shape[0]
+            ik += nframes
+            if ik % (nbatch * 4) == 0:
+                print("%d frames of binary, time %0.2f sec." %
+                      (ik, time.time() - t0))
+
+    # write ops files
+    do_registration = ops1[0]["do_registration"]
+    for ops in ops1:
+        ops["Ly"] = im2write.shape[1]
+        ops["Lx"] = im2write.shape[2]
+        if not do_registration:
+            ops["yrange"] = np.array([0, ops["Ly"]])
+            ops["xrange"] = np.array([0, ops["Lx"]])
+        ops["meanImg"] /= ops["nframes"]
+        if nchannels > 1:
+            ops["meanImg_chan2"] /= ops["nframes"]
+        np.save(ops["ops_path"], ops)
+    # close all binary files and write ops files
+    for j in range(nplanes):
+        reg_file[j].close()
+        if nchannels > 1:
+            reg_file_chan2[j].close()
+    return ops1[0]

suite2p/io/sbx.py

@@ -9,12 +9,10 @@
 
 try:
     from sbxreader import sbx_memmap
+    HAS_SBX = True
 except:
-    print("Could not load the sbx reader, installing with pip.")
-    from subprocess import call
-    call("pip install sbxreader", shell=True)
-    from sbxreader import sbx_memmap
-
+    HAS_SBX = False
+
 
 def sbx_to_binary(ops, ndeadcols=-1, ndeadrows=0):
     """  finds scanbox files and writes them to binaries
@@ -31,6 +29,8 @@ def sbx_to_binary(ops, ndeadcols=-1, ndeadrows=0):
             "Ly", "Lx", ops["reg_file"] or ops["raw_file"] is created binary
 
     """
+    if not HAS_SBX:
+        raise ImportError("sbxreader is required for this file type, please 'pip install sbxreader'")
 
     ops1 = init_ops(ops)
     # the following should be taken from the metadata and not needed but the files are initialized before...

suite2p/io/utils.py

@@ -60,6 +60,41 @@ def get_sbx_list(ops):
         print("** Found %d sbx - converting to binary **" % (len(fsall)))
     return fsall, ops
 
+def get_movie_list(ops):
+    """ make list of movie files to process
+    if ops["subfolders"], then all  ops["data_path"][0] / ops["subfolders"] / *.avi or *.mp4
+    if ops["look_one_level_down"], then all tiffs in all folders + one level down
+    """
+    froot = ops["data_path"]
+    # use a user-specified list of tiffs
+    if len(froot) == 1:
+        if "subfolders" in ops and len(ops["subfolders"]) > 0:
+            fold_list = []
+            for folder_down in ops["subfolders"]:
+                fold = os.path.join(froot[0], folder_down)
+                fold_list.append(fold)
+        else:
+            fold_list = ops["data_path"]
+    else:
+        fold_list = froot
+    fsall = []
+    for k, fld in enumerate(fold_list):
+        try:
+            fs = search_for_ext(fld, extension="mp4",
+                                look_one_level_down=ops["look_one_level_down"])
+            fsall.extend(fs)
+        except:
+            fs = search_for_ext(fld, extension="avi",
+                                look_one_level_down=ops["look_one_level_down"])
+            fsall.extend(fs)
+    if len(fsall) == 0:
+        print(fold_list)
+        raise Exception("No files, check path.")
+    else:
+        print("** Found %d movies - converting to binary **" % (len(fsall)))
+    return fsall, ops
+
+
 
 def list_h5(ops):
     froot = os.path.dirname(ops["h5py"])
@@ -249,6 +284,10 @@ def find_files_open_binaries(ops1, ish5=False):
         fs, ops2 = get_nd2_list(ops1[0])
         print("Nikon files:")
         print("\n".join(fs))
+    elif input_format == "movie":
+        fs, ops2 = get_movie_list(ops1[0])
+        print("Movie files:")
+        print("\n".join(fs))
     else:
         # find tiffs
         fs, ops2 = get_tif_list(ops1[0])