struct2depth_imported_files.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "struct2depth imported files",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/xiajenny/struct2depth2unity/blob/master/struct2depth_imported_files.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "33m0rFwPsWSC",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "318ed13c-3890-4ee5-85e3-36e562d3312e"
      },
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive', force_remount = True)\n",
        "\n",
        "import os\n",
        "from absl import app\n",
        "from absl import flags\n",
        "from absl import logging\n",
        "import numpy as np\n",
        "import cv2\n",
        "import os, glob"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/gdrive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "OytDFpeDsvfp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 170
        },
        "outputId": "5c67a36b-848c-4a00-d298-1963723da22f"
      },
      "cell_type": "code",
      "source": [
        "# alignment.py, gen_data_kitti.py\n",
        "\n",
        "def compute_overlap(mask1, mask2):\n",
        "    # Use IoU here.\n",
        "    return np.sum(mask1 & mask2)/np.sum(mask1 | mask2)\n",
        "\n",
        "def align(seg_img1, seg_img2, seg_img3, threshold_same=0.3):\n",
        "    res_img1 = np.zeros_like(seg_img1)\n",
        "    res_img2 = np.zeros_like(seg_img2)\n",
        "    res_img3 = np.zeros_like(seg_img3)\n",
        "    remaining_objects2 = list(np.unique(seg_img2.flatten()))\n",
        "    remaining_objects3 = list(np.unique(seg_img3.flatten()))\n",
        "    for seg_id in np.unique(seg_img1):\n",
        "        # See if we can find correspondences to seg_id in seg_img2.\n",
        "        max_overlap2 = float('-inf')\n",
        "        max_segid2 = -1\n",
        "        for seg_id2 in remaining_objects2:\n",
        "            overlap = compute_overlap(seg_img1==seg_id, seg_img2==seg_id2)\n",
        "            if overlap>max_overlap2:\n",
        "                max_overlap2 = overlap\n",
        "                max_segid2 = seg_id2\n",
        "        if max_overlap2 > threshold_same:\n",
        "            max_overlap3 = float('-inf')\n",
        "            max_segid3 = -1\n",
        "            for seg_id3 in remaining_objects3:\n",
        "                overlap = compute_overlap(seg_img2==max_segid2, seg_img3==seg_id3)\n",
        "                if overlap>max_overlap3:\n",
        "                    max_overlap3 = overlap\n",
        "                    max_segid3 = seg_id3\n",
        "            if max_overlap3 > threshold_same:\n",
        "                res_img1[seg_img1==seg_id] = seg_id\n",
        "                res_img2[seg_img2==max_segid2] = seg_id\n",
        "                res_img3[seg_img3==max_segid3] = seg_id\n",
        "                remaining_objects2.remove(max_segid2)\n",
        "                remaining_objects3.remove(max_segid3)\n",
        "    return res_img1, res_img2, res_img3\n",
        "  \n",
        "SEQ_LENGTH = 3\n",
        "WIDTH = 416\n",
        "HEIGHT = 128\n",
        "STEPSIZE = 1\n",
        "INPUT_DIR = '/content/gdrive/My Drive/struct2depth2unity/raw_data'\n",
        "OUTPUT_DIR = '/content/gdrive/My Drive/struct2depth2unity/kitti_processed'\n",
        "\n",
        "\n",
        "def get_line(file, start):\n",
        "    file = open(file, 'r')\n",
        "    lines = file.readlines()\n",
        "    lines = [line.rstrip() for line in lines]\n",
        "    ret = None\n",
        "    for line in lines:\n",
        "        nline = line.split(': ')\n",
        "        if nline[0]==start:\n",
        "            ret = nline[1].split(' ')\n",
        "            ret = np.array([float(r) for r in ret], dtype=float)\n",
        "            ret = ret.reshape((3,4))[0:3, 0:3]\n",
        "            break\n",
        "    file.close()\n",
        "    return ret\n",
        "\n",
        "\n",
        "def crop(img, segimg, fx, fy, cx, cy):\n",
        "    # Perform center cropping, preserving 50% vertically.\n",
        "    middle_perc = 0.50\n",
        "    left = 1-middle_perc\n",
        "    half = left/2\n",
        "    a = img[int(img.shape[0]*(half)):int(img.shape[0]*(1-half)), :]\n",
        "    aseg = segimg[int(segimg.shape[0]*(half)):int(segimg.shape[0]*(1-half)), :]\n",
        "    cy /= (1/middle_perc)\n",
        "\n",
        "    # Resize to match target height while preserving aspect ratio.\n",
        "    wdt = int((128*a.shape[1]/a.shape[0]))\n",
        "    x_scaling = float(wdt)/a.shape[1]\n",
        "    y_scaling = 128.0/a.shape[0]\n",
        "    b = cv2.resize(a, (wdt, 128))\n",
        "    bseg = cv2.resize(aseg, (wdt, 128))\n",
        "\n",
        "    # Adjust intrinsics.\n",
        "    fx*=x_scaling\n",
        "    fy*=y_scaling\n",
        "    cx*=x_scaling\n",
        "    cy*=y_scaling\n",
        "\n",
        "    # Perform center cropping horizontally.\n",
        "    remain = b.shape[1] - 416\n",
        "    cx /= (b.shape[1]/416)\n",
        "    c = b[:, int(remain/2):b.shape[1]-int(remain/2)]\n",
        "    cseg = bseg[:, int(remain/2):b.shape[1]-int(remain/2)]\n",
        "\n",
        "    return c, cseg, fx, fy, cx, cy\n",
        "\n",
        "\n",
        "def run_all():\n",
        "  ct = 0\n",
        "if not OUTPUT_DIR.endswith('/'):\n",
        "    OUTPUT_DIR = OUTPUT_DIR + '/'\n",
        "\n",
        "for d in glob.glob(INPUT_DIR + '/*/'):\n",
        "    date = d.split('/')[-2]\n",
        "    print(date)\n",
        "    file_calibration = d + 'calib_cam_to_cam.txt'\n",
        "    calib_raw = [get_line(file_calibration, 'P_rect_02'), get_line(file_calibration, 'P_rect_03')]\n",
        "    print(calib_raw[0], calib_raw[1])\n",
        "    \n",
        "    for d2 in glob.glob(d + '*/'):\n",
        "        seqname = d2.split('/')[-2]\n",
        "        print('Processing sequence', seqname)\n",
        "        for subfolder in ['image_02/data', 'image_03/data']:\n",
        "            ct = 1\n",
        "            seqname = d2.split('/')[-2] + subfolder.replace('image', '').replace('/data', '')\n",
        "            if not os.path.exists(OUTPUT_DIR + seqname):\n",
        "                os.mkdir(OUTPUT_DIR + seqname)\n",
        "\n",
        "            calib_camera = calib_raw[0] if subfolder=='image_02/data' else calib_raw[1]\n",
        "            folder = d2 + subfolder\n",
        "            print(folder)\n",
        "            files = glob.glob(folder + '/*.png')\n",
        "            files = [file for file in files if not 'disp' in file and not 'flip' in file and not 'seg' in file]\n",
        "            files = sorted(files)\n",
        "            for i in range(SEQ_LENGTH, len(files)+1, STEPSIZE):\n",
        "                imgnum = str(ct).zfill(10)\n",
        "                if os.path.exists(OUTPUT_DIR + seqname + '/' + imgnum + '.png'):\n",
        "                    ct+=1\n",
        "                    continue\n",
        "                big_img = np.zeros(shape=(HEIGHT, WIDTH*SEQ_LENGTH, 3))\n",
        "                wct = 0\n",
        "\n",
        "                for j in range(i-SEQ_LENGTH, i):  # Collect frames for this sample.\n",
        "                    img = cv2.imread(files[j])\n",
        "                    ORIGINAL_HEIGHT, ORIGINAL_WIDTH, _ = img.shape\n",
        "\n",
        "                    zoom_x = WIDTH/ORIGINAL_WIDTH\n",
        "                    zoom_y = HEIGHT/ORIGINAL_HEIGHT\n",
        "\n",
        "                    # Adjust intrinsics.\n",
        "                    calib_current = calib_camera.copy()\n",
        "                    calib_current[0, 0] *= zoom_x\n",
        "                    calib_current[0, 2] *= zoom_x\n",
        "                    calib_current[1, 1] *= zoom_y\n",
        "                    calib_current[1, 2] *= zoom_y\n",
        "\n",
        "                    calib_representation = ','.join([str(c) for c in calib_current.flatten()])\n",
        "                    #print(calib_representation)\n",
        "                    #print(calib_current)\n",
        "                    calib_representation.encode(encoding='UTF-8')\n",
        "                    img = cv2.resize(img, (WIDTH, HEIGHT))\n",
        "                    big_img[:,wct*WIDTH:(wct+1)*WIDTH] = img\n",
        "                    wct+=1\n",
        "                cv2.imwrite(OUTPUT_DIR + seqname + '/' + imgnum + '.png', big_img)\n",
        "                f = open(OUTPUT_DIR + seqname + '/' + imgnum + '_cam.txt', 'w')\n",
        "                f.write(calib_representation)\n",
        "                f.close()\n",
        "                ct+=1\n",
        "\n",
        "run_all()\n",
        "\n"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "KITTI_RAW\n",
            "[[721.5377   0.     609.5593]\n",
            " [  0.     721.5377 172.854 ]\n",
            " [  0.       0.       1.    ]] [[721.5377   0.     609.5593]\n",
            " [  0.     721.5377 172.854 ]\n",
            " [  0.       0.       1.    ]]\n",
            "Processing sequence 2011_09_26_drive_0001_sync\n",
            "/content/gdrive/My Drive/struct2depth2unity/raw_data/KITTI_RAW/2011_09_26_drive_0001_sync/image_02/data\n",
            "/content/gdrive/My Drive/struct2depth2unity/raw_data/KITTI_RAW/2011_09_26_drive_0001_sync/image_03/data\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}