{ "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 }