data_visualisation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deepmind/perception_test/blob/main/data_visualisation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "Up7CEnJriAVu"
      },
      "outputs": [],
      "source": [
        "# @title Prerequisites\n",
        "import colorsys\n",
        "import json\n",
        "import os\n",
        "import random\n",
        "from typing import Tuple, List, Dict\n",
        "import zipfile\n",
        "\n",
        "import cv2\n",
        "from google.colab.patches import cv2_imshow\n",
        "import imageio\n",
        "import matplotlib.pyplot as plt\n",
        "import moviepy.editor as mvp\n",
        "import numpy as np\n",
        "import requests\n",
        "from scipy.io import wavfile"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Utility Functions\n",
        "def download_and_unzip(url: str, destination: str):\n",
        "  \"\"\"Downloads and unzips a .zip file to a destination.\n",
        "\n",
        "  Downloads a file from the specified URL, saves it to the destination\n",
        "  directory, and then extracts its contents.\n",
        "\n",
        "  If the file is larger than 1GB, it will be downloaded in chunks,\n",
        "  and the download progress will be displayed.\n",
        "\n",
        "  Args:\n",
        "    url (str): The URL of the file to download.\n",
        "    destination (str): The destination directory to save the file and\n",
        "      extract its contents.\n",
        "  \"\"\"\n",
        "  if not os.path.exists(destination):\n",
        "    os.makedirs(destination)\n",
        "\n",
        "  filename = url.split('/')[-1]\n",
        "  file_path = os.path.join(destination, filename)\n",
        "\n",
        "  if os.path.exists(file_path):\n",
        "    print(f'{filename} already exists. Skipping download.')\n",
        "    return\n",
        "\n",
        "  response = requests.get(url, stream=True)\n",
        "  total_size = int(response.headers.get('content-length', 0))\n",
        "  gb = 1024*1024*1024\n",
        "\n",
        "  if total_size / gb > 1:\n",
        "    print(f'{filename} is larger than 1GB, downloading in chunks')\n",
        "    chunk_flag = True\n",
        "    chunk_size = int(total_size/100)\n",
        "  else:\n",
        "    chunk_flag = False\n",
        "    chunk_size = total_size\n",
        "\n",
        "  with open(file_path, 'wb') as file:\n",
        "    for chunk_idx, chunk in enumerate(\n",
        "        response.iter_content(chunk_size=chunk_size)):\n",
        "      if chunk:\n",
        "        if chunk_flag:\n",
        "          print(f\"\"\"{chunk_idx}% downloading\n",
        "          {round((chunk_idx*chunk_size)/gb, 1)}GB\n",
        "          / {round(total_size/gb, 1)}GB\"\"\")\n",
        "        file.write(chunk)\n",
        "  print(f\"'{filename}' downloaded successfully.\")\n",
        "\n",
        "  with zipfile.ZipFile(file_path, 'r') as zip_ref:\n",
        "    zip_ref.extractall(destination)\n",
        "  print(f\"'{filename}' extracted successfully.\")\n",
        "\n",
        "  os.remove(file_path)\n",
        "\n",
        "\n",
        "def load_db_json(db_file: str) -> Dict:\n",
        "  \"\"\"Loads a JSON file as a dictionary.\n",
        "\n",
        "  Args:\n",
        "    db_file (str): Path to the JSON file.\n",
        "\n",
        "  Returns:\n",
        "    Dict: Loaded JSON data as a dictionary.\n",
        "\n",
        "  Raises:\n",
        "    FileNotFoundError: If the specified file doesn't exist.\n",
        "    TypeError: If the JSON file is not formatted as a dictionary.\n",
        "  \"\"\"\n",
        "  if not os.path.isfile(db_file):\n",
        "    raise FileNotFoundError(f'No such file: {db_file}')\n",
        "\n",
        "  with open(db_file, 'r') as f:\n",
        "    db_file_dict = json.load(f)\n",
        "    if not isinstance(db_file_dict, dict):\n",
        "      raise TypeError('JSON file is not formatted as a dictionary.')\n",
        "    return db_file_dict\n",
        "\n",
        "\n",
        "def load_mp4_to_frames(filename: str) -> np.array:\n",
        "  \"\"\"Loads an MP4 video file and returns its frames as a NumPy array.\n",
        "\n",
        "  Args:\n",
        "    filename (str): Path to the MP4 video file.\n",
        "\n",
        "  Returns:\n",
        "    np.array: Frames of the video as a NumPy array.\n",
        "  \"\"\"\n",
        "  assert os.path.exists(filename), f'File {filename} does not exist.'\n",
        "  cap = cv2.VideoCapture(filename)\n",
        "\n",
        "  vid_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "  height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "  width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "\n",
        "  vid_frames = np.empty((vid_frames, height, width, 3), dtype=np.uint8)\n",
        "\n",
        "  idx = 0\n",
        "  while True:\n",
        "    ret, vid_frame = cap.read()\n",
        "    if not ret:\n",
        "      break\n",
        "\n",
        "    vid_frames[idx] = vid_frame\n",
        "    idx += 1\n",
        "\n",
        "  cap.release()\n",
        "  return vid_frames\n",
        "\n",
        "\n",
        "def get_video_frames(data_item: Dict, vid_path: str) -> np.array:\n",
        "  \"\"\"Loads frames of a video specified by an item dictionary.\n",
        "\n",
        "  Assumes format of annotations used in the Perception Test Dataset.\n",
        "\n",
        "  Args:\n",
        "  \tdata_item (Dict): Item from dataset containing metadata.\n",
        "    vid_path (str): Path to the directory containing videos.\n",
        "\n",
        "  Returns:\n",
        "    np.array: Frames of the video as a NumPy array.\n",
        "  \"\"\"\n",
        "  video_file_path = os.path.join(vid_path,\n",
        "                                 data_item['metadata']['video_id']) + '.mp4'\n",
        "  vid_frames = load_mp4_to_frames(video_file_path)\n",
        "  assert data_item['metadata']['num_frames'] == vid_frames.shape[0]\n",
        "  return vid_frames\n",
        "\n",
        "\n",
        "def get_audio(audio_item: dict, audio_path: str) -> np.array:\n",
        "  \"\"\"Loads audio specified by an item from the dataset.\n",
        "\n",
        "  Args:\n",
        "    audio_item (dict): Item from dataset containing metadata.\n",
        "    audio_path (str): Path to the directory containing audios.\n",
        "\n",
        "  Returns:\n",
        "    np.array: Audio data as a NumPy array.\n",
        "  \"\"\"\n",
        "  audio_file_path = os.path.join(audio_path,\n",
        "  \t\t\t\t\t\t\t\t\taudio_item['metadata']['video_id']) + '.wav'\n",
        "  sample_rate, audio = wavfile.read(audio_file_path)\n",
        "\n",
        "  assert audio_item['metadata']['audio_samples'] == audio.shape[0]\n",
        "  assert audio_item['metadata']['audio_sample_rate'] == sample_rate\n",
        "\n",
        "  return audio.astype(np.float32)"
      ],
      "metadata": {
        "id": "1q_JZ0ueikxp"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Download Dataset Sample\n",
        "data_path = './data/'\n",
        "video_path = './data/videos/'\n",
        "\n",
        "# sample annotations and videos the visualise the annotations later\n",
        "sample_annot_url = 'https://storage.googleapis.com/dm-perception-test/zip_data/sample_annotations.zip'\n",
        "download_and_unzip(sample_annot_url, data_path)\n",
        "\n",
        "sample_videos_url = 'https://storage.googleapis.com/dm-perception-test/zip_data/sample_videos.zip'\n",
        "download_and_unzip(sample_videos_url, data_path)\n",
        "\n",
        "db_json_path = './data/sample.json'\n",
        "db_dict = load_db_json(db_json_path)"
      ],
      "metadata": {
        "id": "fymsuk6HbPsQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualisation functions\n",
        "def get_colors(num_colors: int) -> Tuple[int, int, int]:\n",
        "  \"\"\"Generate random colormaps for visualizing different objects and points.\n",
        "\n",
        "  Args:\n",
        "    num_colors (int): The number of colors to generate.\n",
        "\n",
        "  Returns:\n",
        "    Tuple[int, int, int]: A tuple of RGB values representing the\n",
        "      generated colors.\n",
        "  \"\"\"\n",
        "  colors = []\n",
        "  for j in np.arange(0., 360., 360. / num_colors):\n",
        "    hue = j / 360.\n",
        "    lightness = (50 + np.random.rand() * 10) / 100.\n",
        "    saturation = (90 + np.random.rand() * 10) / 100.\n",
        "    color = colorsys.hls_to_rgb(hue, lightness, saturation)\n",
        "    color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))\n",
        "    colors.append(color)\n",
        "  random.seed(0)\n",
        "  random.shuffle(colors)\n",
        "  return colors\n",
        "\n",
        "\n",
        "def display_video(vid_frames: np.array, fps: int = 30):\n",
        "  \"\"\"Create and display temporary video from numpy array frames.\n",
        "\n",
        "  Args:\n",
        "    vid_frames: (np.array): The frames of the video as a\n",
        "    \tnumpy array. Format of frames should be:\n",
        "    \t(num_frames, height, width, channels)\n",
        "    fps (int): Frames per second for the video playback. Default is 30.\n",
        "  \"\"\"\n",
        "  kwargs = {'macro_block_size': None}\n",
        "  imageio.mimwrite('tmp_video_display.mp4',\n",
        "                   vid_frames[:, :, :, ::-1], fps=fps, **kwargs)\n",
        "  display(mvp.ipython_display('tmp_video_display.mp4'))\n",
        "\n",
        "\n",
        "def display_frame(tmp_frame: np.array):\n",
        "  \"\"\"Display a frame, converting from RGB to BGR for cv2.\n",
        "\n",
        "  Args:\n",
        "    tmp_frame (np.array): The frame to be displayed.\n",
        "  \"\"\"\n",
        "  cv2_imshow(tmp_frame)\n",
        "\n",
        "\n",
        "def paint_box(video: np.array, track: Dict,\n",
        "\t\tcolor: Tuple[int, int, int] = (255, 0, 0),\n",
        "  \taddn_label: str = '') -> np.array:\n",
        "  \"\"\"Paint bounding box and label on video for a given track.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    track (Dict): The track information containing bounding box\n",
        "    and frame information, assumes Perception Test Dataset format.\n",
        "    color (Tuple[int, int, int]): The RGB color values for the bounding box.\n",
        "      Default is red (255, 0, 0).\n",
        "    addn_label (str): Additional label to be added to the track label.\n",
        "      Default is an empty string.\n",
        "\n",
        "  Returns:\n",
        "    np.array: The modified video frames with painted bounding box and\n",
        "      label.\n",
        "  \"\"\"\n",
        "  _, height, width, _ = video.shape\n",
        "  name = str(track['id']) + ' : ' + track['label'] + addn_label\n",
        "  bounding_boxes = np.array(track['bounding_boxes'])\n",
        "\n",
        "  for box, frame_id in zip(bounding_boxes, track['frame_ids']):\n",
        "    frame = np.array(video[frame_id])\n",
        "    x1 = int(round(box[0] * width))\n",
        "    y1 = int(round(box[1] * height))\n",
        "    x2 = int(round(box[2] * width))\n",
        "    y2 = int(round(box[3] * height))\n",
        "    frame = cv2.rectangle(frame, (x1, y1), (x2, y2),\n",
        "                          color=color, thickness=2)\n",
        "    frame = cv2.putText(frame, name, (x1, y1 + 20),\n",
        "                        cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2)\n",
        "    video[frame_id] = frame\n",
        "\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_boxes(video: np.array, tracks: List[Dict]) -> np.array:\n",
        "  \"\"\"Paint bounding boxes and labels on a video for multiple tracks.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    tracks (List): A list of track information,\n",
        "      where each track contains bounding box and frame information.\n",
        "\n",
        "  Returns:\n",
        "    np.array: The modified video frames with painted bounding boxes\n",
        "      and labels.\n",
        "  \"\"\"\n",
        "  for track_idx, track in enumerate(tracks):\n",
        "    video = paint_box(video, track, COLORS[track_idx])\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_point(video: np.array,\n",
        "  \ttrack: dict, color: tuple[int, int, int] = (255, 0, 0)) -> np.array:\n",
        "  \"\"\"Paints a single tracked point on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    track (dict): The track containing frame IDs and corresponding points.\n",
        "    color (tuple, optional): The color of the painted point.\n",
        "      Defaults to (255, 0, 0).\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted points.\n",
        "  \"\"\"\n",
        "  _, height, width, _ = video.shape\n",
        "  for idx, frame_id in enumerate(track['frame_ids']):\n",
        "    vid_frame = video[frame_id]\n",
        "    y = int(round(track['points'][0][idx] * height))\n",
        "    x = int(round(track['points'][1][idx] * width))\n",
        "    vid_frame = cv2.circle(vid_frame, (x, y),\n",
        "    \t\t\t\t\t\tradius=10, color=color, thickness=-1)\n",
        "    video[frame_id] = vid_frame\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_points(video: np.array, tracks: List[dict]) -> np.array:\n",
        "  \"\"\"Paints multiple tracked points on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    tracks (List[dict]): The list of tracks containing\n",
        "      frame IDs and corresponding points.\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted points.\n",
        "  \"\"\"\n",
        "  for idx, track in enumerate(tracks):\n",
        "    video = paint_point(video, track, COLORS[idx])\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_sound(video: np.array,\n",
        "    vid_sound: dict, vid_frames: np.array,\n",
        "    color: tuple[int, int, int] = (0, 0, 255)) -> np.array:\n",
        "  \"\"\"Paints a sound label on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    vid_sound (dict): The sound containing the label,\n",
        "      frame IDs, and visibility.\n",
        "    vid_frames (np.array): The array to keep track of\n",
        "      the number of labels on each frame.\n",
        "    color (tuple, optional): The color of the painted label.\n",
        "      Defaults to (0, 0, 255).\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted labels.\n",
        "  \"\"\"\n",
        "  _, height, width, _ = video.shape\n",
        "  name = f\"\"\"Sound: {vid_sound[\"label\"]}\n",
        "  is_visible: {str(bool(vid_sound[\"is_visible\"]))}\"\"\"\n",
        "  [start_frame, end_frame] = vid_sound['frame_ids']\n",
        "  for frame_id in range(start_frame, end_frame):\n",
        "    vid_frame = np.array(video[frame_id])\n",
        "    y1 = int(round(0.9 * height) - (40 * vid_frames[frame_id]))\n",
        "    x1 = int(round(0.05 * width))\n",
        "\n",
        "    vid_frame = cv2.putText(vid_frame, name, (x1, y1),\n",
        "                cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 2)\n",
        "    video[frame_id] = vid_frame\n",
        "    vid_frames[frame_id] += 1\n",
        "\n",
        "\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_action(video: np.array, vid_action: dict,\n",
        "\t\tvid_frames: np.array, color: tuple[int, int, int] = (0, 255, 0),\n",
        "    ) -> np.array:\n",
        "  \"\"\"Paints an action label on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    vid_action (dict): The action containing the label and frame IDs.\n",
        "    vid_frames (np.array): The array to keep track\n",
        "      of the number of labels on each frame.\n",
        "    color (tuple, optional): The color of the painted label.\n",
        "      Defaults to (0, 255, 0).\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted labels.\n",
        "  \"\"\"\n",
        "  _, height, width, _ = video.shape\n",
        "  name = f\"\"\"Action: {vid_action[\"label\"]}\"\"\"\n",
        "  [start_frame, end_frame] = vid_action['frame_ids']\n",
        "  for frame_id in range(start_frame, end_frame):\n",
        "    vid_frame = np.array(video[frame_id])\n",
        "    y1 = int(round(0.9 * height) - (40 * vid_frames[frame_id]))\n",
        "    x1 = int(round(0.05 * width))\n",
        "\n",
        "    vid_frame = cv2.putText(vid_frame, name, (x1, y1),\n",
        "                cv2.FONT_HERSHEY_SIMPLEX, 1.5, color, 2)\n",
        "    video[frame_id] = vid_frame\n",
        "    vid_frames[frame_id] += 1\n",
        "\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_actions(video: np.array,\n",
        "    vid_actions: List[dict], vid_frames: np.array) -> np.array:\n",
        "  \"\"\"Paints multiple action labels on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    vid_actions (List[dict]): The list of actions containing\n",
        "      the labels and frame IDs.\n",
        "    vid_frames (np.array): The array to keep track\n",
        "      of the number of labels on each frame.\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted labels.\n",
        "  \"\"\"\n",
        "  for vid_action in vid_actions:\n",
        "    video = paint_action(video, vid_action, vid_frames)\n",
        "  return video\n",
        "\n",
        "\n",
        "def paint_sounds(video: np.array,\n",
        "  \tvid_sounds: List[dict], vid_frames: np.array) -> np.array:\n",
        "  \"\"\"Paints multiple sound labels on each frame of a video.\n",
        "\n",
        "  Args:\n",
        "    video (np.array): The video frames as a numpy array.\n",
        "    vid_sounds (List[dict]): The list of sounds containing the labels,\n",
        "      frame IDs, and visibility.\n",
        "    vid_frames (np.array): The array to keep track of the\n",
        "      number of labels on each frame.\n",
        "\n",
        "  Returns:\n",
        "    np.array: The video frames with painted labels.\n",
        "  \"\"\"\n",
        "  for sound in vid_sounds:\n",
        "    video = paint_sound(video, sound, vid_frames)\n",
        "  return video\n",
        "\n",
        "\n",
        "def get_answer_tracks(ex_data: dict, goq_ids: List) -> List[dict]:\n",
        "  \"\"\"Filters and retrieves object tracks based on the given object ids.\n",
        "\n",
        "  Args:\n",
        "    ex_data (dict): The data containing object tracking information.\n",
        "    goq_ids (List): The list of IDs to filter tracks.\n",
        "\n",
        "  Returns:\n",
        "    List[dict]: The filtered tracks matching the goq_ids.\n",
        "  \"\"\"\n",
        "  goq_tracks = []\n",
        "  for track in ex_data['object_tracking']:\n",
        "    if track['id'] in goq_ids:\n",
        "      goq_tracks.append(track)\n",
        "  return goq_tracks"
      ],
      "metadata": {
        "id": "DhLSUAVdUGx2"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Show Example Annotations\n",
        "video_id = list(db_dict.keys())[6]\n",
        "example_data = db_dict[video_id]\n",
        "\n",
        "print('---------------------------------------------------------------------')\n",
        "print('Tasks annotated for this video: ')\n",
        "for k, v in example_data.items():\n",
        "  if v:\n",
        "    print(f'{k} - available: yes - annotations: {len(v)}')\n",
        "  else:\n",
        "    print(f'{k} - available: no')\n",
        "print('---------------------------------------------------------------------')\n",
        "print('Video Metadata')\n",
        "print('---------------------------------------------------------------------')\n",
        "for k, v in example_data['metadata'].items():\n",
        "  print(f'{k} : {v}')\n",
        "print('---------------------------------------------------------------------')\n",
        "print('Object Tracking data')\n",
        "print('---------------------------------------------------------------------')\n",
        "for k, v in example_data['object_tracking'][0].items():\n",
        "  print(f'{k} : {v}')\n",
        "print('---------------------------------------------------------------------')\n",
        "print('Multiple-Choice VQA')\n",
        "print('---------------------------------------------------------------------')\n",
        "for k, v in example_data['mc_question'][0].items():\n",
        "  print(f'{k} : {v}')\n",
        "print('---------------------------------------------------------------------')"
      ],
      "metadata": {
        "id": "nSbbsIDQI6Uf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Object Tracks\n",
        "if example_data['object_tracking']:\n",
        "  frames = get_video_frames(example_data, video_path)\n",
        "\n",
        "  COLORS = get_colors(num_colors=100)\n",
        "  show_all_tracks = True  # @param {type: \"boolean\"}\n",
        "  show_track = 2  # @param {type: \"integer\"}\n",
        "\n",
        "  if show_all_tracks:\n",
        "    frames = paint_boxes(frames, example_data['object_tracking'])\n",
        "  else:\n",
        "    frames = paint_box(frames, example_data['object_tracking'][show_track])\n",
        "\n",
        "  annotated_frames = []\n",
        "  for frame_idx in example_data['object_tracking'][0]['frame_ids']:\n",
        "    annotated_frames.append(frames[frame_idx])\n",
        "\n",
        "  annotated_frames = np.array(annotated_frames)\n",
        "  display_video(annotated_frames, 1)\n",
        "  del frames"
      ],
      "metadata": {
        "id": "hPJkHKWPEvff"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Point Tracks\n",
        "if example_data['point_tracking']:\n",
        "  frames = get_video_frames(example_data, video_path)\n",
        "  COLORS = get_colors(num_colors=100)\n",
        "  frames = paint_points(frames, example_data['point_tracking'])\n",
        "  display_video(frames, example_data['metadata']['frame_rate'])\n",
        "  del frames"
      ],
      "metadata": {
        "id": "cJESCRiMr9i7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Action Segments\n",
        "if example_data['action_localisation']:\n",
        "  frames = get_video_frames(example_data, video_path)\n",
        "  labelled_frames = np.zeros(frames.shape[0])\n",
        "  frames = paint_actions(frames, example_data['action_localisation'],\n",
        "  \t\t\t\t labelled_frames)\n",
        "  display_video(frames, example_data['metadata']['frame_rate'])\n",
        "  del frames"
      ],
      "metadata": {
        "id": "_lB8XcG86YHQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Plotting Action Segments\n",
        "if example_data['action_localisation']:\n",
        "  frames = get_video_frames(example_data, video_path)[:,:,:,::-1]\n",
        "\n",
        "  action_labels = []\n",
        "  action_start_times = []\n",
        "  action_end_times = []\n",
        "\n",
        "  for action in example_data['action_localisation']:\n",
        "    action_labels.append(action['label'])\n",
        "    action_start_times.append(action['timestamps'][0]/1e6)\n",
        "    action_end_times.append(action['timestamps'][1]/1e6)\n",
        "\n",
        "  action_start_times = np.array(action_start_times)\n",
        "  action_end_times = np.array(action_end_times)\n",
        "\n",
        "  plt.figure(figsize=(20, 15))\n",
        "  # Strip of frames\n",
        "  plt.subplot(4, 1, 2)\n",
        "  plt.title('Video Frames')\n",
        "  f_size = frames[0].shape\n",
        "  small = tuple(reversed((np.array(f_size[:2]) / 4).astype(int)))\n",
        "  strip = None\n",
        "  num_frames = example_data['metadata']['num_frames']\n",
        "  for i in range(0, num_frames, int(num_frames/4)):\n",
        "    frame = cv2.resize(frames[i], small)\n",
        "    if strip is None:\n",
        "      strip = np.array(frame)\n",
        "    else:\n",
        "      strip = np.concatenate([strip, frame], axis=1)\n",
        "    plt.imshow(strip)\n",
        "\n",
        "  del frames\n",
        "\n",
        "  plt.subplot(4, 1, 3)\n",
        "  plt.title('Action Events')\n",
        "  plt.barh(range(len(action_start_times)),\n",
        "           action_end_times-action_start_times,\n",
        "           left=action_start_times)\n",
        "  plt.yticks(range(len(action_start_times)), action_labels)\n",
        "\n",
        "  plt.show()"
      ],
      "metadata": {
        "id": "NPEft_TFwcnK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Sound Segments\n",
        "if example_data['sound_localisation']:\n",
        "  frames = get_video_frames(example_data, video_path)\n",
        "  labelled_frames = np.zeros(frames.shape[0])\n",
        "  frames = paint_sounds(frames, example_data['sound_localisation'],\n",
        "  \t\t\t\t labelled_frames)\n",
        "  display_video(frames, example_data['metadata']['frame_rate'])\n",
        "  del frames"
      ],
      "metadata": {
        "id": "NeeAYsGXuByb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Plotting Sound Segments\n",
        "if example_data['sound_localisation']:\n",
        "  frames = get_video_frames(example_data, video_path)[:,:,:,::-1]\n",
        "\n",
        "  audio_labels = []\n",
        "  audio_start_times = []\n",
        "  audio_end_times = []\n",
        "  for audio_event in example_data['sound_localisation']:\n",
        "    audio_labels.append(audio_event['label'])\n",
        "    audio_start_times.append(audio_event['timestamps'][0]/1e6)\n",
        "    audio_end_times.append(audio_event['timestamps'][1]/1e6)\n",
        "\n",
        "  audio_start_times = np.array(audio_start_times)\n",
        "  audio_end_times = np.array(audio_end_times)\n",
        "\n",
        "  plt.figure(figsize=(20, 15))\n",
        "  # Strip of frames\n",
        "  plt.subplot(4, 1, 2)\n",
        "  plt.title('Video Frames')\n",
        "  f_size = frames[0].shape\n",
        "  small = tuple(reversed((np.array(f_size[:2]) / 4).astype(int)))\n",
        "  strip = None\n",
        "  num_frames = example_data['metadata']['num_frames']\n",
        "  for i in range(0, num_frames, int(num_frames/4)):\n",
        "    frame = cv2.resize(frames[i], small)\n",
        "    if strip is None:\n",
        "      strip = np.array(frame)\n",
        "    else:\n",
        "      strip = np.concatenate([strip, frame], axis=1)\n",
        "    plt.imshow(strip)\n",
        "\n",
        "  del frames\n",
        "\n",
        "  # Plot audio events\n",
        "  plt.subplot(4, 1, 3)\n",
        "  plt.title('Audio Events')\n",
        "  plt.barh(range(len(audio_start_times)),\n",
        "           audio_end_times-audio_start_times,\n",
        "           left=audio_start_times)\n",
        "  plt.yticks(range(len(audio_start_times)), audio_labels)\n",
        "\n",
        "  plt.show()"
      ],
      "metadata": {
        "id": "nUZW6HvWwonR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Multiple-Choice Video Question-Answering Annotations\n",
        "if example_data['mc_question']:\n",
        "  for question in example_data['mc_question']:\n",
        "    print('---------------------------------')\n",
        "    print('Question: ', question['question'])\n",
        "    print('Options: ', question['options'])\n",
        "    print('Answer ID: ', question['answer_id'])\n",
        "    print('Answer: ', question['options'][question['answer_id']])\n",
        "    print('Question info: ')\n",
        "    print('Reasoning: ', question['reasoning'])\n",
        "    print('Tag: ', question['tag'])\n",
        "    print('area: ', question['area'])\n",
        "    print('---------------------------------')"
      ],
      "metadata": {
        "id": "RR9Mm_0T6iBn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Visualising Grounded Video Question-Answering\n",
        "\n",
        "# loading an example that has grounded question annotations\n",
        "video_id = list(db_dict.keys())[7]\n",
        "example_data = db_dict[video_id]\n",
        "\n",
        "if example_data['grounded_question']:\n",
        "  question = example_data['grounded_question'][0]\n",
        "  print('---------------------------------')\n",
        "  print('Question: ', question['question'])\n",
        "  print('Answer IDs: ', question['answers'])\n",
        "  print('Question info: ')\n",
        "  print('Reasoning: ', question['reasoning'])\n",
        "  print('area: ', question['area'])\n",
        "  print('---------------------------------')\n",
        "\n",
        "  frames = get_video_frames(example_data, video_path)\n",
        "  answer_tracks = get_answer_tracks(example_data, question['answers'])\n",
        "  frames = paint_boxes(frames, answer_tracks)\n",
        "\n",
        "  annotated_frames = []\n",
        "  for frame_idx in answer_tracks[0]['frame_ids']:\n",
        "    annotated_frames.append(frames[frame_idx])\n",
        "\n",
        "  annotated_frames = np.array(annotated_frames)\n",
        "  display_video(annotated_frames, 1)\n",
        "  del frames"
      ],
      "metadata": {
        "id": "3iFMOZgShAqn"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}