dsp-85746.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://dsp.stackexchange.com/questions/85746/how-can-i-analyse-the-motion-of-plasma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2 as cv\n",
    "import imageio.v2 as iio2\n",
    "import imageio.v3 as iio3\n",
    "import matplotlib.pyplot as plt\n",
    "import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"Plasma_Motion_numkmi.mp4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_map(x, xmin, xmax, ymin, ymax):\n",
    "    return (x - xmin) / (xmax - xmin) * (ymax - ymin) + ymin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_onto(canvas, values):\n",
    "    (h,w) = canvas.shape[:2]\n",
    "    N = len(values)\n",
    "    # just markers for now (circle)\n",
    "\n",
    "    ii = np.arange(len(values))\n",
    "    xx = linear_map(ii+0.5, 0, N, 0, w-1)\n",
    "    yy = linear_map(values, 0.0, 1.0, h-1, 0)\n",
    "\n",
    "    shift = 4\n",
    "\n",
    "    for pt in np.vstack((xx,yy)).T:\n",
    "        (x,y) = pt\n",
    "        if 0 <= y < h:\n",
    "            cv.circle(\n",
    "                img=canvas,\n",
    "                center=(pt * 2**shift).round().astype(int), radius=3 * 2**shift,\n",
    "                color=(255,255,255), thickness=cv.FILLED, lineType=cv.LINE_AA, shift=shift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'audio_codec': 'aac',\n",
      " 'codec': 'h264',\n",
      " 'duration': 45.91,\n",
      " 'ffmpeg_version': '4.2.2 built with gcc 9.2.1 (GCC) 20200122',\n",
      " 'fps': 30.0,\n",
      " 'nframes': inf,\n",
      " 'pix_fmt': 'yuv420p',\n",
      " 'plugin': 'ffmpeg',\n",
      " 'rotate': 0,\n",
      " 'size': (1920, 1080),\n",
      " 'source_size': (1920, 1080)}\n"
     ]
    }
   ],
   "source": [
    "meta = iio3.immeta(fname)\n",
    "pprint.pprint(dict(meta))\n",
    "\n",
    "(width, height) = meta['size']"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: this loop will get progressively slower ***only*** because `cv.circle()` calls cost something... and the number of circles increases as the loop iterates. The actual background subtraction always costs the same, proportional to frame size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Chris\\AppData\\Local\\Temp\\ipykernel_21524\\255809645.py:24: RuntimeWarning: divide by zero encountered in log10\n",
      "  plot_onto(frame, linear_map(np.log10(history), 0, -3, 1.0, 0.0))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n"
     ]
    }
   ],
   "source": [
    "with iio3.imopen(\"result.mov\", 'w', plugin='pyav') as outvid:\n",
    "    invid = iio3.imiter(fname, plugin=\"pyav\", format=\"gray\")\n",
    "    outvid.init_video_stream(\"h264\", fps=meta[\"fps\"])\n",
    "\n",
    "    bs = cv.createBackgroundSubtractorMOG2()\n",
    "\n",
    "    history = [] # per frame: foreground fraction\n",
    "\n",
    "    for i, frame in enumerate(invid):\n",
    "        fgmask = bs.apply(frame, learningRate=0.1)\n",
    "\n",
    "        fgfraction = np.count_nonzero(fgmask) / fgmask.size\n",
    "        history.append(fgfraction)\n",
    "\n",
    "        fgmask = cv.pyrDown(fgmask)\n",
    "        frame = cv.pyrDown(frame)\n",
    "\n",
    "        cv.putText(img=frame,\n",
    "            text=f\"t = {i/meta['fps']:5.2f}, fg = {fgfraction:.6f}\",\n",
    "            org=(10, height//2-10),\n",
    "            fontFace=cv.FONT_HERSHEY_SIMPLEX, fontScale=0.8,\n",
    "            color=(255,255,255), thickness=1)\n",
    "\n",
    "        plot_onto(frame, linear_map(np.log10(history), 0, -3, 1.0, 0.0))\n",
    "\n",
    "        composite = np.vstack([frame, fgmask])\n",
    "        outvid.write_frame(composite[:,:,None])\n",
    "\n",
    "        cv.imshow(\"vid\", frame)\n",
    "        #cv.setWindowTitle(\"vid\", f\"t = {i/fps:.2f}\")\n",
    "\n",
    "        cv.imshow(\"mask\", fgmask)\n",
    "\n",
    "        if cv.pollKey() in (13, 27):\n",
    "            break\n",
    "\n",
    "    else:\n",
    "        print(\"done\")\n",
    "        cv.waitKey()\n",
    "\n",
    "    cv.destroyAllWindows()"
   ]
  }
 ],
 "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.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "369f2c481f4da34e4445cda3fffd2e751bd1c4d706f27375911949ba6bb62e1c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}