{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. and affiliates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Segment Anything Video (SA-V) Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook presents an example showing how to load and visualize the SA-V annotations.\n",
    "To run this notebook, you will need to have the following installed:\n",
    "- jupyter\n",
    "- matplotlib\n",
    "- opencv-python\n",
    "- numpy\n",
    "- pycocoevalcap\n",
    "- pandas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import pandas as pd\n",
    "from utils.sav_utils import SAVDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example video and annotations from SA-V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 483 frames decoded from example/sav_000001.mp4 (24fps).\n",
      "Videos are annotated every 4 frames. To align with the annotations, downsample the video to 121 frames.\n"
     ]
    }
   ],
   "source": [
    "sav_dataset = SAVDataset(sav_dir=\"example/\")\n",
    "frames, manual_annot, auto_annot = sav_dataset.get_frames_and_annotations(\"sav_000001\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the SA-V annotations in frame 0 - auto + manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 162076 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sav_dataset.visualize_annotation(\n",
    "    frames, manual_annot, auto_annot, \n",
    "    annotated_frame_id=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the SA-V annotations in another frame - auto + manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 172352 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sav_dataset.visualize_annotation(\n",
    "    frames, manual_annot, auto_annot, \n",
    "    annotated_frame_id=30,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the SA-V annotations in frame 0 - manual only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 163584 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sav_dataset.visualize_annotation(\n",
    "    frames, manual_annot, auto_annot,\n",
    "    annotated_frame_id=0,\n",
    "    show_auto=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Show the SA-V annotations in frame 0 - auto only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 162748 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sav_dataset.visualize_annotation(\n",
    "    frames, manual_annot, auto_annot,\n",
    "    annotated_frame_id=0,\n",
    "    show_manual=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Masklet annotations and Metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Manual annotations and metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>video_id</th>\n",
       "      <th>video_duration</th>\n",
       "      <th>video_frame_count</th>\n",
       "      <th>video_height</th>\n",
       "      <th>video_width</th>\n",
       "      <th>video_resolution</th>\n",
       "      <th>video_environment</th>\n",
       "      <th>video_split</th>\n",
       "      <th>masklet</th>\n",
       "      <th>masklet_id</th>\n",
       "      <th>masklet_size_rel</th>\n",
       "      <th>masklet_size_abs</th>\n",
       "      <th>masklet_size_bucket</th>\n",
       "      <th>masklet_visibility_changes</th>\n",
       "      <th>masklet_first_appeared_frame</th>\n",
       "      <th>masklet_frame_count</th>\n",
       "      <th>masklet_edited_frame_count</th>\n",
       "      <th>masklet_type</th>\n",
       "      <th>masklet_stability_score</th>\n",
       "      <th>masklet_num</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>sav_000001</td>\n",
       "      <td>20.125</td>\n",
       "      <td>483.0</td>\n",
       "      <td>848.0</td>\n",
       "      <td>480.0</td>\n",
       "      <td>407040.0</td>\n",
       "      <td>Indoor</td>\n",
       "      <td>train</td>\n",
       "      <td>[[{'size': [848, 480], 'counts': 'i\\Y4&lt;Qj05K4L...</td>\n",
       "      <td>[0, 1, 2, 3, 4]</td>\n",
       "      <td>[0.0035249812, 0.0946159778, 0.011285757, 0.00...</td>\n",
       "      <td>[1434.8083333333, 38512.4876033058, 4593.75454...</td>\n",
       "      <td>[medium, large, medium, medium, medium]</td>\n",
       "      <td>[2, 0, 10, 0, 0]</td>\n",
       "      <td>[0.0, 0.0, 0.0, 113.0, 0.0]</td>\n",
       "      <td>[121, 121, 121, 121, 121]</td>\n",
       "      <td>[41, 11, 22, 4, 115]</td>\n",
       "      <td>[manual, manual, manual, manual, manual]</td>\n",
       "      <td>[None, None, None, None, None]</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     video_id  video_duration  video_frame_count  video_height  video_width  \\\n",
       "0  sav_000001          20.125              483.0         848.0        480.0   \n",
       "\n",
       "   video_resolution video_environment video_split  \\\n",
       "0          407040.0            Indoor       train   \n",
       "\n",
       "                                             masklet       masklet_id  \\\n",
       "0  [[{'size': [848, 480], 'counts': 'i\\Y4<Qj05K4L...  [0, 1, 2, 3, 4]   \n",
       "\n",
       "                                    masklet_size_rel  \\\n",
       "0  [0.0035249812, 0.0946159778, 0.011285757, 0.00...   \n",
       "\n",
       "                                    masklet_size_abs  \\\n",
       "0  [1434.8083333333, 38512.4876033058, 4593.75454...   \n",
       "\n",
       "                       masklet_size_bucket masklet_visibility_changes  \\\n",
       "0  [medium, large, medium, medium, medium]           [2, 0, 10, 0, 0]   \n",
       "\n",
       "  masklet_first_appeared_frame        masklet_frame_count  \\\n",
       "0  [0.0, 0.0, 0.0, 113.0, 0.0]  [121, 121, 121, 121, 121]   \n",
       "\n",
       "  masklet_edited_frame_count                              masklet_type  \\\n",
       "0       [41, 11, 22, 4, 115]  [manual, manual, manual, manual, manual]   \n",
       "\n",
       "          masklet_stability_score  masklet_num  \n",
       "0  [None, None, None, None, None]            5  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame([manual_annot])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Auto annotations and metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>video_id</th>\n",
       "      <th>video_duration</th>\n",
       "      <th>video_frame_count</th>\n",
       "      <th>video_height</th>\n",
       "      <th>video_width</th>\n",
       "      <th>video_resolution</th>\n",
       "      <th>video_environment</th>\n",
       "      <th>video_split</th>\n",
       "      <th>masklet</th>\n",
       "      <th>masklet_id</th>\n",
       "      <th>masklet_size_rel</th>\n",
       "      <th>masklet_size_abs</th>\n",
       "      <th>masklet_size_bucket</th>\n",
       "      <th>masklet_visibility_changes</th>\n",
       "      <th>masklet_first_appeared_frame</th>\n",
       "      <th>masklet_frame_count</th>\n",
       "      <th>masklet_edited_frame_count</th>\n",
       "      <th>masklet_type</th>\n",
       "      <th>masklet_stability_score</th>\n",
       "      <th>masklet_num</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>sav_000001</td>\n",
       "      <td>20.125</td>\n",
       "      <td>483.0</td>\n",
       "      <td>848.0</td>\n",
       "      <td>480.0</td>\n",
       "      <td>407040.0</td>\n",
       "      <td>Indoor</td>\n",
       "      <td>train</td>\n",
       "      <td>[[{'size': [848, 480], 'counts': 'ka0e8ka001O1...</td>\n",
       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8]</td>\n",
       "      <td>[0.010841089678796047, 0.038489445267425544, 0...</td>\n",
       "      <td>[4412.757142857143, 15666.743801652892, 7663.1...</td>\n",
       "      <td>[medium, large, medium, large, medium, medium,...</td>\n",
       "      <td>[5, 0, 0, 0, 0, 0, 10, 3, 12]</td>\n",
       "      <td>[0, 0, 0, 0, 0, 0, 0, 0, 0]</td>\n",
       "      <td>[121, 121, 121, 121, 121, 121, 121, 121, 121]</td>\n",
       "      <td>[0, 0, 0, 0, 0, 0, 0, 0, 0]</td>\n",
       "      <td>[auto, auto, auto, auto, auto, auto, auto, aut...</td>\n",
       "      <td>[[1.0, 0.999616265296936, 1.0, 1.0, 1.0, 1.0, ...</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     video_id  video_duration  video_frame_count  video_height  video_width  \\\n",
       "0  sav_000001          20.125              483.0         848.0        480.0   \n",
       "\n",
       "   video_resolution video_environment video_split  \\\n",
       "0          407040.0            Indoor       train   \n",
       "\n",
       "                                             masklet  \\\n",
       "0  [[{'size': [848, 480], 'counts': 'ka0e8ka001O1...   \n",
       "\n",
       "                    masklet_id  \\\n",
       "0  [0, 1, 2, 3, 4, 5, 6, 7, 8]   \n",
       "\n",
       "                                    masklet_size_rel  \\\n",
       "0  [0.010841089678796047, 0.038489445267425544, 0...   \n",
       "\n",
       "                                    masklet_size_abs  \\\n",
       "0  [4412.757142857143, 15666.743801652892, 7663.1...   \n",
       "\n",
       "                                 masklet_size_bucket  \\\n",
       "0  [medium, large, medium, large, medium, medium,...   \n",
       "\n",
       "      masklet_visibility_changes masklet_first_appeared_frame  \\\n",
       "0  [5, 0, 0, 0, 0, 0, 10, 3, 12]  [0, 0, 0, 0, 0, 0, 0, 0, 0]   \n",
       "\n",
       "                             masklet_frame_count   masklet_edited_frame_count  \\\n",
       "0  [121, 121, 121, 121, 121, 121, 121, 121, 121]  [0, 0, 0, 0, 0, 0, 0, 0, 0]   \n",
       "\n",
       "                                        masklet_type  \\\n",
       "0  [auto, auto, auto, auto, auto, auto, auto, aut...   \n",
       "\n",
       "                             masklet_stability_score  masklet_num  \n",
       "0  [[1.0, 0.999616265296936, 1.0, 1.0, 1.0, 1.0, ...            9  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.DataFrame([auto_annot])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Video info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sav_000001 is 20.125 seconds long with 483.0 frames. The video resolution is 848.0 x 480.0.\n",
      "This video is captured in Indoor environment.\n"
     ]
    }
   ],
   "source": [
    "video_id = manual_annot[\"video_id\"]\n",
    "video_duration = manual_annot[\"video_duration\"]\n",
    "video_frame_count = manual_annot[\"video_frame_count\"]\n",
    "H = manual_annot[\"video_height\"]\n",
    "W = manual_annot[\"video_width\"]\n",
    "environment = manual_annot[\"video_environment\"]\n",
    "print(\n",
    "    f\"{video_id} is {video_duration} seconds long with {video_frame_count} frames. The video resolution is {H} x {W}.\"\n",
    ")\n",
    "print(f\"This video is captured in {environment} environment.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Masklet info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 5 manually labeled masklets and 9 automatically generated masklets.\n",
      "In SA-V, videos are annotated every 4 frames. Therefore, there are 121 frames being annotated.\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    f\"There are {manual_annot['masklet_num']} manually labeled masklets and {auto_annot['masklet_num']} automatically generated masklets.\"\n",
    ")\n",
    "print(\n",
    "    f\"In SA-V, videos are annotated every 4 frames. Therefore, there are {manual_annot['masklet_frame_count'][0]} frames being annotated.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'size': [848, 480],\n",
       " 'counts': 'i\\\\Y4<Qj05K4L4M3M2N3O000010O00001O00010O00000O101O0000000000000O010000000O10O10000O01000O100O010O1O1O1N2O1O0O2O1N20O01O1001N101O1N100O100O10O01O2O00000O1O10O1O11N1O1O100O1O2N1O1O3L4HTXk5'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the rle of the manual masklet with masklet_id=0 in frame 0\n",
    "masklet_id = 0\n",
    "annotated_frame_id = 0\n",
    "manual_annot[\"masklet\"][annotated_frame_id][masklet_id]\n",
    "# decode the rle using `mask_util.decode(rle)>0` to get the binary segmentation mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'size': [848, 480],\n",
       " 'counts': 'Q_T6S1Xh0X1eNY1[Od0E;M4N10000O101O00000000000000O0100000000001M2O1O1N3N1M4H8B?@e0POc1jMfZ[5'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the rle of the masklet with masklet_id=5 in frame 100\n",
    "masklet_id = 5\n",
    "annotated_frame_id = 100\n",
    "auto_annot[\"masklet\"][annotated_frame_id][masklet_id]\n",
    "# decode the rle using `mask_util.decode(rle)>0` to get the binary segmentation mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "onevision_ta_2_pseudo_labeling",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
