Trackpy Click to Annotate V2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from pandas import DataFrame, Series  # Not used right now\n",
    "\n",
    "import trackpy as tp\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "# Optionally, tweak styles.\n",
    "mpl.rc('figure',  figsize=(8, 8))\n",
    "mpl.rc('image', cmap='gray')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "file_name = \"Pool.mp4\"\n",
    "#file_name = \"LampVid2.mp4\"\n",
    "cap = cv2.VideoCapture(file_name)\n",
    "\n",
    "imgs = []\n",
    "cimgs = []\n",
    "mimgs = []\n",
    "\n",
    "ret, img = cap.read()\n",
    "r = cv2.selectROI('Select your ROI with the mouse.  Press Enter/Space to finalize',img)\n",
    "# Space / Enter to confirm selection\n",
    "#cv2.waitKey(0)\n",
    "cv2.destroyWindow('Select your ROI with the mouse.  Press Enter/Space to finalize')\n",
    "img = img[int(r[1]):int(r[1]+r[3]), int(r[0]):int(r[0]+r[2])]\n",
    "\n",
    "imgs.append(img)\n",
    "\n",
    "#change range to load more frames\n",
    "for x in range(100):\n",
    "    ret, img = cap.read()\n",
    "    img = img[int(r[1]):int(r[1]+r[3]), int(r[0]):int(r[0]+r[2])]\n",
    "    imgs.append(img)\n",
    "\n",
    "'''\n",
    "cimgs.append(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))\n",
    "img_b,img_g,img_r = cv2.split(img)\n",
    "mimgs.append(cv2.addWeighted(img_r,.8,img_b,.8,10))\n",
    "\n",
    "for x in range(100):\n",
    "    ret, img = cap.read()\n",
    "    img = img[int(r[1]):int(r[1]+r[3]), int(r[0]):int(r[0]+r[2])]\n",
    "    imgs.append(img)\n",
    "    cimgs.append(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))\n",
    "    img_b,img_g,img_r = cv2.split(img)\n",
    "    mimgs.append(cv2.addWeighted(img_r,.8,img_b,.8,10))\n",
    "'''\n",
    "\n",
    "fig,ax = plt.subplots()\n",
    "ax.imshow(imgs[0])\n",
    "plt.show()\n",
    "\n",
    "print(\"ROI selected\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_img = np.mean(imgs[:], axis=0).astype(np.uint8)   #Averages all imgs to get background\n",
    "bkg_img = cv2.blur(avg_img,(2,2)) #blurs the averaged image using a 25x25 square on each pixel.  \n",
    "#Helps ensure that the particles themselves are less likely to be counted in the background\n",
    "\n",
    "diff_video = []\n",
    "display_video = []\n",
    "\n",
    "# The following code is for gamma correction for humans, since it can be really hard to see the particles after background subtraction.\n",
    "gamma = 0.2\n",
    "inv_gamma = 1.0 / gamma\n",
    "table = np.array([((i / 255.0) ** inv_gamma) * 255\n",
    "    for i in np.arange(0, 256)]).astype(\"uint8\")\n",
    "\n",
    "for item in imgs:\n",
    "    diff_video.append(cv2.bitwise_not(cv2.absdiff(item, bkg_img)))\n",
    "    display_video.append(cv2.LUT(cv2.bitwise_not(cv2.absdiff(item, bkg_img)), table))\n",
    "print(\"Background Generated\")\n",
    "\n",
    "lazy_subtraction = cv2.bitwise_not(cv2.absdiff(imgs[0], avg_img))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "from ipywidgets import Button, GridBox, Layout, ButtonStyle\n",
    "\n",
    "up_control = widgets.Button(description=\"🠉\",layout=Layout(width='auto',grid_area='up'))\n",
    "down_control = widgets.Button(description='🠋',layout=Layout(width='auto',grid_area='down'))\n",
    "left_control = widgets.Button(description='🠈',layout=Layout(width='auto',grid_area='left'))\n",
    "right_control = widgets.Button(description='🠊',layout=Layout(width='auto',grid_area='right'))\n",
    "deselect = widgets.Button(description='⊘',layout=Layout(width='auto',grid_area='deselect'))\n",
    "next_frame = widgets.Button(description='⇒',layout=Layout(width='auto',grid_area='n_frame'))\n",
    "prev_frame = widgets.Button(description='⇐',layout=Layout(width='auto',grid_area='p_frame'))\n",
    "\n",
    "\n",
    "arrow_control = GridBox(children = [up_control,down_control,left_control,right_control,deselect,next_frame,prev_frame],\n",
    "       layout=Layout(\n",
    "       width='100px',\n",
    "       height='200px',\n",
    "       grid_template_rows = \"100px auto auto auto\",\n",
    "       grid_template_columns = \"33% 33% 33%\",\n",
    "       grid_template_areas = '''\n",
    "       \" . . .\"\n",
    "       \". up .\"\n",
    "       \"left deselect right\"\n",
    "       \"p_frame down n_frame\"\n",
    "       ''')\n",
    "       )\n",
    "\n",
    "out = widgets.Output()\n",
    "\n",
    "select_add = widgets.Button(description=\"Add Points\")\n",
    "select_modify = widgets.Button(description=\"Modify Points\")\n",
    "status = widgets.Label(value=\"foo\")\n",
    "\n",
    "show_points_selector = widgets.ToggleButton(description = \"Show Points\",value=True)\n",
    "\n",
    "control_panel = widgets.VBox(children=[arrow_control,select_add,select_modify,show_points_selector ,status])\n",
    "\n",
    "image_to_view = display_video\n",
    "\n",
    "all_widgets = widgets.HBox(children=[control_panel, out])\n",
    "display(all_widgets)\n",
    "\n",
    "control_state = \"Add\"\n",
    "click_to_advance = True\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from ipywidgets import widgets\n",
    "%matplotlib notebook\n",
    "image_index = 0\n",
    "\n",
    "with out:\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_title('click on points')\n",
    "    ax.imshow(image_to_view[image_index])\n",
    "    plot_points = ax.plot([], 'o', picker=5,alpha = 1)  # 5 points tolerance\n",
    "    global highlight\n",
    "    highlight, = ax.plot([], [], 'o', color='orange')\n",
    "\n",
    "\n",
    "selected = None\n",
    "\n",
    "def onpick(event):\n",
    "    this_point = event.artist\n",
    "    x_value = this_point.get_xdata()\n",
    "    y_value = this_point.get_ydata()\n",
    "    ind = event.ind\n",
    "    global selected\n",
    "    selected = ind\n",
    "    highlight.set_data(x_value[ind][0],y_value[ind][0])\n",
    "    fig.canvas.draw_idle()\n",
    "    plt.gca().set_title(\"{},{}\".format(x_value[ind][0],y_value[ind][0]))\n",
    "    #plt.gca().set_title(\"Click at {:.2f}/{:.2f}\\non {:s}\".format(event.mouseevent.x, event.mouseevent.y, event.artist.__repr__()))\n",
    "    plt.gcf().canvas.draw_idle() \n",
    "       \n",
    "   \n",
    "    \n",
    "def onclick(event):\n",
    "    global status\n",
    "    status.value = \"click\"\n",
    "    global control_state\n",
    "    if(control_state == \"Add\"):\n",
    "        newx = np.append(ax.lines[0].get_xdata(),event.xdata)\n",
    "        newy = np.append(ax.lines[0].get_ydata(),event.ydata)\n",
    "        status.value = \"clack\"\n",
    "        ax.lines[0].set_data(newx,newy)\n",
    "        if click_to_advance:\n",
    "            global image_index\n",
    "            image_index = (image_index+1)%len(image_to_view)\n",
    "            ax.imshow(image_to_view[image_index])\n",
    "        ax.set_title(\"Frame {} / {}\".format(image_index,len(image_to_view)))\n",
    "        plt.canvas.draw_idle()\n",
    "        \n",
    "click_control = fig.canvas.mpl_connect('button_press_event', onclick)  \n",
    "\n",
    "def move_up(event):\n",
    "    a,b = ax.get_ylim()\n",
    "    delta = 0.005*(b-a)\n",
    "    highlight.set_data(highlight.get_xdata(),highlight.get_ydata()+delta)\n",
    "    fig.canvas.draw_idle()\n",
    "    if selected is not None:\n",
    "        foo = ax.lines[0].get_ydata()\n",
    "        mod = np.zeros_like(foo)\n",
    "        mod[selected] = delta\n",
    "        ax.lines[0].set_ydata(foo+mod)\n",
    "    \n",
    "def move_down(event):\n",
    "    a,b = ax.get_ylim()\n",
    "    delta = 0.005*(b-a)\n",
    "    highlight.set_data(highlight.get_xdata(),highlight.get_ydata()-delta)\n",
    "    fig.canvas.draw_idle() \n",
    "    if selected is not None:\n",
    "        foo = ax.lines[0].get_ydata()\n",
    "        mod = np.zeros_like(foo)\n",
    "        mod[selected] = delta\n",
    "        ax.lines[0].set_ydata(foo-mod)\n",
    "    \n",
    "def move_left(event):\n",
    "    a,b = ax.get_xlim()\n",
    "    delta = 0.005*(b-a)\n",
    "    highlight.set_data(highlight.get_xdata()-delta,highlight.get_ydata())\n",
    "    fig.canvas.draw_idle()\n",
    "    if selected is not None:\n",
    "        foo = ax.lines[0].get_xdata()\n",
    "        mod = np.zeros_like(foo)\n",
    "        mod[selected] = delta\n",
    "        ax.lines[0].set_xdata(foo-mod)\n",
    "    \n",
    "def move_right(event):\n",
    "    a,b = ax.get_xlim()\n",
    "    delta = 0.005*(b-a)\n",
    "    highlight.set_data(highlight.get_xdata()+delta,highlight.get_ydata())\n",
    "    fig.canvas.draw_idle()\n",
    "    if selected is not None:\n",
    "        foo = ax.lines[0].get_xdata()\n",
    "        mod = np.zeros_like(foo)\n",
    "        mod[selected] = delta\n",
    "        ax.lines[0].set_xdata(foo+mod)\n",
    "\n",
    "def remove_selection(event):\n",
    "    global selected\n",
    "    selected = None\n",
    "    highlight.set_data([],[])\n",
    "    \n",
    "def set_state(b):\n",
    "    global control_state\n",
    "    global pick_control\n",
    "    if b.description == \"Add Points\":\n",
    "        control_state = \"Add\"\n",
    "        highlight.set_data([],[])\n",
    "        fig.canvas.draw_idle()\n",
    "        try:\n",
    "            fig.canvas.mpl_disconnect(pick_control);\n",
    "        except:\n",
    "            pass\n",
    "    elif b.description == \"Modify Points\":\n",
    "        control_state = \"Modify\"\n",
    "        pick_control = fig.canvas.mpl_connect('pick_event', onpick);\n",
    "    else:\n",
    "        pass\n",
    "    global status\n",
    "    status.value = control_state\n",
    "    \n",
    "def set_next_frame(event):\n",
    "    global image_index\n",
    "    image_index = (image_index+1)%len(image_to_view)\n",
    "    ax.imshow(image_to_view[image_index])\n",
    "    fig.canvas.draw_idle()\n",
    "\n",
    "def set_prev_frame(event):\n",
    "    global image_index\n",
    "    image_index = (image_index-1)%len(image_to_view)\n",
    "    ax.imshow(image_to_view[image_index])\n",
    "    fig.canvas.draw_idle()    \n",
    "    \n",
    "def show_points_toggle(event):\n",
    "    status.value = str(event['new'])\n",
    "    if(event['new']):\n",
    "        ax.lines[0].set_alpha(1)\n",
    "    else:\n",
    "        ax.lines[0].set_alpha(0)\n",
    "    fig.canvas.draw_idle()\n",
    "    \n",
    "up_control.on_click(move_up)\n",
    "down_control.on_click(move_down)\n",
    "left_control.on_click(move_left)\n",
    "right_control.on_click(move_right)\n",
    "deselect.on_click(remove_selection)\n",
    "next_frame.on_click(set_next_frame)\n",
    "prev_frame.on_click(set_prev_frame)\n",
    "show_points_selector.observe(show_points_toggle, 'value')\n",
    "\n",
    "\n",
    "select_add.on_click(set_state)\n",
    "select_modify.on_click(set_state)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Particle = ax.lines[0].get_data()\n",
    "figg, (axx,ayy) = plt.subplots(2)\n",
    "axx.plot(np.arange(len(Particle[0])),Particle[0],label = \"position\",marker='o')\n",
    "ayy.plot(np.arange(len(Particle[1])),Particle[1],label = \"position\",marker='o')\n",
    "velocityx = np.gradient(Particle[0])\n",
    "axx.plot(np.arange(len(velocityx)),5*velocityx,label = \"velocity\",marker='o')\n",
    "velocityy = np.gradient(Particle[1])\n",
    "ayy.plot(np.arange(len(velocityy)),5*velocityy,label = \"velocity\",marker='o')\n",
    "axx.set_title(\"X coordinate\")\n",
    "ayy.set_title(\"Y coordinate\")\n",
    "axx.legend()\n",
    "ayy.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig2, (ax2,ay2) = plt.subplots(2)\n",
    "ax2.hist(velocityx)\n",
    "ay2.hist(velocityy)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "----\n",
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(velocityx,'\\n\\n',velocityy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import optimize\n",
    "# Generate numpy arrays for the histogram values and bin centers\n",
    "dxhist = np.histogram(velocityx, bins=10)\n",
    "errdx = np.sqrt(dxhist[0])\n",
    "errdx[errdx<1] = 1.4 # This sets the correct Poisson uncertainty when N = 0\n",
    "cbins = (dxhist[1][:-1]+dxhist[1][1:])/2\n",
    "\n",
    "# Define a Gaussian fit function and weighted residual\n",
    "def fitfunc(p,x):\n",
    "    return p[0]/(p[2]*np.sqrt(2*np.pi))*np.exp(-(x-p[1])**2/(2*p[2]**2))\n",
    "def residual(p, x, y, err):\n",
    "    return (fitfunc(p, x)-y)/err\n",
    "\n",
    "#Provide an initial guess for fit parameters and perform the fit\n",
    "p0 = [3000., 0., 1.]\n",
    "pf, cov, info, mesg, success = optimize.leastsq(residual, p0, args=(cbins, dxhist[0], errdx),\n",
    "                                                full_output=1)\n",
    "\n",
    "if cov is None:\n",
    "    print('Fit did not converge')\n",
    "    print('Success code:', success)\n",
    "    print(mesg)\n",
    "else:\n",
    "    chisq = sum(info['fvec']*info['fvec'])\n",
    "    dof = len(cbins) - len(pf)\n",
    "    pferr = [np.sqrt(cov[i,i]) for i in range(len(pf))]\n",
    "    print('Converged with chi-squared', chisq)\n",
    "    print('Number of degrees of freedom, dof =', dof)\n",
    "    print('Reduced chi-squared ', chisq/dof)\n",
    "    print('Inital guess values:')\n",
    "    print('  p0 =', p0)\n",
    "    print('Best fit values:')\n",
    "    print('  pf =', pf)\n",
    "    print('Uncertainties in the best fit values:')\n",
    "    print('  pferr =', pferr)\n",
    "    \n",
    "    #Let's save the standard deviation values for use later\n",
    "    sigma_x = pf[2]\n",
    "    dsigma_x = pferr[2]\n",
    "    \n",
    "    \n",
    "print()   \n",
    "dyhist = np.histogram(velocityy, bins=10)\n",
    "errdy = np.sqrt(dyhist[0])\n",
    "errdy[errdy<1] = 1\n",
    "cbins = (dyhist[1][:-1]+dyhist[1][1:])/2\n",
    "\n",
    "p0 = [3000., 0., 1.]\n",
    "pfy, cov, info, mesg, success = optimize.leastsq(residual, p0, args=(cbins, dyhist[0], errdy),\n",
    "                                                full_output=1)\n",
    "\n",
    "if cov is None:\n",
    "    print('Fit did not converge')\n",
    "    print('Success code:', success)\n",
    "    print(mesg)\n",
    "else:\n",
    "    chisq = sum(info['fvec']*info['fvec'])\n",
    "    dof = len(cbins) - len(pf)\n",
    "    pferr = [np.sqrt(cov[i,i]) for i in range(len(pf))]\n",
    "    print('Converged with chi-squared', chisq)\n",
    "    print('Number of degrees of freedom, dof =', dof)\n",
    "    print('Reduced chi-squared ', chisq/dof)\n",
    "    print('Inital guess values:')\n",
    "    print('  p0 =', p0)\n",
    "    print('Best fit values:')\n",
    "    print('  pf =', pf)\n",
    "    print('Uncertainties in the best fit values:')\n",
    "    print('  pferr =', pferr)\n",
    "    \n",
    "    #Let's save the standard deviation values for use later\n",
    "    sigma_y = pf[2]\n",
    "    dsigma_y = pferr[2]    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,(ax,ay) = plt.subplots(2)\n",
    "ax.errorbar(cbins, dxhist[0], errdx, fmt='k.', label='Data')\n",
    "\n",
    "X = np.linspace(min(cbins), max(cbins), 5000)\n",
    "ax.plot(X, fitfunc(pf, X), 'r-', label='Fit')\n",
    "\n",
    "ax.set_title('Horizontal displacement between consecutive frames')\n",
    "ax.set_xlabel('$\\Delta x$ (pixels)')\n",
    "ax.set_ylabel('Counts')\n",
    "ax.legend()\n",
    "\n",
    "textfit1 = '$C(x) = \\\\frac{N}{\\sigma\\sqrt{2\\pi}}\\\\exp\\\\left(-\\\\frac{(x-x_0)^2}{2\\sigma^2}\\\\right)$ \\n' \\\n",
    "               '$N = %.0f \\pm %.0f$ \\n' \\\n",
    "               '$x_0 = %.3f \\pm %.3f$ pixels\\n' \\\n",
    "               '$\\sigma= %.3f \\pm %.3f$ pixels \\n' \\\n",
    "               '$\\chi^2/N = % .2f$' \\\n",
    "                %(pf[0], pferr[0], pf[1], pferr[1], pf[2], pferr[2], chisq/dof)\n",
    "\n",
    "ax.text(0.02, .94, textfit1, transform=ax.transAxes, fontsize=12,\n",
    "             verticalalignment='top');\n",
    "\n",
    "ay.errorbar(cbins, dyhist[0], errdy, fmt='k.', label='Data')\n",
    "\n",
    "Y = np.linspace(min(cbins), max(cbins), 5000)\n",
    "ay.plot(Y, fitfunc(pfy, Y), 'r-', label='Fit')\n",
    "\n",
    "ay.set_title('Vertical displacement between consecutive frames')\n",
    "fig.tight_layout()\n",
    "ay.set_xlabel('$\\Delta y$ (pixels)')\n",
    "ay.set_ylabel('Counts')\n",
    "ay.legend()\n",
    "\n",
    "textfit1 = '$C(x) = \\\\frac{N}{\\sigma\\sqrt{2\\pi}}\\\\exp\\\\left(-\\\\frac{(x-x_0)^2}{2\\sigma^2}\\\\right)$ \\n' \\\n",
    "               '$N = %.0f \\pm %.0f$ \\n' \\\n",
    "               '$x_0 = %.3f \\pm %.3f$ pixels\\n' \\\n",
    "               '$\\sigma= %.3f \\pm %.3f$ pixels \\n' \\\n",
    "               '$\\chi^2/N = % .2f$' \\\n",
    "                %(pf[0], pferr[0], pf[1], pferr[1], pf[2], pferr[2], chisq/dof)\n",
    "\n",
    "#ax.set_xlim(left=-2.3)\n",
    "ay.text(0.02, .94, textfit1, transform=ax.transAxes, fontsize=12,\n",
    "             verticalalignment='top');\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Particle radius\n",
    "a = 0.96*10**(-6)/2  # m\n",
    "\n",
    "# Video frame rate and time between frames\n",
    "framerate = 10 # frames per second\n",
    "dt = 1/framerate # s\n",
    "\n",
    "# Absolute temperature and Boltzmann constant\n",
    "T = 297 # K\n",
    "k_B = 1.3806*10**(-23) # m^2 kg/(s^2 K)\n",
    "\n",
    "# Conversion factor between real length and pixel length\n",
    "conversion = 0.01*10**(-3)/61.0 # m/pixel\n",
    "\n",
    "# Viscosity of water (at temperature T)\n",
    "eta = 8.90*10**(-4)  # Pa s = kg/(s m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Take the average of your two sigma values, and convert to real length units\n",
    "sigma = (sigma_x + sigma_y)/2 * conversion                 # m\n",
    "dsigma = np.sqrt(dsigma_x**2 + dsigma_y**2)/2 * conversion # m\n",
    "\n",
    "#Compute the diffusion coefficient\n",
    "D = sigma**2/2/dt     # m^2/s\n",
    "dD = D*(dsigma/sigma) # m^2/s\n",
    "\n",
    "#Compute the expected diffusion coefficient\n",
    "D_predicted = k_B*T/(6*np.pi*eta*a)\n",
    "\n",
    "print('D_predicted = {:0.4f} x 10^(-13) m^2/s'.format((D_predicted/10**(-13))))\n",
    "print('D_exp = ({:0.4f} +/- {:0.4f}) x 10^(-13) m^2/s'.format(D/10**(-13), dD/10**(-13)))"
   ]
  }
 ],
 "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.8.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}