updating code for avi / mp4 support

Author: carsen-stringer

notebooks/example_notebook.ipynb

@@ -8,14 +8,22 @@
    "source": [
     "import suite2p\n",
     "from suite2p import default_settings\n",
+    "from suite2p.run_s2p import logger_setup\n",
+    "\n",
     "db = {\"data_path\": [\"/media/carsen/ssd2/suite2p_paper/GT1/\"], \n",
     "      \"file_list\": [\"input_1500.tif\"], # (optional) Specify files you'd like to specifically use in the data_path\n",
     "      \"input_format\": \"tif\",\n",
     "      \"nplanes\": 1, # each tiff has these many planes in sequence\n",
     "      \"nchannels\": 1, # each tiff has these many channels per plane\n",
-    "      \"keep_movie_raw\": True,\n",
+    "      \"keep_movie_raw\": True, # optional, to keep the raw movie after registration\n",
     "}\n",
     "\n",
+    "db[\"save_path0\"] = db[\"data_path\"][0]\n",
+    "\n",
+    "# initialize logger\n",
+    "logger_setup(db[\"save_path0\"])\n",
+    "\n",
+    "# initialize settings\n",
     "settings = default_settings()\n",
     "settings[\"registration\"][\"align_by_chan2\"] = False\n",
     "settings[\"run\"][\"do_registration\"] = 1\n",
@@ -24,7 +32,7 @@
     "settings[\"detection\"][\"threshold_scaling\"] = 1.0\n",
     "settings[\"detection\"][\"algorithm\"] = \"sparsery\"\n",
     "settings[\"torch_device\"] = \"mps\" # use mps for mac, cuda for nvidia gpu, or \"cpu\"\n",
-    "output_settings = suite2p.run_s2p(settings=settings, db=db)"
+    "db_paths = suite2p.run_s2p(settings=settings, db=db)"
    ]
   },
   {

notebooks/run_movie_avi_mp4.ipynb

@@ -0,0 +1,65 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import suite2p\n",
+    "from suite2p import default_settings\n",
+    "from suite2p.run_s2p import logger_setup\n",
+    "\n",
+    "# ** can use *.mp4 or *.avi files \n",
+    "db = {\"data_path\": [\"/media/carsen/disk2/test_suite2p/M0/\"], \n",
+    "      \"input_format\": \"movie\",\n",
+    "      \"nplanes\": 1, \n",
+    "      \"nchannels\": 1, \n",
+    "      \"keep_movie_raw\": False, # optional, to keep the raw movie after registration\n",
+    "}\n",
+    "\n",
+    "db[\"save_path0\"] = db[\"data_path\"][0]\n",
+    "\n",
+    "# initialize logger\n",
+    "logger_setup(db[\"save_path0\"])\n",
+    "\n",
+    "# initialize settings\n",
+    "settings = default_settings()\n",
+    "settings[\"fs\"] = 30\n",
+    "settings[\"tau\"] = 1.25 # timescale of gcamp to use for deconvolution\n",
+    "settings[\"detection\"][\"threshold_scaling\"] = 1.0\n",
+    "settings[\"detection\"][\"algorithm\"] = \"sparsery\"\n",
+    "settings[\"torch_device\"] = \"mps\" # use mps for mac, cuda for nvidia gpu, or \"cpu\"\n",
+    "db_paths = suite2p.run_s2p(settings=settings, db=db)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "s2p",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

suite2p/io/movie.py

@@ -114,46 +114,50 @@ def get_frames(self, cframes):
             nk += nt0
         return im
 
-def movie_to_binary(settings):
+def movie_to_binary(dbs, settings, reg_file, reg_file_chan2):
     """  finds movie files and writes them to binaries
 
     Parameters
     ----------
-    settings : dictionary
-        "nplanes", "data_path", "save_path", "save_folder", "fast_disk",
-        "nchannels", "keep_movie_raw", "look_one_level_down" (optional: "subfolders")
+    dbs : list of dict
+        Database dictionaries for each plane. Must contain keys "file_list",
+        "nplanes", "nchannels", "batch_size", "h5py_key", and "functional_chan".
+        Updated in-place with "Ly", "Lx", "nframes", "nframes_per_folder",
+        "meanImg", and "meanImg_chan2".
+    settings : dict
+        Suite2p settings dictionary, saved alongside each plane's database.
+    reg_file : list of file objects
+        Opened binary files for writing each plane's functional channel data.
+    reg_file_chan2 : list of file objects
+        Opened binary files for writing each plane's second channel data
+        (used only when nchannels > 1).
 
     Returns
     -------
-        settings : dictionary of first plane
-            "Ly", "Lx", settings["reg_file"] or settings["raw_file"] is created binary
-
+    dbs : list of dict
+        Updated database dictionaries with image dimensions, frame counts, and
+        mean images populated.
     """
     if not HAS_CV2:
         raise ImportError("cv2 is required for this file type, please 'pip install opencv-python-headless'")
 
-    settings1 = init_settings(settings)
-
-    nplanes = settings1[0]["nplanes"]
-    nchannels = settings1[0]["nchannels"]
-    
-    # open all binary files for writing
-    settings1, filenames, reg_file, reg_file_chan2 = find_files_open_binaries(settings1)
 
-    ik = 0
-    for j in range(settings["nplanes"]):
-        settings1[j]["nframes_per_folder"] = np.zeros(len(filenames), np.int32)
+    nplanes = dbs[0]["nplanes"]
+    nchannels = dbs[0]["nchannels"]
+    flist = dbs[0]["file_list"]
 
+    # todo: implement for multiple movies
+    # for j in range(nplanes):
+    #    dbs[j]["nframes_per_folder"] = np.zeros(len(flist), np.int32)
 
     ncp = nplanes * nchannels
-    nbatch = ncp * int(np.ceil(settings1[0]["batch_size"] / ncp))
-    logger.info(filenames)
+    nbatch = ncp * int(np.ceil(dbs[0]["batch_size"] / ncp))
+    logger.info(flist)
     t0 = time.time()
-    with VideoReader(filenames=filenames) as vr:
-        if settings1[0]["fs"]<=0:
-            for settings in settings1:
-                settings["fs"] = vr.fs
-
+    with VideoReader(filenames=flist) as vr:
+        if settings["fs"]<=0:
+            settings["fs"] = vr.fs
+            
         nframes_all = vr.cumframes[-1]
         nbatch = min(nbatch, nframes_all)
         nfunc = settings["functional_chan"] - 1 if nchannels > 1 else 0
@@ -167,47 +171,46 @@ def movie_to_binary(settings):
             nframes = im.shape[0]
             for j in range(0, nplanes):
                 if ik == 0:
-                    settings1[j]["meanImg"] = np.zeros((im.shape[1], im.shape[2]),
+                    dbs[j]["meanImg"] = np.zeros((im.shape[1], im.shape[2]),
                                                     np.float32)
                     if nchannels > 1:
-                        settings1[j]["meanImg_chan2"] = np.zeros(
+                        dbs[j]["meanImg_chan2"] = np.zeros(
                             (im.shape[1], im.shape[2]), np.float32)
-                    settings1[j]["nframes"] = 0
+                    dbs[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))
-                settings1[j]["meanImg"] += im2write.astype(np.float32).sum(axis=0)
+                dbs[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))
-                    settings1[j]["meanImg_chan2"] += im2write.astype(
+                    dbs[j]["meanImg_chan2"] += im2write.astype(
                         np.float32).sum(axis=0)
-                settings1[j]["nframes"] += im2write.shape[0]
-                #settings1[j]["nframes_per_folder"][ih5] += im2write.shape[0]
+                dbs[j]["nframes"] += im2write.shape[0]
+                
             ik += nframes
             if ik % (nbatch * 4) == 0:
                 logger.info("%d frames of binary, time %0.2f sec." %
                       (ik, time.time() - t0))
 
     # write settings files
-    do_registration = settings1[0]["do_registration"]
-    for settings in settings1:
-        settings["Ly"] = im2write.shape[1]
-        settings["Lx"] = im2write.shape[2]
+    # update dbs with image dimensions and mean images
+    do_registration = settings["run"]["do_registration"]
+    for db in dbs:
+        db["Ly"] = im2write.shape[1]
+        db["Lx"] = im2write.shape[2]
         if not do_registration:
-            settings["yrange"] = np.array([0, settings["Ly"]])
-            settings["xrange"] = np.array([0, settings["Lx"]])
-        settings["meanImg"] /= settings["nframes"]
-        if nchannels > 1:
-            settings["meanImg_chan2"] /= settings["nframes"]
-        np.save(settings["settings_path"], settings)
-    # close all binary files and write settings files
-    for j in range(nplanes):
-        reg_file[j].close()
+            db["yrange"] = np.array([0, db["Ly"]])
+            db["xrange"] = np.array([0, db["Lx"]])
+        db["meanImg"] /= db["nframes"]
         if nchannels > 1:
-            reg_file_chan2[j].close()
-    return settings1[0]
+            db["meanImg_chan2"] /= db["nframes"]
+        # Save db and settings to each plane folder
+        np.save(db["db_path"], db)
+        np.save(db["settings_path"], settings)
+
+    return dbs