Track_with_background_subtraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Set up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, datetime, re\n",
    "from os.path import join\n",
    "from glob import glob\n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "# sys.path.append(join(root_dir,'cvtracer'))\n",
    "from cvt.TrAQ.Trial import Trial\n",
    "from cvt.TrAQ.Tank import Tank\n",
    "from cvt.TrAQ.Group import Group\n",
    "from cvt.TrAQ.CVTracer import CVTracer, create_named_window, wait_on_named_window"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Telling python where things are.\n",
    "\n",
    "The location of a file or directory is called its path.  \n",
    "\n",
    "An absolute path tells your computer how to get to the file or directory from the root of the filesystem (roughly speaking the place you land if you keep pressing \"parent directory\" until it no longer does anything). Depending on your system it should look something like this:  \n",
    "\n",
    "- Windows: `C:\\\\Users\\yaouen\\larval_schooling\\data\\Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi`  \n",
    "- MacOS/Linux: `/home/yaouen/larval_schooling/data/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi`\n",
    "\n",
    "A relative path tells your computer how to get to the file of directory from wherever you are right now. When running a jupyter notebook, relative means relative to whatever directory the notebook is in.  Depending on your system it should look something like this:  \n",
    "\n",
    "- Windows: `..\\data\\Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi`  \n",
    "- MacOS/Linux: `../data/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi`\n",
    "\n",
    "Here `..` means \"go up one directory\".\n",
    "\n",
    "\n",
    "### Suggested directory layout\n",
    "\n",
    "All having the same directory layout makes it easier to exchange data and python scripts. I suggest you set things up like this:\n",
    "\n",
    "<pre>\n",
    "--larval_project\n",
    "   |\n",
    "   |--cvtracer\n",
    "   |\n",
    "   |--data\n",
    "   |   |--Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi\n",
    "   |   |--Pa_Fri_7dpf_GroupA_n5_2020-06-05-083453-0000.avi\n",
    "   |\n",
    "   |--tracking\n",
    "</pre>\n",
    "\n",
    "\n",
    "### Pick an output directory\n",
    "\n",
    "All the output of the tracking will go in the `tracking` directory, in a subdirectory named after the original video file. If you use the suggested layout above, just execute the cell below. If you want the output of the tracking to go somewhere else, edit the `tracking_dir` variable accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tracking_dir = '../tracking' # This means \"starting from where the notebook is, \n",
    "                             # go up one directory, then find a directory called \"tracking\".\n",
    "\n",
    "# This creates the directory if it doesn't already exist.\n",
    "if not os.path.exists(tracking_dir):\n",
    "    os.mkdir(tracking_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-danger\">\n",
    "<p style=\"font-size:200%\"><b>\n",
    "Pick a video to analyze\n",
    "</b></style>\n",
    "</div>\n",
    "\n",
    "Indicate the path to the video (absolute or relative)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_file = '../data/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi'\n",
    "input_file = '../data/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000.avi'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare for tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "''' Extract trial info from the filename and video itself. '''\n",
    "\n",
    "filename,ext = os.path.splitext(os.path.basename(input_file))\n",
    "pop,_,age,group,Nfish,date = filename.split('_')\n",
    "Nfish    = int(re.findall('\\d+',Nfish)[0])\n",
    "date     = ''.join(date.split('-')[:3])\n",
    "\n",
    "cap      = cv2.VideoCapture(input_file)\n",
    "cap.read()\n",
    "Nframes  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "fps      = int(cap.get(cv2.CAP_PROP_FPS))\n",
    "fourcc   = int(cap.get(cv2.CAP_PROP_FOURCC))\n",
    "width    = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
    "height   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
    "cap.release()\n",
    "\n",
    "#==========================================================================\n",
    "\n",
    "''' Define and create necessary folders/files/links. '''\n",
    "\n",
    "output_dir = join(tracking_dir,filename)\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)\n",
    "\n",
    "new_input_file = join(output_dir,'raw'+ext)\n",
    "if not os.path.exists(new_input_file):\n",
    "    os.symlink(os.path.relpath(input_file,output_dir),new_input_file)\n",
    "    \n",
    "bkg_file  = join(output_dir,'background.npy')\n",
    "tank_file = join(output_dir,'tank.pik')\n",
    "\n",
    "settings_file = join(output_dir,'tracking_settings.txt')\n",
    "\n",
    "#==========================================================================\n",
    "\n",
    "''' Set default tracking parameters. '''\n",
    "\n",
    "# By default only track 10 seconds 1 minute in.\n",
    "t_start,t_end  = 60,70 # Times between which to track, in seconds.\n",
    "\n",
    "# Tank.\n",
    "tank_radius    = 10    # Radius of the tank, in centimeters.\n",
    "\n",
    "# Background subtraction (for naive background subtraction only).\n",
    "bkg_frame_skip = 50    # Using every frame of the video to compute the background takes a while.\n",
    "                       # Instead we only use one frame in bkg_frame_skip.\n",
    "bkg_sub_amp    = 4     # Contrast amplification factor applied after background subtraction.\n",
    "\n",
    "# Contour detection.\n",
    "n_pixel_blur   =  3    # square-root of n-pixels for threshold blurring\n",
    "block_size     = 15    # contour block size\n",
    "thresh_offset  = 13    # threshold offset for contour-finding\n",
    "min_area       = 20    # minimum area for threhold detection\n",
    "max_area       = 500   # maximum area for threhold detection\n",
    "RGB            = True  # track in color, false does greyscale\n",
    "online_viewer  = True  # Toggle live preview of tracking.\n",
    "\n",
    "#==========================================================================\n",
    "\n",
    "''' Define and create necessary folders/files/links. '''\n",
    "\n",
    "output_dir = join(tracking_dir,filename)\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)\n",
    "\n",
    "new_input_file = join(output_dir,'raw'+ext)\n",
    "if not os.path.exists(new_input_file):\n",
    "    os.symlink(os.path.relpath(input_file,output_dir),new_input_file)\n",
    "    \n",
    "settings_file = join(output_dir,'tracking_settings.txt')\n",
    "tank_file = join(output_dir,'tank.pik')\n",
    "bkg_file  = join(output_dir,f'background-{bkg_frame_skip}.npy')\n",
    "\n",
    "#==========================================================================\n",
    "\n",
    "''' Prepare to save tracking settings. '''\n",
    "\n",
    "def save_settings():\n",
    "    tracking_info = { k:globals()[k] for k in ['bkg_file','tank_file','t_start','t_end',\n",
    "                                               'tank_radius','bkg_frame_skip','bkg_sub_amp',\n",
    "                                               'n_pixel_blur','block_size','thresh_offset',\n",
    "                                               'min_area','max_area','RGB'] }\n",
    "    trial_info = { k:globals()[k] for k in ['ext','filename','pop','age',\n",
    "                                            'group','Nfish','date',\n",
    "                                            'Nframes','fps','fourcc',\n",
    "                                            'width','height'] }\n",
    "    with open(settings_file,'w') as fh:\n",
    "        for d in trial_info,tracking_info:\n",
    "            for k,v in d.items():\n",
    "                print(f'{k} = {v}',file=fh)\n",
    "    \n",
    "    return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-danger\">\n",
    "<p style=\"font-size:200%\"><b>\n",
    "Locate the tank\n",
    "</b></style>\n",
    "</div>\n",
    "\n",
    "You only need to do this once for each video. Click on three points on the edge of the tank, then press the space bar twice. This will create the file `tank.pik` in the output directory (the one named after the original video). If you run this notebook again on the same video file it load the pre-existing `tank.pik` file. If you want to locate the tank again, delete the file `tank.pik` in the output directory then execute this cell again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "        Tank object loaded from ../tracking/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000/tank.pik \n"
     ]
    }
   ],
   "source": [
    "tank = Tank(new_input_file,tank_radius)\n",
    "if os.path.exists(tank_file):\n",
    "    tank.load(tank_file)\n",
    "else:\n",
    "    tank.locate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Track\n",
    "\n",
    "<div class=\"alert alert-block alert-danger\">\n",
    "<p style=\"font-size:200%\"><b>\n",
    "Pick a tracking method, tweak parameters, then execute.\n",
    "</b></style>\n",
    "</div>\n",
    "\n",
    "## 2.1. No background subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "        Trial loaded from /media/data/work/fau/research/trilab/astyanax_larval_schooling/tracking/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000/trial.pik \n",
      "2 individuals in trial\n",
      "4\n",
      "Using Gaussian Adaptive Threshold\n",
      "Using Inverted Binary Threshold with [0, 100].\n",
      " Group of 2\n",
      "       Current tracking time: 00:01:01:10 \n",
      "       Video capture released.\n",
      "\n",
      "        Trial object saved as /media/data/work/fau/research/trilab/astyanax_larval_schooling/tracking/Pa_Fri_7dpf_GroupA_n2_2020-06-05-120920-0000/trial.pik \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#==========================================================================\n",
    "# Parameter tweaking area.\n",
    "\n",
    "t_start,t_end = 60,70 # Times between which to track, in seconds.\n",
    "# bkg_sub_amp   = 4    # Amplification factor applied after background subtraction.\n",
    "n_pixel_blur  = 5    # square-root of n-pixels for threshold blurring\n",
    "block_size    = 15   # contour block size\n",
    "thresh_offset = 4    # threshold offset for contour-finding\n",
    "min_area      = 50   # minimum area for threhold detection\n",
    "max_area      = 100  # maximum area for threhold detection\n",
    "online_viewer = True # False # \n",
    "\n",
    "#==========================================================================\n",
    "# Actual tracking.\n",
    "\n",
    "save_settings()\n",
    "frame_start,frame_end = int(t_start*fps),int(t_end*fps)\n",
    "\n",
    "trial = Trial(fvideo=new_input_file, n=Nfish, t=pop, date=date, fps=fps, \n",
    "              tank_radius=tank_radius, t_start=t_start, t_end=t_end, reorganize=False)\n",
    "\n",
    "cvt = CVTracer(trial, n_pixel_blur=n_pixel_blur, block_size=block_size, \n",
    "               threshold_offset=thresh_offset, min_area=min_area, view_scale=1, RGB=True,\n",
    "               frame_start=frame_start, frame_end=frame_end, \n",
    "               online=online_viewer)\n",
    "\n",
    "cvt.set_frame(cvt.frame_start)\n",
    "for i_frame in range(cvt.frame_start, cvt.frame_end+1):\n",
    "    if cvt.get_frame():\n",
    "        cvt.mask_tank()\n",
    "        cvt.detect_contours()\n",
    "        cvt.analyze_contours()\n",
    "        cvt.connect_frames()\n",
    "        cvt.update_trial()\n",
    "        cvt.draw()\n",
    "        cvt.write_frame()\n",
    "#         if cvt.online_viewer:\n",
    "#             cvt.draw_contours()\n",
    "        if not cvt.post_frame(delay=1):\n",
    "            break\n",
    "        cvt.print_current_frame()\n",
    "cvt.release()\n",
    "cvt.trial.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvt.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Simple background subtraction\n",
    "\n",
    "## 2.1. Compute the background\n",
    "\n",
    "Compute the background by averaging frames over the entire video. Save it as `background.npy` in the output directory. Use the pre-existing file if there is one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#==========================================================================\n",
    "''' If you  want to change the number of frames used to compute the background do it here.'''\n",
    "# Example: \"bkg_frame_skip = 50\" means \"only use every 50th frame\".\n",
    "bkg_frame_skip = 50\n",
    "#==========================================================================\n",
    "\n",
    "bkg_file  = join(output_dir,f'background-{bkg_frame_skip}.npy')\n",
    "\n",
    "\n",
    "if os.path.exists(bkg_file):\n",
    "    bkg  = np.load(bkg_file)\n",
    "else:\n",
    "    t0    = datetime.datetime.now()\n",
    "    cap   = cv2.VideoCapture(new_input_file)\n",
    "    _,frame = cap.read()\n",
    "    bkg   = np.zeros(frame.shape,dtype=float)\n",
    "    count = 0\n",
    "    # If bkg_frame_skip is small (<10) it may be faster to use\n",
    "    # cap.grab instead of cap.set.\n",
    "    for n in range(0,Nframes,bkg_frame_skip):\n",
    "        cap.set(cv2.CAP_PROP_POS_FRAMES,n)\n",
    "        ret,frame = cap.read()\n",
    "        bkg      += frame\n",
    "        count    += 1\n",
    "    bkg   = bkg / count\n",
    "    np.save(bkg_file,bkg)\n",
    "    print(datetime.datetime.now()-t0)\n",
    "\n",
    "\n",
    "def subtract_background(frame,bkg=bkg,bkg_sub_amp=bkg_sub_amp):\n",
    "    return 255-np.minimum(255,bkg_sub_amp*np.absolute(frame-bkg)).astype(np.uint8)    \n",
    "\n",
    "\n",
    "# Show the background.\n",
    "plt.figure(figsize=(9,9))\n",
    "plt.imshow(bkg.astype(np.uint))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2. Preview simple background subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_skip = 100 # frame skipping rate\n",
    "\n",
    "window_name = 'background subtraction preview'\n",
    "create_named_window(window_name)\n",
    "cap = cv2.VideoCapture(new_input_file)\n",
    "n   = 0\n",
    "while True:\n",
    "#     for i in range(n_skip-1):\n",
    "#         cap.grab()\n",
    "    n += n_skip\n",
    "    cap.set(cv2.CAP_PROP_POS_FRAMES,n)\n",
    "    ret,frame = cap.read()\n",
    "    if not ret:\n",
    "        break\n",
    "    diff = subtract_background(frame)\n",
    "#     diff = cv2.GaussianBlur(diff,(5,)*2,0)\n",
    "    cv2.imshow(window_name,diff)\n",
    "    if wait_on_named_window(window_name,10)==-2:\n",
    "        break\n",
    "cap.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3. Track with simple background subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#==========================================================================\n",
    "# Parameter tweaking area.\n",
    "\n",
    "t_start,t_end = 60,70 # Times between which to track, in seconds.\n",
    "bkg_sub_amp   = 4     # Amplification factor applied after background subtraction.\n",
    "n_pixel_blur  = 7     # square-root of n-pixels for threshold blurring\n",
    "block_size    = 15    # contour block size\n",
    "thresh_offset = 20    # threshold offset for contour-finding\n",
    "min_area      = 20    # minimum area for threhold detection\n",
    "max_area      = 90    # maximum area for threhold detection\n",
    "online_viewer = True  # False # \n",
    "\n",
    "#==========================================================================\n",
    "# Actual tracking.\n",
    "\n",
    "save_settings()\n",
    "frame_start,frame_end = int(t_start*fps),int(t_end*fps)\n",
    "\n",
    "trial = Trial(fvideo=new_input_file, n=Nfish, t=pop, date=date, fps=fps, \n",
    "              tank_radius=tank_radius, t_start=t_start, t_end=t_end, reorganize=False)\n",
    "\n",
    "cvt = CVTracer(trial, n_pixel_blur=n_pixel_blur, block_size=block_size, \n",
    "               threshold_offset=thresh_offset, min_area=min_area, view_scale=1, RGB=True,\n",
    "               frame_start=frame_start, frame_end=frame_end, \n",
    "               online=online_viewer)\n",
    "\n",
    "cvt.set_frame(cvt.frame_start)\n",
    "for i_frame in range(cvt.frame_start, cvt.frame_end+1):\n",
    "    if cvt.get_frame():\n",
    "        cvt.frame = subtract_background(cvt.frame,bkg_sub_amp=bkg_sub_amp)\n",
    "        cvt.mask_tank()\n",
    "        cvt.detect_contours()\n",
    "        cvt.analyze_contours()\n",
    "        cvt.connect_frames()\n",
    "        cvt.update_trial()\n",
    "        cvt.draw()\n",
    "        cvt.write_frame()\n",
    "        if not cvt.post_frame(delay=1):\n",
    "            break\n",
    "        cvt.print_current_frame()\n",
    "cvt.release()\n",
    "cvt.trial.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cvt.release()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4.3. MOG2 background subtraction\n",
    "\n",
    "Uses openCV's MOG2 algorithm. Added to cvtracer by Adam for the dark astyanax schooling trials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#==========================================================================\n",
    "# Parameter tweaking area.\n",
    "\n",
    "t_start,t_end = 60,70 # Times between which to track, in seconds.\n",
    "n_pixel_blur  = 3     # square-root of n-pixels for threshold blurring\n",
    "block_size    = 15    # contour block size\n",
    "thresh_offset = 3     # threshold offset for contour-finding\n",
    "min_area      = 15    # minimum area for threhold detection\n",
    "max_area      = 150   # maximum area for threhold detection\n",
    "online_viewer = True  # False # \n",
    "\n",
    "#==========================================================================\n",
    "# Actual tracking.\n",
    "\n",
    "save_settings()\n",
    "frame_start,frame_end = int(t_start*fps),int(t_end*fps)\n",
    "\n",
    "trial = Trial(fvideo=new_input_file, n=Nfish, t=pop, date=date, fps=fps, \n",
    "              tank_radius=tank_radius, t_start=t_start, t_end=t_end, reorganize=False)\n",
    "\n",
    "cvt = CVTracer(trial, n_pixel_blur=n_pixel_blur, block_size=block_size, \n",
    "               threshold_offset=thresh_offset, min_area=min_area, view_scale=1, RGB=True,\n",
    "               frame_start=frame_start, frame_end=frame_end, \n",
    "               online=online_viewer, \n",
    "               MOG2=True)\n",
    "\n",
    "cvt.set_frame(cvt.frame_start)\n",
    "for i_frame in range(cvt.frame_start, cvt.frame_end+1):\n",
    "    if cvt.get_frame():\n",
    "        cvt.mask_tank()\n",
    "        cvt.mask_background()\n",
    "        cvt.detect_contours()\n",
    "        cvt.analyze_contours()\n",
    "        cvt.connect_frames()\n",
    "        cvt.update_trial()\n",
    "        cvt.draw()\n",
    "        cvt.write_frame()\n",
    "        if not cvt.post_frame(delay=1):\n",
    "            break\n",
    "        cvt.print_current_frame()\n",
    "cvt.release()\n",
    "cvt.trial.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvt.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}