{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e9f6a44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "import cv2\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision.transforms.functional as F\n",
    "from torchvision.utils import draw_bounding_boxes\n",
    "from torchvision.io import read_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e169a705",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"visualization/frames/\", exist_ok=True)\n",
    "os.makedirs(\"visualization/annotated_frames/\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c1e603",
   "metadata": {},
   "outputs": [],
   "source": [
    "def video2frame(\n",
    "    tfile,frame_width, frame_height, interval):\n",
    "    \"\"\"\n",
    "    Extract frame from video by interval\n",
    "    :param video_src_path: video src path\n",
    "    :param video:　video file name\n",
    "    :param frame_width:　frame width\n",
    "    :param frame_height:　frame height\n",
    "    :param interval:　interval for frame to extract\n",
    "    :return:　list of numpy.ndarray \n",
    "    \"\"\"\n",
    "    video_frames = []\n",
    "    cap = cv2.VideoCapture(tfile)\n",
    "    frame_index = 0\n",
    "    frame_count = 0\n",
    "    if cap.isOpened():\n",
    "        success = True\n",
    "    else:\n",
    "        success = False\n",
    "        print(\"Read failed!\")\n",
    "\n",
    "    while success:\n",
    "        success, frame = cap.read()\n",
    "\n",
    "        if frame_index % interval == 0:\n",
    "            print(\"---> Reading the %d frame:\" % frame_index, success)\n",
    "            resize_frame = cv2.resize(\n",
    "                frame, (frame_width, frame_height), interpolation=cv2.INTER_AREA\n",
    "            )\n",
    "            video_frames.append(resize_frame)\n",
    "            frame_count += 1\n",
    "\n",
    "        frame_index += 1\n",
    "\n",
    "    cap.release()\n",
    "    \n",
    "    print('Number of frames')\n",
    "    print(frame_count)\n",
    "    return video_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfb14262",
   "metadata": {},
   "outputs": [],
   "source": [
    "video_frames = video2frame('videos/ducks.mp4', 1024, 1024, 30)\n",
    "for i in range(len(video_frames)):\n",
    "    cv2.imwrite(f\"visualization/frames/image-{i}.jpg\", video_frames[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692fe844",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams[\"savefig.bbox\"] = 'tight'\n",
    "\n",
    "\n",
    "def save(imgs, img_num):\n",
    "    if not isinstance(imgs, list):\n",
    "        imgs = [imgs]\n",
    "    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)\n",
    "    for i, img in enumerate(imgs):\n",
    "        img = img.detach()\n",
    "        img = F.to_pil_image(img)\n",
    "        \n",
    "        axs[0, i].imshow(np.asarray(img))\n",
    "        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])\n",
    "    plt.savefig(f\"visualization/annotated_frames/{str(img_num).zfill(3)}.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b49141d",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"visualization/output.json\", \"r\") as read_file:\n",
    "    data = json.load(read_file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1878aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(data)):\n",
    "    score_threshold = .9\n",
    "    scores = torch.from_numpy(np.array(data[i]['scores']))\n",
    "    boxes = torch.from_numpy(np.array(data[i]['boxes']))\n",
    "\n",
    "    birds_with_boxes = [\n",
    "        draw_bounding_boxes(read_image(f'visualization/frames/image-{i}.jpg'), boxes=boxes[scores > score_threshold], width=10)\n",
    "     ]\n",
    "    save(birds_with_boxes, i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff5a0748",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "from PIL import Image\n",
    "\n",
    "# filepaths\n",
    "fp_in = \"visualization/annotated_frames/*.jpg\"\n",
    "fp_out = \"visualization/annotated_frames/birds.gif\"\n",
    "\n",
    "\n",
    "img, *imgs = [Image.open(f) for f in sorted(glob.glob(fp_in))]\n",
    "img.save(fp=fp_out, format='GIF', append_images=imgs,\n",
    "         save_all=True, duration=200, loop=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4ba837f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}