stripe_fixer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## load modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from dask.distributed import Client, LocalCluster\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt \n",
    "from scipy.signal import periodogram, find_peaks\n",
    "import re\n",
    "import warnings\n",
    "from os import listdir\n",
    "from pathlib import Path\n",
    "import cv2\n",
    "import dask as da\n",
    "import math\n",
    "import dask.array as darr\n",
    "import xarray as xr\n",
    "import zarr as zr\n",
    "from natsort import natsorted\n",
    "from tifffile import TiffFile, imread\n",
    "import matplotlib\n",
    "import shutil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## set path and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "lines_to_next_cell": 2,
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Set up Initial Basic Parameters#\n",
    "dpath = \"FILE_PATH_HERE\"\n",
    "dpath = os.path.abspath(dpath)\n",
    "framesPerFile = 1000\n",
    "\n",
    "# Pre-processing Parameters#\n",
    "param_load_videos = {\n",
    "    \"pattern\": \"[0-9]+\\.avi$\", \n",
    "    \"dtype\": np.uint8,\n",
    "    \"downsample\": dict(frame=1, height=1, width=1),\n",
    "    \"downsample_strategy\": \"subset\",\n",
    "}\n",
    "\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"\n",
    "os.environ[\"NUMBA_NUM_THREADS\"] = \"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## functions for loading in all videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_videos(\n",
    "    vpath,\n",
    "    pattern=r\"msCam[0-9]+\\.avi$\",\n",
    "    dtype=np.float64,\n",
    "    in_memory=False,\n",
    "    downsample=None,\n",
    "    downsample_strategy=\"subset\",\n",
    "    post_process=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Load videos from the folder specified in `vpath` and according to the regex\n",
    "    `pattern`, then concatenate them together across time and return a\n",
    "    `xarray.DataArray` representation of the concatenated videos. The default\n",
    "    assumption is video filenames start with ``msCam`` followed by at least a\n",
    "    number, and then followed by ``.avi``. In addition, it is assumed that the\n",
    "    name of the folder correspond to a recording session identifier.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    vpath : str\n",
    "        The path to search for videos\n",
    "    pattern : str, optional\n",
    "        The pattern that describes filenames of videos. (Default value =\n",
    "        'msCam[0-9]+\\.avi')\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    xarray.DataArray or None\n",
    "        The labeled 3-d array representation of the videos with dimensions:\n",
    "        ``frame``, ``height`` and ``width``. Returns ``None`` if no data was\n",
    "        found in the specified folder.\n",
    "    \"\"\"\n",
    "    vpath = os.path.normpath(vpath)\n",
    "    ssname = os.path.basename(vpath)\n",
    "    vlist = natsorted(\n",
    "        [vpath + os.sep + v for v in os.listdir(vpath) if re.search(pattern, v)]\n",
    "    )\n",
    "    if not vlist:\n",
    "        raise FileNotFoundError(\n",
    "            \"No data with pattern {}\"\n",
    "            \" found in the specified folder {}\".format(pattern, vpath)\n",
    "        )\n",
    "    print(\"loading {} videos in folder {}\".format(len(vlist), vpath))\n",
    "\n",
    "    file_extension = os.path.splitext(vlist[0])[1]\n",
    "    if file_extension in (\".avi\", \".mkv\"):\n",
    "        movie_load_func = load_avi_lazy\n",
    "    elif file_extension == \".tif\":\n",
    "        movie_load_func = load_tif_lazy\n",
    "    else:\n",
    "        raise ValueError(\"Extension not supported.\")\n",
    "\n",
    "    varr_list = [movie_load_func(v) for v in vlist]\n",
    "    varr = darr.concatenate(varr_list, axis=0)\n",
    "    varr = xr.DataArray(\n",
    "        varr,\n",
    "        dims=[\"frame\", \"height\", \"width\"],\n",
    "        coords=dict(\n",
    "            frame=np.arange(varr.shape[0]),\n",
    "            height=np.arange(varr.shape[1]),\n",
    "            width=np.arange(varr.shape[2]),\n",
    "        ),\n",
    "    )\n",
    "    if dtype:\n",
    "        varr = varr.astype(dtype)\n",
    "    if downsample:\n",
    "        bin_eg = {d: np.arange(0, varr.sizes[d], w) for d, w in downsample.items()}\n",
    "        if downsample_strategy == \"mean\":\n",
    "            varr = (\n",
    "                varr.coarsen(**downsample, boundary=\"trim\")\n",
    "                .mean()\n",
    "                .assign_coords(**bin_eg)\n",
    "            )\n",
    "        elif downsample_strategy == \"subset\":\n",
    "            varr = varr.sel(**bin_eg)\n",
    "        else:\n",
    "            warnings.warn(\"unrecognized downsampling strategy\", RuntimeWarning)\n",
    "    varr = varr.rename(\"fluorescence\")\n",
    "    if post_process:\n",
    "        varr = post_process(varr, vpath, ssname, vlist, varr_list)\n",
    "    return varr\n",
    "\n",
    "def load_tif_lazy(fname):\n",
    "    data = TiffFile(fname)\n",
    "    f = len(data.pages)\n",
    "\n",
    "    fmread = da.delayed(load_tif_perframe)\n",
    "    flist = [fmread(fname, i) for i in range(f)]\n",
    "\n",
    "    sample = flist[0].compute()\n",
    "    arr = [\n",
    "        da.array.from_delayed(fm, dtype=sample.dtype, shape=sample.shape)\n",
    "        for fm in flist\n",
    "    ]\n",
    "    return da.array.stack(arr, axis=0)\n",
    "\n",
    "\n",
    "def load_tif_perframe(fname, fid):\n",
    "    return imread(fname, key=fid)\n",
    "\n",
    "\n",
    "def load_avi_lazy(fname):\n",
    "    cap = cv2.VideoCapture(fname)\n",
    "    f = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "    fmread = da.delayed(load_avi_perframe)\n",
    "    flist = [fmread(fname, i) for i in range(f)]\n",
    "    sample = flist[0].compute()\n",
    "    arr = [\n",
    "        da.array.from_delayed(fm, dtype=sample.dtype, shape=sample.shape)\n",
    "        for fm in flist\n",
    "    ]\n",
    "    return da.array.stack(arr, axis=0)\n",
    "\n",
    "\n",
    "def load_avi_perframe(fname, fid):\n",
    "    cap = cv2.VideoCapture(fname)\n",
    "    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "    cap.set(cv2.CAP_PROP_POS_FRAMES, fid)\n",
    "    ret, fm = cap.read()\n",
    "    if ret:\n",
    "        return np.flip(cv2.cvtColor(fm, cv2.COLOR_RGB2GRAY), axis=0)\n",
    "    else:\n",
    "        print(\"frame read failed for frame {}\".format(fid))\n",
    "        return np.zeros((h, w))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## start cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster = LocalCluster(n_workers=1, memory_limit=\"8GB\")\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Misalignment detection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## loading videos and visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "autoscroll": false,
    "ein.tags": "worksheet-0",
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "varr = load_videos(dpath, **param_load_videos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_ref = varr.chunk({\"frame\": 20, \"height\": -1, \"width\": -1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## get rid of V4 stripe noise (present on earlier releases of the V4 scope)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# f = 0\n",
    "# im = varr_ref[f,:,:].copy().values\n",
    "# im_fft = fftpack.fft2(im.astype(np.float64))\n",
    "# im_fft2 = im_fft.copy()\n",
    "# n = 5\n",
    "# y = 2\n",
    "# im_fft2[n:im_fft2.shape[0]-n, :y] = 0\n",
    "# im_fft2[n:im_fft2.shape[0]-n, -y:] = 0\n",
    "# im_new = fftpack.ifft2(im_fft2).real\n",
    "# fig, ax = plt.subplots(ncols=4, figsize=(50,15))\n",
    "# ax[0].imshow(np.abs(im_fft), norm=LogNorm(vmin=5))\n",
    "# ax[0].set_title('Fourier transform', fontsize=30)\n",
    "# ax[1].imshow(np.abs(im_fft2), norm=LogNorm(vmin=5))\n",
    "# ax[1].set_title('Filtered spectrum', fontsize=30)\n",
    "# ax[2].imshow(im, cmap='binary_r', vmin=0, vmax=255, origin='lower')\n",
    "# ax[2].set_title('Old image', fontsize=30)\n",
    "# ax[3].imshow(im_new, cmap='binary_r', vmin=0, vmax=255, origin='lower')\n",
    "# ax[3].set_title('New image', fontsize=30)\n",
    "# plt.show()\n",
    "# im_opts = dict(frame_width=500, aspect=608/608, cmap='Spectral_r', colorbar=True)\n",
    "# hv.Image(np.log(np.abs(im_fft)), ['width', 'height'], label='before_mc').opts(**im_opts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def sensor_denoise(varr, n, y):\n",
    "#     return xr.apply_ufunc(\n",
    "#         sensor_denoise_perframe,\n",
    "#         varr,\n",
    "#         input_core_dims=[['height', 'width']],\n",
    "#         output_core_dims=[['height', 'width']],\n",
    "#         vectorize=True,\n",
    "#         dask='parallelized',\n",
    "#         output_dtypes=[np.float64],\n",
    "#         kwargs=dict(n=n, y=y))\n",
    "\n",
    "# def sensor_denoise_perframe(f, n, y):\n",
    "#     im_fft = fftpack.fft2(f)\n",
    "#     im_fft[n:im_fft.shape[0]-n, :y] = 0\n",
    "#     im_fft[n:im_fft.shape[0]-n, -y:] = 0\n",
    "#     im_new = fftpack.ifft2(im_fft).real\n",
    "#     return im_new\n",
    "\n",
    "# varr_ref = sensor_denoise(varr_ref, n=5, y=2)\n",
    "# min_val = varr_ref.min().compute()\n",
    "# max_val = varr_ref.max().compute()\n",
    "# varr_ref = ((varr_ref - min_val) / (max_val - min_val) * 255)\n",
    "# varr_ref = varr_ref.astype(np.uint8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## find frames with misaligned rows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, find the frames where the stripes exist.\n",
    "\n",
    "from scipy.signal import periodogram, find_peaks\n",
    "def no_stripes_frames(varr):\n",
    "    return xr.apply_ufunc(\n",
    "        no_stripes_in_frame,\n",
    "        varr,\n",
    "        input_core_dims=[['height', 'width']],\n",
    "        vectorize=True,\n",
    "        dask='parallelized',\n",
    "        output_dtypes=[bool])\n",
    "\n",
    "def no_stripes_in_frame(x, thresh=20):\n",
    "    x = x[:,0]\n",
    "    f, Pxx_spec = periodogram(x, 1)\n",
    "    peaks = find_peaks(np.sqrt(Pxx_spec)[:int(len(Pxx_spec)/5)], height=thresh)[0]\n",
    "    return not any(f[peaks] > 0.03)\n",
    "\n",
    "frames_without_stripes = no_stripes_frames(varr_ref).values\n",
    "bad_frames = np.asarray(varr_ref[~frames_without_stripes].frame)\n",
    "\n",
    "print('bad frames:')\n",
    "bad_frames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot what the minimum projection would look like without the misaligned frames. \n",
    "If stripes still appear, try lowering spec_thresh in no_stripes_in_frame() and rerunning it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_min = varr_ref[frames_without_stripes].min(\"frame\").compute()\n",
    "plt.imshow(varr_min)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## get file names that need to be replaced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_filename(frame_list):\n",
    "    # Get the file names associated with each frame. \n",
    "    vid_numbers = np.unique([math.floor(f/framesPerFile) for f in frame_list])\n",
    "    fnames = [os.path.join(dpath, str(n) + '.avi') for n in vid_numbers]\n",
    "    \n",
    "    # Get the frame number within that video file. \n",
    "    relative_frame_numbers = []\n",
    "    for n in vid_numbers:\n",
    "        quotient, remainder = np.divmod(frame_list, n*framesPerFile)\n",
    "        \n",
    "        relative_frame_numbers.append(remainder[(quotient==1) & (remainder < framesPerFile)])\n",
    "    \n",
    "    return fnames, relative_frame_numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames, frame_numbers = get_filename(bad_frames)\n",
    "fnames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Realignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## move the original files to a new folder and then write a new file with fixed frames\\\n",
    "\n",
    "The function below makes a folder called \"originals\" and moves the original file there. Then it writes a new video file with the same name to replace it, but with the frames fixed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_frame(frame, shift_amount=8184*2, show_plot=False, ax=None):\n",
    "    buffer_size = 8184\n",
    "    flattened_frame = frame.flatten()\n",
    "    frame_size = frame.shape\n",
    "    n_pixels = len(flattened_frame)\n",
    "    \n",
    "    for pixel_number in range(n_pixels):\n",
    "        buf_num = int(pixel_number/buffer_size)\n",
    "        \n",
    "        if ((buf_num % 2) == 0):\n",
    "            if ((pixel_number + shift_amount) < n_pixels):\n",
    "                flattened_frame[pixel_number] = flattened_frame[pixel_number + shift_amount]\n",
    "    \n",
    "    fixed_frame = flattened_frame.reshape(frame_size)\n",
    "    \n",
    "    if show_plot:\n",
    "        if ax is None:\n",
    "            fig, ax = plt.subplots(figsize=(24,24))\n",
    "        ax.imshow(fixed_frame)\n",
    "    \n",
    "    return fixed_frame\n",
    "\n",
    "def fix_video(fnames, frame_numbers):\n",
    "    folder = os.path.join(os.path.split(fnames[0])[0], 'originals')\n",
    "    if not os.path.exists(folder):\n",
    "        os.mkdir(folder)\n",
    "        print(f'Created {folder}')\n",
    "        \n",
    "    compressionCodec = \"FFV1\"\n",
    "    codec = cv2.VideoWriter_fourcc(*compressionCodec)\n",
    "    \n",
    "    buffer_size = 8184\n",
    "    shift_amount = buffer_size*2\n",
    "    \n",
    "    # For each video...\n",
    "    for video, bad_frame_numbers in zip(fnames, frame_numbers):\n",
    "        print(f'Rewriting {video}')\n",
    "        cap = cv2.VideoCapture(video)\n",
    "        rows, cols = int(cap.get(4)), int(cap.get(3))\n",
    "        \n",
    "        fname = os.path.split(video)[1]\n",
    "        move_fpath = os.path.join(folder, fname)\n",
    "        \n",
    "        # Move the original file. \n",
    "        if not os.path.exists(move_fpath):\n",
    "            shutil.move(video, move_fpath)\n",
    "            print(f'Moved {video} to {move_fpath}')\n",
    "        else:\n",
    "            raise FileExistsError('The folder is already storing an original file. Aborting to prevent overwrite.')\n",
    "\n",
    "        writeFile = cv2.VideoWriter(video, codec, 60, (cols,rows), isColor=False)\n",
    "        \n",
    "        for frame_number in tqdm(range(int(cap.get(7)))):\n",
    "            ret, frame = cap.read()\n",
    "\n",
    "            if ret:\n",
    "                write_frame = frame[:,:,0]\n",
    "\n",
    "                if frame_number in bad_frame_numbers:\n",
    "                    write_frame = fix_frame(write_frame, shift_amount, show_plot=False)\n",
    "\n",
    "                writeFile.write(np.uint8(write_frame))\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        writeFile.release()\n",
    "        cap.release()\n",
    "    cv2.destroyAllWindows()\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fix_video(fnames, frame_numbers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Further corrections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, inspect the new videos and compare them to the ones in the \"original\" folder. If the new ones look pristine, you're done! Skip the remaining steps and proceed to Minian.\n",
    "\n",
    "If the new videos still have stripes, you have no choice but to replace those frames with the most recent good frame. Usually these stripes are \"nonstationary\", as in they sweep across the image rather than being static. I haven't figured out how to correct these, so the workaround is to replace them with the most recent good frame. Proceed ahead. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Reload the whole again movie. Make sure you replaced the files with the repaired ones!! \n",
    "varr = load_videos(dpath, **param_load_videos)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "varr_ref = varr.chunk({\"frame\": 20, \"height\": -1, \"width\": -1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# varr_ref = sensor_denoise(varr_ref, n=5, y=2)\n",
    "# min_val = varr_ref.min().compute()\n",
    "# max_val = varr_ref.max().compute()\n",
    "# varr_ref = ((varr_ref - min_val) / (max_val - min_val) * 255)\n",
    "# varr_ref = varr_ref.astype(np.uint8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Look for striped frames again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames_without_stripes = no_stripes_frames(varr_ref).values\n",
    "bad_frames = np.asarray(varr_ref[~frames_without_stripes].frame)\n",
    "bad_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inspect frames by eye and optionally, save the frames to a folder.\n",
    "\n",
    "def plot_bad_frame(bad_frame, save=False, show_plot=False, dpath=dpath):\n",
    "    x = varr_ref.sel(frame=bad_frame)\n",
    "    f, Pxx_spec = periodogram(x[:,0], 1)\n",
    "    peaks = find_peaks(np.sqrt(Pxx_spec)[:int(len(Pxx_spec)/5)], height=20)[0]\n",
    "    fig, ax = plt.subplots(ncols=2, figsize=(40,10))\n",
    "    ax[0].imshow(x, cmap='binary_r', aspect='equal', origin='lower')\n",
    "    ax[1].semilogy(f, np.sqrt(Pxx_spec), c='k')\n",
    "    ax[1].scatter(f[peaks], np.sqrt(Pxx_spec)[peaks], s=100, c='brown')\n",
    "    if len(peaks) > 1:\n",
    "        ax[1].text(x = 0.01, y = 0.1, s='Frame has stripes', color='brown', fontsize=20)\n",
    "    ax[1].set_ylim([1e-2, 1e3])\n",
    "    ax[1].margins(x=0.01)\n",
    "    ax[1].set_xlabel('Frequency [Hz]', fontsize=30)\n",
    "    ax[1].set_ylabel('Linear spectrum [V RMS]', fontsize=30)\n",
    "    \n",
    "    if show_plot:\n",
    "        fig.show()\n",
    "    else:\n",
    "        plt.close(fig)\n",
    "    \n",
    "    if save:\n",
    "        fig.savefig(os.path.join(dpath, 'bad_frames', str(bad_frame) + '.png'))\n",
    "        \n",
    "    return f, peaks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bad_frames_folder = os.path.join(dpath, 'bad_frames')\n",
    "if not os.path.exists(bad_frames_folder):\n",
    "    os.mkdir(bad_frames_folder)\n",
    "for bad_frame in tqdm(bad_frames):\n",
    "    plot_bad_frame(bad_frame, show_plot=False, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Identify chunks where the stripes exist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_consecutives(vals, step=1):\n",
    "    \"\"\"Return list of consecutive lists of numbers from vals (number list).\"\"\"\n",
    "    run = []\n",
    "    result = [run]\n",
    "    expect = None\n",
    "    for v in vals:\n",
    "        if (v == expect) or (expect is None):\n",
    "            run.append(v)\n",
    "        else:\n",
    "            run = [v]\n",
    "            result.append(run)\n",
    "        expect = v + step\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rewrite_video(fnames, frame_numbers):\n",
    "    folder = os.path.join(os.path.split(fnames[0])[0], 'failed_to_fix')\n",
    "    if not os.path.exists(folder):\n",
    "        os.mkdir(folder)\n",
    "        print(f'Created {folder}')\n",
    "\n",
    "    compressionCodec = \"FFV1\"\n",
    "    codec = cv2.VideoWriter_fourcc(*compressionCodec)\n",
    "\n",
    "    # For each video...\n",
    "    for video, frames in zip(fnames, frame_numbers):\n",
    "        print(f'Rewriting {video}')\n",
    "        cap = cv2.VideoCapture(video)\n",
    "        rows, cols = int(cap.get(4)), int(cap.get(3))\n",
    "        chunk_number = 0\n",
    "\n",
    "        fname = os.path.split(video)[1]\n",
    "        move_fpath = os.path.join(folder, fname)\n",
    "        \n",
    "        # Move the original file. \n",
    "        if not os.path.exists(move_fpath):\n",
    "            shutil.move(video, move_fpath)\n",
    "            print(f'Moved {video} to {move_fpath}')\n",
    "        else:\n",
    "            raise FileExistsError('The folder is already storing an original file. Aborting to prevent overwrite.')\n",
    "        writeFile = cv2.VideoWriter(video, codec, 60, (cols,rows), isColor=False)\n",
    "\n",
    "        # Group the frame numbers into chunks. Get the frame number that will replace all the bad frames.\n",
    "        # This is the frame right before the first bad frame in each chunk.\n",
    "        frame_chunks = group_consecutives(frames)\n",
    "        frame_chunks = [np.arange(chunk[0], chunk[-1]+2) for chunk in frame_chunks]\n",
    "        replacement_frame_number = [frames[0]-1 for frames in frame_chunks]\n",
    "\n",
    "        # For each frame in the video...\n",
    "        for frame_number in tqdm(range(int(cap.get(7)))):\n",
    "            ret, frame = cap.read()\n",
    "            frame = frame[:,:,1]\n",
    "\n",
    "            if ret:\n",
    "                # If it's a replacement frame, store it. \n",
    "                if frame_number in replacement_frame_number:\n",
    "                    #print('found replacement')\n",
    "                    replacement_frame = frame\n",
    "\n",
    "                    chunk_number = replacement_frame_number.index(frame_number)\n",
    "\n",
    "                # If it's a bad frame, replace it with the replacement frame. \n",
    "                elif frame_number in frame_chunks[chunk_number]:\n",
    "                    #print('replacing a frame')\n",
    "                    frame = replacement_frame\n",
    "\n",
    "                #print('writing')\n",
    "                writeFile.write(np.uint8(frame))\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        writeFile.release()\n",
    "        cap.release()\n",
    "    cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames, frame_numbers = get_filename(bad_frames)\n",
    "fnames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## write new videos by replacing the bad frames with the last good one\n",
    "Running the cell below will create another directory \"failed_to_fix\". It will move the files you recently created above that are still not okay, and it will move those to the \"failed_to_fix\" folder. Then it will write yet more new videos that will replace the striped frames with the last good one. After this finishes, check the new video files to see if they look okay. If so, you should be good to proceed to Minian. If not, you may have to adjust some parameters for stripe detection and try again. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rewrite_video(fnames, frame_numbers)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "minian-perf",
   "language": "python",
   "name": "minian-perf"
  },
  "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.8.1"
  },
  "name": "pipeline.ipynb"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}