KITTI_get_item.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "KITTI_get_item",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QgAhLIuJEruY"
      },
      "source": [
        "## Load data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fx5JPhE7sL4h",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "591569c5-95a9-44db-e9af-3c988a14b7cf"
      },
      "source": [
        "#camera calibration\n",
        "!wget https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_calib.zip\n",
        "!unzip data_object_calib.zip -d ./data_object_calib | grep \"I love DL\"\n",
        "!rm data_object_calib.zip"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-12-06 12:27:15--  https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_calib.zip\n",
            "Resolving s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)... 52.219.47.207\n",
            "Connecting to s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)|52.219.47.207|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 26854811 (26M) [application/zip]\n",
            "Saving to: ‘data_object_calib.zip’\n",
            "\n",
            "data_object_calib.z 100%[===================>]  25.61M  81.7MB/s    in 0.3s    \n",
            "\n",
            "2020-12-06 12:27:15 (81.7 MB/s) - ‘data_object_calib.zip’ saved [26854811/26854811]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aHOa4bZnsqG4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dd8abc47-111a-4164-cf9b-d97369e8691a"
      },
      "source": [
        "!wget  https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_label_2.zip\n",
        "!unzip data_object_label_2.zip -d ./data_object_label_2 | grep \"I love DL\"\n",
        "!rm data_object_label_2.zip "
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-12-06 12:27:17--  https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_label_2.zip\n",
            "Resolving s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)... 52.219.47.207\n",
            "Connecting to s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)|52.219.47.207|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 5601213 (5.3M) [application/zip]\n",
            "Saving to: ‘data_object_label_2.zip’\n",
            "\n",
            "\rdata_object_label_2   0%[                    ]       0  --.-KB/s               \rdata_object_label_2 100%[===================>]   5.34M  --.-KB/s    in 0.08s   \n",
            "\n",
            "2020-12-06 12:27:17 (70.5 MB/s) - ‘data_object_label_2.zip’ saved [5601213/5601213]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XkK5U4kTtSJc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "3c5f43e7-f442-4da5-f980-8eee87d987d2"
      },
      "source": [
        "!wget  https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_image_2.zip\n",
        "!unzip data_object_image_2.zip -d ./data_object_image_2 | grep \"I love DL\"\n",
        "!rm data_object_image_2.zip"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-12-06 12:27:20--  https://s3.eu-central-1.amazonaws.com/avg-kitti/data_object_image_2.zip\n",
            "Resolving s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)... 52.219.74.199\n",
            "Connecting to s3.eu-central-1.amazonaws.com (s3.eu-central-1.amazonaws.com)|52.219.74.199|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 12569945557 (12G) [application/zip]\n",
            "Saving to: ‘data_object_image_2.zip’\n",
            "\n",
            "data_object_image_2 100%[===================>]  11.71G  43.1MB/s    in 4m 45s  \n",
            "\n",
            "2020-12-06 12:32:05 (42.0 MB/s) - ‘data_object_image_2.zip’ saved [12569945557/12569945557]\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iHpta6rkxZhY"
      },
      "source": [
        "import os\n",
        "import cv2\n",
        "import math\n",
        "import random\n",
        "import torch.utils.data as data\n",
        "import cv2\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.patches as patches\n",
        "\n",
        "import os\n",
        "\n",
        "from os import listdir\n",
        "from os.path import isfile, join\n",
        "import pandas as pd\n",
        "from sklearn.model_selection import train_test_split\n",
        "import gc; gc.enable() # memory is tight\n",
        "import torch\n",
        "\n",
        "dtype = \"float32\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-1HR953GBadl"
      },
      "source": [
        "## Create file path list"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zjgKIDdtxroE"
      },
      "source": [
        "# img_dir = 'kitti dataset/'\n",
        "# ../input/kitti_single/training/label_2/\n",
        "train_label_dir = './data_object_label_2/training/label_2/'\n",
        "# label_test_image_dir = os.path.join(os.getcwd(), 'streets\\\\test\\\\labels\\\\')\n",
        "\n",
        "train_image_dir = './data_object_image_2/training/image_2/'\n",
        "# test_image_dir = os.path.join(os.getcwd(), 'streets\\\\test\\\\images\\\\')\n",
        "\n",
        "train_calib_dir = './data_object_calib/training/calib/'\n",
        "\n",
        "images =  [(train_image_dir+f) for f in listdir(train_image_dir) if isfile(join(train_image_dir, f))]\n",
        "labels = [(train_label_dir+f) for f in listdir(train_label_dir) if isfile(join(train_label_dir, f))]\n",
        "calibs = [(train_calib_dir+f) for f in listdir(train_calib_dir) if isfile(join(train_calib_dir, f))]\n",
        "\n",
        "df = pd.DataFrame(np.column_stack([images, labels, calibs]), columns=['images', 'labels', 'calibs'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mC-FPVpLzpgF",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 426
        },
        "outputId": "5c21e267-0f24-4a7a-afe6-cb3da66edacb"
      },
      "source": [
        "df"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "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>images</th>\n",
              "      <th>labels</th>\n",
              "      <th>calibs</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>./data_object_image_2/training/image_2/005644.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/002003.txt</td>\n",
              "      <td>./data_object_calib/training/calib/002003.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>./data_object_image_2/training/image_2/004044.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/004047.txt</td>\n",
              "      <td>./data_object_calib/training/calib/004047.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>./data_object_image_2/training/image_2/003676.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/000614.txt</td>\n",
              "      <td>./data_object_calib/training/calib/000614.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>./data_object_image_2/training/image_2/004612.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/000513.txt</td>\n",
              "      <td>./data_object_calib/training/calib/000513.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>./data_object_image_2/training/image_2/004063.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/006156.txt</td>\n",
              "      <td>./data_object_calib/training/calib/006156.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>...</th>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "      <td>...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7476</th>\n",
              "      <td>./data_object_image_2/training/image_2/002843.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/003855.txt</td>\n",
              "      <td>./data_object_calib/training/calib/003855.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7477</th>\n",
              "      <td>./data_object_image_2/training/image_2/000345.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/005322.txt</td>\n",
              "      <td>./data_object_calib/training/calib/005322.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7478</th>\n",
              "      <td>./data_object_image_2/training/image_2/001058.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/003709.txt</td>\n",
              "      <td>./data_object_calib/training/calib/003709.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7479</th>\n",
              "      <td>./data_object_image_2/training/image_2/002608.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/003799.txt</td>\n",
              "      <td>./data_object_calib/training/calib/003799.txt</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>7480</th>\n",
              "      <td>./data_object_image_2/training/image_2/002601.png</td>\n",
              "      <td>./data_object_label_2/training/label_2/005114.txt</td>\n",
              "      <td>./data_object_calib/training/calib/005114.txt</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "<p>7481 rows × 3 columns</p>\n",
              "</div>"
            ],
            "text/plain": [
              "                                                 images  ...                                         calibs\n",
              "0     ./data_object_image_2/training/image_2/005644.png  ...  ./data_object_calib/training/calib/002003.txt\n",
              "1     ./data_object_image_2/training/image_2/004044.png  ...  ./data_object_calib/training/calib/004047.txt\n",
              "2     ./data_object_image_2/training/image_2/003676.png  ...  ./data_object_calib/training/calib/000614.txt\n",
              "3     ./data_object_image_2/training/image_2/004612.png  ...  ./data_object_calib/training/calib/000513.txt\n",
              "4     ./data_object_image_2/training/image_2/004063.png  ...  ./data_object_calib/training/calib/006156.txt\n",
              "...                                                 ...  ...                                            ...\n",
              "7476  ./data_object_image_2/training/image_2/002843.png  ...  ./data_object_calib/training/calib/003855.txt\n",
              "7477  ./data_object_image_2/training/image_2/000345.png  ...  ./data_object_calib/training/calib/005322.txt\n",
              "7478  ./data_object_image_2/training/image_2/001058.png  ...  ./data_object_calib/training/calib/003709.txt\n",
              "7479  ./data_object_image_2/training/image_2/002608.png  ...  ./data_object_calib/training/calib/003799.txt\n",
              "7480  ./data_object_image_2/training/image_2/002601.png  ...  ./data_object_calib/training/calib/005114.txt\n",
              "\n",
              "[7481 rows x 3 columns]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gRqZc8_wzrDq"
      },
      "source": [
        "df1 = df.sort_values(by='images')['images'].reset_index()\n",
        "# df1 = df.sort_values(by='a')['a']\n",
        "df2 = df.sort_values(by='labels')['labels'].reset_index()\n",
        "# df2 = df.sort_values(by='b')['b']\n",
        "df3 = df.sort_values(by='calibs')['calibs'].reset_index()\n",
        "\n",
        "df['images'] = df1['images']\n",
        "df['labels'] = df2['labels']\n",
        "df['calibs'] = df3['calibs']\n",
        "del df1, df2, df3\n",
        "df\n",
        "df_in_list = (df).values.tolist()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WT7jJaBgBj2u"
      },
      "source": [
        "## HeatMap Genrating Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pmNvJdksjX56"
      },
      "source": [
        "def gaussian_radius(det_size, min_overlap=0.7):\n",
        "    height, width = det_size\n",
        "    a1  = 1\n",
        "    b1  = (height + width)\n",
        "    c1  = width * height * (1 - min_overlap) / (1 + min_overlap)\n",
        "    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)\n",
        "    r1  = (b1 + sq1) / 2\n",
        "    a2  = 4\n",
        "    b2  = 2 * (height + width)\n",
        "    c2  = (1 - min_overlap) * width * height\n",
        "    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)\n",
        "    r2  = (b2 + sq2) / 2\n",
        "    a3  = 4 * min_overlap\n",
        "    b3  = -2 * min_overlap * (height + width)\n",
        "    c3  = (min_overlap - 1) * width * height\n",
        "    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)\n",
        "    r3  = (b3 + sq3) / 2\n",
        "    return min(r1, r2, r3)\n",
        "\n",
        "def gaussian2D(shape, sigma=1):\n",
        "    m, n = [(ss - 1.) / 2. for ss in shape]\n",
        "    y, x = np.ogrid[-m:m+1,-n:n+1]\n",
        "    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))\n",
        "    h[h < np.finfo(h.dtype).eps * h.max()] = 0\n",
        "    return h\n",
        "\n",
        "def draw_umich_gaussian(heatmap, center, radius, k=1):\n",
        "    diameter = 2 * radius + 1\n",
        "    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)\n",
        "  \n",
        "    x, y = int(center[0]), int(center[1])\n",
        "\n",
        "    height, width = heatmap.shape[0:2]\n",
        "    left, right = min(x, radius), min(width - x, radius + 1) \n",
        "    top, bottom = min(y, radius), min(height - y, radius + 1)\n",
        "\n",
        "    masked_heatmap  = heatmap[y - top:y + bottom, x - left:x + right] \n",
        "    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right] \n",
        "    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug\n",
        "        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)\n",
        "    return heatmap"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hGR1CKlhBpZo"
      },
      "source": [
        "## Get Item"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ogRU5IgGj5kj"
      },
      "source": [
        "def getitem(index):\n",
        "    image_path, label_path, cali_path = df_in_list[index]\n",
        "    #get the image (375, 1242, 3)\n",
        "    img = cv2.imread(image_path) \n",
        "    default_resolution = [375, 1242]\n",
        "    #get the labels\n",
        "    with open(label_path) as f:\n",
        "        content = f.readlines()\n",
        "    content = [x.split() for x in content]\n",
        "    # print(content)\n",
        "\n",
        "    #transform to 512 * 512\n",
        "    height, width = img.shape[0], img.shape[1]  \n",
        "    input_h, input_w = 512, 512 \n",
        "    inp = cv2.resize(img,(input_w, input_h))\n",
        "    scale_h, scale_w = input_h/height, input_w/width\n",
        "    # print(scale_h, scale_w)\n",
        "\n",
        "    \n",
        "    inp = (inp.astype(np.float32) / 255.)  \n",
        "\n",
        "    inp = inp.transpose(2, 0, 1) \n",
        "\n",
        "    \n",
        "    max_objs = 128\n",
        "    down_ratio = 4 \n",
        "    output_h = input_h // down_ratio\n",
        "    output_w = input_w // down_ratio\n",
        "    num_classes = 1\n",
        "    draw_gaussian = draw_umich_gaussian\n",
        "    hm = np.zeros((num_classes, output_h, output_w), dtype=np.float32)   \n",
        "    reg_mask = np.zeros((max_objs), dtype=np.uint8) \n",
        "    wh = np.zeros((max_objs, 2), dtype=np.float32)\n",
        "    reg = np.zeros((max_objs, 2), dtype=np.float32) \n",
        "    ind = np.zeros((max_objs), dtype=np.int64) \n",
        "    \n",
        "    count = 0\n",
        "    for c in content:\n",
        "      if (c[0] == \"Car\"):\n",
        "        bbox = np.array(c[4:8], dtype = \"float32\")\n",
        "        bbox[1::2] *= scale_h\n",
        "        bbox[0::2] *= scale_w\n",
        "        bbox = bbox/down_ratio\n",
        "        h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]\n",
        "        if h > 0 and w > 0:\n",
        "          radius = gaussian_radius((math.ceil(h), math.ceil(w)))  \n",
        "          radius = max(0, int(radius))\n",
        "          ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32) \n",
        "          ct_int = ct.astype(np.int32) \n",
        "          draw_gaussian(hm[0], ct_int, radius)\n",
        "          wh[count] = 1. * w, 1. * h\n",
        "          ind[count] = ct_int[1] * output_w + ct_int[0]  \n",
        "          reg[count] = ct - ct_int\n",
        "          reg_mask[count] = 1\n",
        "          count = count + 1\n",
        "\n",
        "    res = {'image': img, 'input': inp, 'hm': hm, 'reg_mask': reg_mask, 'ind': ind, 'wh': wh, 'reg':reg}\n",
        "    return res\n",
        "\n",
        "def reverse_res_to_bbox(res):\n",
        "    down_ratio = res['input'].shape[1] / res['hm'].shape[1]\n",
        "    output_h, output_w =  res['hm'].shape[0],  res['hm'].shape[1]\n",
        "    num_objs = np.sum(res['reg_mask'])\n",
        "    bbox = np.zeros((num_objs, 4), dtype = 'float32' )\n",
        "    ct = res['reg'][:num_objs]\n",
        "    ct[:, 0] += res['ind'][:num_objs] % output_w\n",
        "    ct[:,1] += res['ind'][:num_objs] // output_w\n",
        "    h, w = res['wh'][:num_objs,1],  res['wh'][:num_objs,0]\n",
        "    bbox[:,0] = (ct[:,0] * 2 - w) /2\n",
        "    bbox[:,2] = (ct[:,0] * 2 + w) /2\n",
        "    bbox[:,1] = (ct[:,1] * 2 - h) /2\n",
        "    bbox[:,3] = (ct[:,1] * 2 + h) /2\n",
        "    bbox *= down_ratio\n",
        "    return bbox\n",
        "\n",
        "\n",
        "    \n",
        "      \n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qGNUeZNQkHvD"
      },
      "source": [
        "im_idx = 500\n",
        "res = getitem(im_idx)\n",
        "\n",
        "img = res['image']\n",
        "inp = res['input'].transpose(1,2,0)\n",
        "hm = res['hm']\n",
        "print(inp.dtype)\n",
        "\n",
        "plt.title(\"Original Image\")\n",
        "plt.imshow(img)\n",
        "plt.show()\n",
        "\n",
        "\n",
        "plt.title(\"Ground Truth Heat Map\")\n",
        "im = plt.imshow(hm[0])\n",
        "plt.colorbar(im)\n",
        "plt.show()\n",
        "\n",
        "\n",
        "bbox = reverse_res_to_bbox(res)\n",
        "for b in bbox:\n",
        "  cv2.rectangle(inp, (b[0], b[1]), (b[2], b[3]), (0,0,1), 2) \n",
        "\n",
        "plt.title(\"Calculated bounding box positions\")\n",
        "plt.imshow(inp)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}