{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e0dd1dae-47ab-4832-a6d6-17ca23a4b4f4",
   "metadata": {},
   "source": [
    "# Loading DAVIS recordings with events and frames\n",
    "Let's load a sample that contains a tuple of events, inertial measurement unit (IMU) recordings and images. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d72047c6-e505-490c-8069-645a6eb7339f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tonic\n",
    "\n",
    "dataset = tonic.datasets.DAVISDATA(save_to=\"data\", recording=\"shapes_6dof\")\n",
    "\n",
    "data, targets = dataset[0]\n",
    "events, imu, images = data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d1c73a-43c1-4a7f-bd92-4eb07b5c7c79",
   "metadata": {},
   "source": [
    "The timestamps for events are from 0 to some 3.4 seconds. We also have timestamps for images, which are regularly sampled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "822187be-4415-4ce5-9b35-5a35a6734668",
   "metadata": {},
   "outputs": [],
   "source": [
    "events[\"t\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fae77b07-2bee-4179-bf54-db1de4079c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "images[\"ts\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc5be6d-6414-4e47-8f66-c4ac757c8b6b",
   "metadata": {},
   "source": [
    "Let's bin our events into roughly the same time bins. The sampling frequency for images in microseconds can be calculated easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4276644-9848-435d-8de8-a139c57fb17f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "mean_diff = np.diff(list(zip(images[\"ts\"], images[\"ts\"][1:]))).mean()\n",
    "print(f\"Average difference in image timestamps in microseconds: {mean_diff}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a2a0089-a666-4d87-9a2b-58985920e6f8",
   "metadata": {},
   "source": [
    "Say we wanted to apply the same transform to the event frames and grey-level images at the same time. Since the DAVIS dataset has a tuple of data, we have to make use of a helper function to apply our transforms specifically, as we can not apply any single transform to the data tuple at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32ab581e-209e-478d-a075-f43275299d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "sensor_size = tonic.datasets.DAVISDATA.sensor_size\n",
    "frame_transform = tonic.transforms.ToFrame(\n",
    "    sensor_size=sensor_size, time_window=mean_diff\n",
    ")\n",
    "\n",
    "image_center_crop = torchvision.transforms.Compose(\n",
    "    [torch.tensor, torchvision.transforms.CenterCrop((100, 100))]\n",
    ")\n",
    "\n",
    "\n",
    "def data_transform(data):\n",
    "    # first we have to unpack our data\n",
    "    events, imu, images = data\n",
    "    # we bin events to event frames\n",
    "    frames = frame_transform(events)\n",
    "    # then we can apply frame transforms to both event frames and images at the same time\n",
    "    frames_cropped = image_center_crop(frames)\n",
    "    images_cropped = image_center_crop(images[\"frames\"])\n",
    "    return frames_cropped, imu, images_cropped"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efff5fff-a943-4b8e-b149-3cec88449ded",
   "metadata": {},
   "source": [
    "Now we can load the same sample file again, this time with our custom transform function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635d338f-a2cf-4e22-be96-c303e9d65789",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = tonic.datasets.DAVISDATA(\n",
    "    save_to=\"./data\", recording=\"slider_depth\", transform=data_transform\n",
    ")\n",
    "\n",
    "data, targets = dataset[0]\n",
    "frames_cropped, imu, images_cropped = data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af1524f8-d9e9-43aa-998a-ce83ac0d3139",
   "metadata": {},
   "source": [
    "All what's left is to plot binned event frame and image next to each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55111905-d20f-4342-9ab7-2911990008cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "event_frame = frames_cropped[10]\n",
    "ax1.imshow(event_frame[0] - event_frame[1])\n",
    "ax1.set_title(\"event frame\")\n",
    "ax2.imshow(images_cropped[10], cmap=mpl.cm.gray)\n",
    "ax2.set_title(\"grey level image\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b83ccc6-1c3b-4739-8776-1b0fe755ae67",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}