gt_corner.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Errno 2] No such file or directory: 'RTM/src'\n",
      "/home/kaixin1/Project/RTM/src\n"
     ]
    }
   ],
   "source": [
    "cd RTM/src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import pickle\n",
    "import json\n",
    "import numpy as np\n",
    "import cv2\n",
    "DEBUG = False\n",
    "# VAL_PATH = DATA_PATH + 'training/label_val/'\n",
    "import os\n",
    "SPLITS = ['3dop', 'subcnn'] \n",
    "import _init_paths\n",
    "from utils.ddd_utils import compute_box_3d, project_to_image, alpha2rot_y\n",
    "from utils.ddd_utils import draw_box_3d, unproject_2d_to_3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = '../data/kitti/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _bbox_to_coco_bbox(bbox):\n",
    "  return [(bbox[0]), (bbox[1]),\n",
    "          (bbox[2] - bbox[0]), (bbox[3] - bbox[1])]\n",
    "\n",
    "def read_clib(calib_path):\n",
    "  f = open(calib_path, 'r')\n",
    "  for i, line in enumerate(f):\n",
    "    if i == 2:\n",
    "      calib = np.array(line[:-1].split(' ')[1:], dtype=np.float32)\n",
    "      calib = calib.reshape(3, 4)\n",
    "      return calib\n",
    "\n",
    "cats = ['Pedestrian', 'Car', 'Cyclist', 'Van', 'Truck',  'Person_sitting',\n",
    "        'Tram', 'Misc', 'DontCare']\n",
    "cat_ids = {cat: i + 1 for i, cat in enumerate(cats)}\n",
    "# cat_info = [{\"name\": \"pedestrian\", \"id\": 1}, {\"name\": \"vehicle\", \"id\": 2}]\n",
    "F = 721\n",
    "H = 384 # 375\n",
    "W = 1248 # 1242\n",
    "EXT = [45.75, -0.34, 0.005]\n",
    "CALIB = np.array([[F, 0, W / 2, EXT[0]], [0, F, H / 2, EXT[1]], \n",
    "                  [0, 0, 1, EXT[2]]], dtype=np.float32)\n",
    "\n",
    "cat_info = []\n",
    "for i, cat in enumerate(cats):\n",
    "  cat_info.append({'name': cat, 'id': i + 1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'name': 'Pedestrian', 'id': 1},\n",
       " {'name': 'Car', 'id': 2},\n",
       " {'name': 'Cyclist', 'id': 3},\n",
       " {'name': 'Van', 'id': 4},\n",
       " {'name': 'Truck', 'id': 5},\n",
       " {'name': 'Person_sitting', 'id': 6},\n",
       " {'name': 'Tram', 'id': 7},\n",
       " {'name': 'Misc', 'id': 8},\n",
       " {'name': 'DontCare', 'id': 9}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEBUG  = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "for SPLIT in SPLITS:\n",
    "  image_set_path = DATA_PATH + 'ImageSets_{}/'.format(SPLIT)\n",
    "  ann_dir = DATA_PATH + 'training/label_2/'\n",
    "  calib_dir = DATA_PATH + '{}/calib/'\n",
    "#   splits = ['train', 'val']\n",
    "  splits = ['train']\n",
    "  # splits = ['trainval', 'test']\n",
    "  calib_type = {'train': 'training', 'val': 'training', 'trainval': 'training',\n",
    "                'test': 'testing'}\n",
    "#   print(\"image_set_path\",image_set_path)\n",
    "#   print(\"ann_dir\",ann_dir)\n",
    "#   print(\"calib_dir\",calib_dir)\n",
    "  for split in splits:\n",
    "    ret = {'images': [], 'annotations': [], \"categories\": cat_info}\n",
    "    image_set = open(image_set_path + '{}.txt'.format(split), 'r')\n",
    "    image_to_id = {}\n",
    "    for line in image_set:\n",
    "      if line[-1] == '\\n':\n",
    "        line = line[:-1]\n",
    "      image_id = int(line)\n",
    "      calib_path = calib_dir.format(calib_type[split]) + '{}.txt'.format(line)\n",
    "      calib = read_clib(calib_path)\n",
    "      image_info = {'file_name': '{}.png'.format(line),\n",
    "                    'id': int(image_id),\n",
    "                    'calib': calib.tolist()}\n",
    "      ret['images'].append(image_info)\n",
    "      if split == 'test':\n",
    "        continue\n",
    "      ann_path = ann_dir + '{}.txt'.format(line)\n",
    "      # if split == 'val':\n",
    "      #   os.system('cp {} {}/'.format(ann_path, VAL_PATH))\n",
    "      anns = open(ann_path, 'r')\n",
    "      if DEBUG:\n",
    "        image = cv2.imread(\n",
    "          DATA_PATH + 'images/trainval/' + image_info['file_name'])\n",
    "\n",
    "      for ann_ind, txt in enumerate(anns):\n",
    "        tmp = txt[:-1].split(' ')\n",
    "        cat_id = cat_ids[tmp[0]]\n",
    "        truncated = int(float(tmp[1]))\n",
    "        occluded = int(tmp[2])\n",
    "        alpha = float(tmp[3])\n",
    "        bbox = [float(tmp[4]), float(tmp[5]), float(tmp[6]), float(tmp[7])]\n",
    "        dim = [float(tmp[8]), float(tmp[9]), float(tmp[10])]\n",
    "        location = [float(tmp[11]), float(tmp[12]), float(tmp[13])]\n",
    "        rotation_y = float(tmp[14])\n",
    "\n",
    "        ann = {'image_id': image_id,\n",
    "               'id': int(len(ret['annotations']) + 1),\n",
    "               'category_id': cat_id,\n",
    "               'dim': dim,\n",
    "               'bbox': _bbox_to_coco_bbox(bbox),\n",
    "               'depth': location[2],\n",
    "               'alpha': alpha,\n",
    "               'truncated': truncated,\n",
    "               'occluded': occluded,\n",
    "               'location': location,\n",
    "               'rotation_y': rotation_y}\n",
    "        ret['annotations'].append(ann)\n",
    "        if DEBUG and tmp[0] != 'DontCare':\n",
    "          box_3d = compute_box_3d(dim, location, rotation_y)\n",
    "          box_2d = project_to_image(box_3d, calib)\n",
    "          # print('box_2d', box_2d)\n",
    "#           image = draw_box_3d(image, box_2d)\n",
    "          ann['vertex'] = box_2d.tolist()\n",
    "          x = (bbox[0] + bbox[2]) / 2\n",
    "          '''\n",
    "          print('rot_y, alpha2rot_y, dlt', tmp[0], \n",
    "                rotation_y, alpha2rot_y(alpha, x, calib[0, 2], calib[0, 0]),\n",
    "                np.cos(\n",
    "                  rotation_y - alpha2rot_y(alpha, x, calib[0, 2], calib[0, 0])))\n",
    "          '''\n",
    "          depth = np.array([location[2]], dtype=np.float32)\n",
    "          pt_2d = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2],\n",
    "                            dtype=np.float32)\n",
    "          pt_3d = unproject_2d_to_3d(pt_2d, depth, calib)\n",
    "          pt_3d[1] += dim[0] / 2\n",
    "    out_path = '{}/annotations/kitti_{}_{}.json'.format(DATA_PATH, SPLIT, split)\n",
    "    json.dump(ret, open(out_path, 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "vertexes = box_2d.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "ct = [395,196]\n",
    "ct = np.array(ct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "vc = np.zeros(18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.array(vertexes) -ct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "vc = r.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  8.55215454,  14.06947327, -20.89697266,  14.10087585,\n",
       "        -4.94778442,  11.40611267,  22.35986328,  11.37910461,\n",
       "         8.55215454, -10.31129456, -20.89697266, -10.30046082,\n",
       "        -4.94778442, -11.22981262,  22.35986328, -11.23912048,\n",
       "         1.55163574,   0.95039368])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "def compute_box_3d(dim, location, rotation_y):\n",
    "  # dim: 3\n",
    "  # location: 3\n",
    "  # rotation_y: 1\n",
    "  # return: 8 x 3\n",
    "  c, s = np.cos(rotation_y), np.sin(rotation_y)\n",
    "  R = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=np.float32)\n",
    "  l, w, h = dim[2], dim[1], dim[0]\n",
    "  x_corners = [l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2,0]\n",
    "  y_corners = [0,0,0,0,-h,-h,-h,-h,-h/2]\n",
    "  z_corners = [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2,0]\n",
    "\n",
    "  corners = np.array([x_corners, y_corners, z_corners], dtype=np.float32)\n",
    "  corners_3d = np.dot(R, corners) \n",
    "  corners_3d = corners_3d + np.array(location, dtype=np.float32).reshape(3, 1)\n",
    "  return corners_3d.transpose(1, 0)\n",
    "\n",
    "def project_to_image(pts_3d, P):\n",
    "  # pts_3d: n x 3\n",
    "  # P: 3 x 4\n",
    "  # return: n x 2\n",
    "  pts_3d_homo = np.concatenate(\n",
    "    [pts_3d, np.ones((pts_3d.shape[0], 1), dtype=np.float32)], axis=1)\n",
    "  pts_2d = np.dot(P, pts_3d_homo.transpose(1, 0)).transpose(1, 0)\n",
    "  pts_2d = pts_2d[:, :2] / pts_2d[:, 2:]\n",
    "  # import pdb; pdb.set_trace()\n",
    "  return pts_2d\n",
    "\n",
    "def compute_orientation_3d(dim, location, rotation_y):\n",
    "  # dim: 3\n",
    "  # location: 3\n",
    "  # rotation_y: 1\n",
    "  # return: 2 x 3\n",
    "  c, s = np.cos(rotation_y), np.sin(rotation_y)\n",
    "  R = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=np.float32)\n",
    "  orientation_3d = np.array([[0, dim[2]], [0, 0], [0, 0]], dtype=np.float32)\n",
    "  orientation_3d = np.dot(R, orientation_3d)\n",
    "  orientation_3d = orientation_3d + \\\n",
    "                   np.array(location, dtype=np.float32).reshape(3, 1)\n",
    "  return orientation_3d.transpose(1, 0)\n",
    "\n",
    "def draw_box_3d(image, corners, c=(0, 0, 255)):\n",
    "  face_idx = [[0,1,5,4],\n",
    "              [1,2,6, 5],\n",
    "              [2,3,7,6],\n",
    "              [3,0,4,7]]\n",
    "  for ind_f in range(3, -1, -1):\n",
    "    f = face_idx[ind_f]\n",
    "    for j in range(4):\n",
    "      cv2.line(image, (corners[f[j], 0], corners[f[j], 1]),\n",
    "               (corners[f[(j+1)%4], 0], corners[f[(j+1)%4], 1]), c, 2, lineType=cv2.LINE_AA)\n",
    "    if ind_f == 0:\n",
    "      cv2.line(image, (corners[f[0], 0], corners[f[0], 1]),\n",
    "               (corners[f[2], 0], corners[f[2], 1]), c, 1, lineType=cv2.LINE_AA)\n",
    "      cv2.line(image, (corners[f[1], 0], corners[f[1], 1]),\n",
    "               (corners[f[3], 0], corners[f[3], 1]), c, 1, lineType=cv2.LINE_AA)\n",
    "  return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}