ffmpeg-python/examples/ffmpeg-numpy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact\n",
    "from matplotlib import pyplot as plt\n",
    "import ffmpeg\n",
    "import ipywidgets as widgets\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "probe = ffmpeg.probe('in.mp4')\n",
    "video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')\n",
    "width = int(video_info['width'])\n",
    "height = int(video_info['height'])\n",
    "num_frames = int(video_info['nb_frames'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5f63dc164956464c994ec58d86ee7cd9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='frame', max=209), Output()), _dom_classes=('widget-inter…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "out, err = (\n",
    "    ffmpeg\n",
    "    .input('in.mp4')\n",
    "    .output('pipe:', format='rawvideo', pix_fmt='rgb24')\n",
    "    .run(capture_stdout=True)\n",
    ")\n",
    "video = (\n",
    "    np\n",
    "    .frombuffer(out, np.uint8)\n",
    "    .reshape([-1, height, width, 3])\n",
    ")\n",
    "\n",
    "@interact(frame=(0, num_frames))\n",
    "def show_frame(frame=0):\n",
    "    plt.imshow(video[frame,:,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "84bcac52195f47f8854f09acd7666b84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Checkbox(value=True, description='enable_overlay'), Checkbox(value=True, description='en…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def extract_frame(stream, frame_num):\n",
    "    while isinstance(stream, ffmpeg.nodes.OutputStream):\n",
    "        stream = stream.node.incoming_edges[0].upstream_node.stream()\n",
    "    out, _ = (\n",
    "        stream\n",
    "        .filter_('select', 'gte(n,{})'.format(frame_num))\n",
    "        .output('pipe:', format='rawvideo', pix_fmt='rgb24', vframes=1)\n",
    "        .run(capture_stdout=True, capture_stderr=True)\n",
    "    )\n",
    "    return np.frombuffer(out, np.uint8).reshape([height, width, 3])\n",
    "\n",
    "\n",
    "def png_to_np(png_bytes):\n",
    "    buffer = BytesIO(png_bytes)\n",
    "    pil_image = Image.open(buffer)\n",
    "    return np.array(pil_image)\n",
    "    \n",
    "\n",
    "def build_graph(\n",
    "        enable_overlay, flip_overlay, enable_box, box_x, box_y,\n",
    "        thickness, color):\n",
    "\n",
    "    stream = ffmpeg.input('in.mp4')\n",
    "\n",
    "    if enable_overlay:\n",
    "        overlay = ffmpeg.input('overlay.png')\n",
    "        if flip_overlay:\n",
    "            overlay = overlay.hflip()\n",
    "        stream = stream.overlay(overlay)\n",
    "\n",
    "    if enable_box:\n",
    "        stream = stream.drawbox(\n",
    "            box_x, box_y, 120, 120, color=color, t=thickness)\n",
    "\n",
    "    return stream.output('out.mp4')\n",
    "\n",
    "\n",
    "def show_image(ax, stream, frame_num):\n",
    "    try:\n",
    "        image = extract_frame(stream, frame_num)\n",
    "        ax.imshow(image)\n",
    "        ax.axis('off')\n",
    "    except ffmpeg.Error as e:\n",
    "        print(e.stderr.decode())\n",
    "\n",
    "\n",
    "def show_graph(ax, stream, detail):\n",
    "    data = ffmpeg.view(stream, detail=detail, pipe=True)\n",
    "    image = png_to_np(data)\n",
    "    ax.imshow(image, aspect='equal', interpolation='hanning')\n",
    "    ax.set_xlim(0, 1100)\n",
    "    ax.axis('off')\n",
    "\n",
    "\n",
    "@interact(\n",
    "    frame_num=(0, num_frames),\n",
    "    box_x=(0, 200),\n",
    "    box_y=(0, 200),\n",
    "    thickness=(1, 40),\n",
    "    color=['red', 'green', 'magenta', 'blue'],\n",
    ")\n",
    "def f(\n",
    "        enable_overlay=True,\n",
    "        enable_box=True,\n",
    "        flip_overlay=True,\n",
    "        graph_detail=False,\n",
    "        frame_num=0,\n",
    "        box_x=50,\n",
    "        box_y=50,\n",
    "        thickness=5,\n",
    "        color='red'):\n",
    "\n",
    "    stream = build_graph(\n",
    "        enable_overlay,\n",
    "        flip_overlay,\n",
    "        enable_box,\n",
    "        box_x,\n",
    "        box_y,\n",
    "        thickness,\n",
    "        color\n",
    "    )\n",
    "\n",
    "    fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(15,4))\n",
    "    plt.tight_layout()\n",
    "    show_image(ax0, stream, frame_num)\n",
    "    show_graph(ax1, stream, graph_detail)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}