notebooks/Perception_Statistics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset Statistics for Perception Package Projects\n",
    "This example notebook shows how to use datasetinsights to load synthetic datasets generated from the [Perception package](https://github.com/Unity-Technologies/com.unity.perception) and visualize dataset statistics. It includes statistics and visualizations of the outputs built into the Perception package and should give a good idea of how to use datasetinsights to visualize custom annotations and metrics.\n",
    "\n",
    "## Setup dataset\n",
    "If the dataset was generated locally, point `data_root` below to the path of the dataset. The `GUID` folder suffix should be changed accordingly.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_root = \"/data/<GUID>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset metadata\n",
    "Once the dataset metadata is downloaded, it can be loaded for statistics using `datasetinsights.data.simulation`. Annotation and metric definitions are loaded into pandas dataframes using `AnnotationDefinitions` and `MetricDefinitions` respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasetinsights.datasets.unity_perception import AnnotationDefinitions, MetricDefinitions\n",
    "ann_def = AnnotationDefinitions(data_root)\n",
    "ann_def.table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric_def = MetricDefinitions(data_root)\n",
    "metric_def.table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Built-in Statistics\n",
    "The following tables and charts are supplied by `datasetinsights.data.datasets.statistics.RenderedObjectInfo` on datasets that include the \"rendered object info\" metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasetinsights.stats.statistics import RenderedObjectInfo\n",
    "import datasetinsights.datasets.unity_perception.metrics as metrics\n",
    "from datasetinsights.datasets.unity_perception.exceptions import DefinitionIDError\n",
    "from datasetinsights.stats import bar_plot, histogram_plot, rotation_plot\n",
    "\n",
    "max_samples = 10000          # maximum number of samples points used in histogram plots\n",
    "\n",
    "rendered_object_info_definition_id = \"5ba92024-b3b7-41a7-9d3f-c03a6a8ddd01\"\n",
    "roinfo = None\n",
    "try:\n",
    "    roinfo = RenderedObjectInfo(data_root=data_root, def_id=rendered_object_info_definition_id)\n",
    "except DefinitionIDError:\n",
    "    print(\"No RenderedObjectInfo in this dataset\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Descriptive Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if roinfo is not None:\n",
    "    print(roinfo.num_captures())\n",
    "    roinfo.raw_table.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Total Object Count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if roinfo is not None:\n",
    "    total_count = roinfo.total_counts()\n",
    "    display(total_count)\n",
    "    \n",
    "    display(bar_plot(\n",
    "        total_count, \n",
    "        x=\"label_id\", \n",
    "        y=\"count\", \n",
    "        x_title=\"Label Name\",\n",
    "        y_title=\"Count\",\n",
    "        title=\"Total Object Count in Dataset\",\n",
    "        hover_name=\"label_name\"\n",
    "    ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Per Capture Object Count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if roinfo is not None:\n",
    "    per_capture_count = roinfo.per_capture_counts()\n",
    "    display(per_capture_count.head(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if roinfo is not None:\n",
    "    display(histogram_plot(\n",
    "        per_capture_count, \n",
    "        x=\"count\",  \n",
    "        x_title=\"Object Counts Per Capture\",\n",
    "        y_title=\"Frequency\",\n",
    "        title=\"Distribution of Object Counts Per Capture\",\n",
    "        max_samples=max_samples\n",
    "    ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Object Visible Pixels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if roinfo is not None:\n",
    "    display(histogram_plot(\n",
    "        roinfo.raw_table, \n",
    "        x=\"visible_pixels\",  \n",
    "        x_title=\"Visible Pixels Per Object\",\n",
    "        y_title=\"Frequency\",\n",
    "        title=\"Distribution of Visible Pixels Per Object\",\n",
    "        max_samples=max_samples\n",
    "    ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annotation Visualization\n",
    "In the following sections we show how to load annotations from the Captures object and visualize them. Similar code can be used to consume annotations for model training or visualize and train on custom annotations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unity Simulation [Optional]\n",
    "If the dataset was generated on Unity Simulation, the following cells can be used to download the images, captures and annotations in the dataset. Make sure you have enough disk space to store all files. For example, a dataset with 100K captures requires roughly 300GiB storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# downloader.download(source_uri=source_uri, output=data_root, include_binary=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load captures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasetinsights.datasets.unity_perception.captures import Captures\n",
    "cap = Captures(data_root)\n",
    "cap.captures.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bounding Boxes\n",
    "In this section we render 2d bounding boxes on top of the captured images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "def cleanup(catalog):\n",
    "    catalog = remove_captures_with_missing_files(data_root, catalog)\n",
    "    catalog = remove_captures_without_bboxes(catalog)\n",
    "    return catalog\n",
    "\n",
    "def remove_captures_without_bboxes(catalog):\n",
    "    keep_mask = catalog[\"annotation.values\"].apply(len) > 0\n",
    "    return catalog[keep_mask]\n",
    "\n",
    "def remove_captures_with_missing_files(root, catalog):\n",
    "    def exists(capture_file):\n",
    "        path = Path(root) / capture_file\n",
    "        return path.exists()\n",
    "    keep_mask = catalog.filename.apply(exists)\n",
    "    return catalog[keep_mask]\n",
    "\n",
    "def capture_df(def_id):\n",
    "    captures = Captures(data_root)\n",
    "    catalog = captures.filter(bounding_box_definition_id)\n",
    "    catalog=cleanup(catalog)\n",
    "    return catalog\n",
    "\n",
    "def label_mappings_dict(def_id):\n",
    "    annotation_def = AnnotationDefinitions(data_root)\n",
    "    init_definition = annotation_def.get_definition(bounding_box_definition_id)\n",
    "    label_mappings = {\n",
    "        m[\"label_id\"]: m[\"label_name\"] for m in init_definition[\"spec\"]\n",
    "    }\n",
    "    return label_mappings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "from PIL import Image\n",
    "\n",
    "from datasetinsights.stats.visualization.plots import plot_bboxes\n",
    "from datasetinsights.datasets.synthetic import read_bounding_box_2d\n",
    "\n",
    "bounding_box_definition_id = \"c31620e3-55ff-4af6-ae86-884aa0daa9b2\"\n",
    "\n",
    "try:\n",
    "    catalog= capture_df(bounding_box_definition_id)\n",
    "    label_mappings=label_mappings_dict(bounding_box_definition_id)\n",
    "except DefinitionIDError:\n",
    "    print(\"No bounding boxes found\")\n",
    "    \n",
    "def draw_bounding_boxes(index):\n",
    "    cap = catalog.iloc[index]\n",
    "    capture_file = cap.filename\n",
    "    ann = cap[\"annotation.values\"]\n",
    "    capture = Image.open(os.path.join(data_root, capture_file))\n",
    "    image = capture.convert(\"RGB\")  # Remove alpha channel\n",
    "    bboxes = read_bounding_box_2d(ann, label_mappings)\n",
    "    return plot_bboxes(image, bboxes, label_mappings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from ipywidgets import interact\n",
    "\n",
    "# pick an index and visualize\n",
    "interact(draw_bounding_boxes, index=list(range(len(catalog))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3D Ground Truth Bounding Boxes\n",
    "In this section we render 3d ground truth bounding boxes on top of the captured images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from ipywidgets import interact\n",
    "from PIL import Image\n",
    "from datasetinsights.stats.visualization.plots import plot_bboxes3d\n",
    "from datasetinsights.datasets.synthetic import read_bounding_box_3d\n",
    "\n",
    "bounding_box_3d_defintion_id = \"0bfbe00d-00fa-4555-88d1-471b58449f5c\"\n",
    "def draw_bounding_boxes3d(index):\n",
    "    filename = os.path.join(data_root, box_captures.loc[index, \"filename\"])\n",
    "    annotations = box_captures.loc[index, \"annotation.values\"]\n",
    "    sensor = box_captures.loc[index, \"sensor\"]\n",
    "\n",
    "    if 'camera_intrinsic' in sensor:\n",
    "        projection = np.array(sensor[\"camera_intrinsic\"])\n",
    "    else:\n",
    "        projection = np.array([[1,0,0],[0,1,0],[0,0,1]])\n",
    "\n",
    "    image = Image.open(filename)\n",
    "    boxes = read_bounding_box_3d(annotations)\n",
    "    img_with_boxes = plot_bboxes3d(image, boxes, projection)\n",
    "    img_with_boxes.thumbnail([1024,1024], Image.ANTIALIAS)\n",
    "    display(img_with_boxes)\n",
    "\n",
    "try:\n",
    "    box_captures = cap.filter(def_id=bounding_box_3d_defintion_id)\n",
    "    interact(draw_bounding_boxes3d, index=(0, box_captures.shape[0]))\n",
    "except DefinitionIDError:\n",
    "    print(\"No bounding boxes found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Semantic Segmentation\n",
    "In this section we render the semantic segmentation images on top of the captured images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_with_segmentation(index, opacity):\n",
    "    filename = os.path.join(data_root, seg_captures.loc[index, \"filename\"])\n",
    "    seg_filename = os.path.join(data_root, seg_captures.loc[index, \"annotation.filename\"])\n",
    "    \n",
    "    image = Image.open(filename)\n",
    "    seg = Image.open(seg_filename)\n",
    "    img_with_seg = Image.blend(image, seg, opacity)\n",
    "    img_with_seg.thumbnail([1024,1024], Image.ANTIALIAS)\n",
    "    display(img_with_seg)\n",
    "    \n",
    "try:\n",
    "    semantic_segmentation_definition_id = \"12f94d8d-5425-4deb-9b21-5e53ad957d66\"\n",
    "    seg_captures = cap.filter(def_id=semantic_segmentation_definition_id)\n",
    "    interact(draw_with_segmentation, index=(0, seg_captures.shape[0]), opacity=(0.0, 1.0))\n",
    "except DefinitionIDError:\n",
    "    print(\"No semantic segmentation images found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instance Segmentation\n",
    "In this section we render the instance segmentation images on top of the captured images. Image IDs are mapped to an RGBA color value, below the image we include a preview of the mapping between colors and IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def instance_sorter(instance):\n",
    "    return instance[\"instance_id\"]\n",
    "\n",
    "def draw_with_instance_segmentation(index, opacity):\n",
    "    filename = os.path.join(data_root, inst_caps.loc[index, \"filename\"])\n",
    "    seg_filename = os.path.join(data_root, inst_caps.loc[index, \"annotation.filename\"])\n",
    "\n",
    "    image = Image.open(filename)\n",
    "    seg = Image.open(seg_filename)\n",
    "    img_with_seg = Image.blend(image, seg, opacity)\n",
    "    img_with_seg.thumbnail([1024,1024], Image.ANTIALIAS)\n",
    "    display(img_with_seg)\n",
    "\n",
    "    anns = inst_caps.loc[index, \"annotation.values\"].copy()\n",
    "    anns.sort(key=instance_sorter)\n",
    "\n",
    "    count = min(5, len(anns))\n",
    "    print(\"First {} ID entries:\".format(count))\n",
    "\n",
    "    for i in range(count):\n",
    "        color = anns[i].get(\"color\")\n",
    "        print (\"{} => Color({:>3}, {:>3}, {:>3})\".format(anns[i].get(\"instance_id\"), color.get(\"r\"), color.get(\"g\"), color.get(\"b\")))\n",
    "\n",
    "try:\n",
    "    inst_seg_def_id = \"1ccebeb4-5886-41ff-8fe0-f911fa8cbcdf\"\n",
    "    inst_caps = cap.filter(def_id=inst_seg_def_id)\n",
    "    interact(draw_with_instance_segmentation, index=(0, inst_caps.shape[0]), opacity=(0.0, 1.0))\n",
    "except DefinitionIDError:\n",
    "    print(\"No instance segmentation images found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keypoints\n",
    "In this section we render the keypoint labeled data for the captured frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from datasetinsights.stats.visualization.plots import plot_keypoints\n",
    "\n",
    "\n",
    "def draw_human_pose(index):\n",
    "    filename = os.path.join(data_root, keypoint_caps.loc[index, \"filename\"])\n",
    "    annotations = keypoint_caps.loc[index, \"annotation.values\"]\n",
    "    templates = ann_def.get_definition(keypoint_def_id)['spec']\n",
    "    img = Image.open(filename)\n",
    "    img_with_pose = plot_keypoints(img, annotations, templates)\n",
    "    img_with_pose.thumbnail([1024,1024], Image.ANTIALIAS)\n",
    "    display(img_with_pose)\n",
    "\n",
    "\n",
    "try:\n",
    "    keypoint_def_id = \"8b3ef246-daa7-4dd5-a0e8-a943f6e7f8c2\"\n",
    "    keypoint_caps = cap.filter(def_id=keypoint_def_id)\n",
    "    interact(draw_human_pose, index=(0, keypoint_caps.shape[0] - 1))\n",
    "except DefinitionIDError:\n",
    "    print(\"No keypoint data found\")\n"
   ]
  }
 ],
 "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}