""" Helper methods for loading and parsing KITTI data.

Author: <NAME>
Date: September 2017
"""
from __future__ import print_function

import tensorflow as tf
import numpy as np
import cv2
import os


class Object3d(object):
    """ 3d object label """
    def __init__(self, label_file_line=None):
        # each object represents a line in the %06d.txt
        # that is each object3D represent a instance object in the image
        # Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 
        # 1.89 0.48 1.20 1.84 1.47 8.41 0.01
        data = label_file_line.split(' ')
        data[1:] = [float(x) for x in data[1:]]

        # extract label, truncation, occlusion
        self.type = data[0] # 'Car', 'Pedestrian', ...
        self.truncation = data[1] # truncated pixel ratio [0..1]
        self.occlusion = int(data[2]) # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown
        self.alpha = data[3] # object observation angle [-pi..pi]

        # extract 2d bounding box in 0-based coordinates
        self.xmin = data[4] # left
        self.ymin = data[5] # top
        self.xmax = data[6] # right
        self.ymax = data[7] # bottom
        self.box2d = np.array([self.xmin,self.ymin,self.xmax,self.ymax])
        
        # extract 3d bounding box information
        self.h = data[8] # box height
        self.w = data[9] # box width
        self.l = data[10] # box length (in meters)
        self.t = (data[11],data[12],data[13]) # location (x,y,z) in camera coord.
        # self.ry is the rotation angle around the y-axis
        self.ry = data[14] # yaw angle (around Y-axis in camera coordinates) [-pi..pi]
        if len(data) == 16:
            self.score = data[15]

    def print_object(self):
        print('Type, truncation, occlusion, alpha: %s, %d, %d, %f' % \
            (self.type, self.truncation, self.occlusion, self.alpha))
        print('2d bbox (x0,y0,x1,y1): %f, %f, %f, %f' % \
            (self.xmin, self.ymin, self.xmax, self.ymax))
        print('3d bbox h,w,l: %f, %f, %f' % \
            (self.h, self.w, self.l))
        print('3d bbox location, ry: (%f, %f, %f), %f' % \
            (self.t[0],self.t[1],self.t[2],self.ry))


class Calibration(object):
    ''' Calibration matrices and utils
        # this is used to multi the 
        3d XYZ in <label>.txt are in rect camera coord.
        2d box xy are in image2 coord
        Points in <lidar>.bin are in Velodyne coord.

        y_image2 = P^2_rect * x_rect
        y_image2 = P^2_rect * R0_rect * Tr_velo_to_cam * x_velo
        x_ref = Tr_velo_to_cam * x_velo
        x_rect = R0_rect * x_ref

        P^2_rect = [f^2_u,  0,      c^2_u,  -f^2_u b^2_x;
                    0,      f^2_v,  c^2_v,  -f^2_v b^2_y;
                    0,      0,      1,      0]
                 = K * [1|t]

        image2 coord:
         ----> x-axis (u)
        |
        |
        v y-axis (v)

        velodyne coord:
        front x, left y, up z

        rect/ref camera coord:
        right x, down y, front z

        Ref (KITTI paper): http://www.cvlibs.net/publications/Geiger2013IJRR.pdf

        TODO(rqi): do matrix multiplication only once for each projection.
    '''
    def __init__(self, calib_filepath, from_video=False, calib=None):
        if calib:
            calibs = calib
        elif from_video:
            calibs = self.read_calib_from_video(calib_filepath)
        else:
            calibs = self.read_calib_file(calib_filepath)
        # Projection matrix from rect camera coord to image2 coord
        # P2 means the left color image in the kitti dataset
        # which is the images we used
        self.P = calibs['P2'] 
        self.P = np.reshape(self.P, [3, 4])
        # Rigid transform from Velodyne coord to reference camera coord
        # the velo_to_cam matrix
        self.V2C = calibs['Tr_velo_to_cam']
        self.V2C = np.reshape(self.V2C, [3,4])
        # the cam_to_velo matrix [3, 4]
        self.C2V = inverse_rigid_trans(self.V2C)
        # Rotation from reference camera coord to rect camera coord
        # rotating the reference camera coord from non-rect coord to rect camera coord
        # thus, we could print the lidar point by multi the rect lidar points
        self.R0 = calibs['R0_rect']
        self.R0 = np.reshape(self.R0,[3,3])

        # Camera intrinsics and extrinsics
        self.c_u = self.P[0,2]
        self.c_v = self.P[1,2]
        self.f_u = self.P[0,0]
        self.f_v = self.P[1,1]
        self.b_x = self.P[0,3]/(-self.f_u) # relative 
        self.b_y = self.P[1,3]/(-self.f_v)

    def read_calib_file(self, filepath):
        ''' Read in a calibration file and parse into a dictionary.
        Ref: https://github.com/utiasSTARS/pykitti/blob/master/pykitti/utils.py
        read_calib_file: return the matrix in the %06d.txt
        '''
        data = {}
        with open(filepath, 'r') as f:
            for line in f.readlines():
                line = line.rstrip()
                if len(line)==0: continue
                key, value = line.split(':', 1)
                # The only non-float values in these files are dates, which
                # we don't care about anyway
                try:
                    data[key] = np.array([float(x) for x in value.split()])
                except ValueError:
                    pass

        return data
    
    def read_calib_from_video(self, calib_root_dir):
        ''' Read calibration for camera 2 from video calib files.
            there are calib_cam_to_cam and calib_velo_to_cam under the calib_root_dir
        '''
        data = {}
        cam2cam = self.read_calib_file(os.path.join(calib_root_dir, 'calib_cam_to_cam.txt'))
        velo2cam = self.read_calib_file(os.path.join(calib_root_dir, 'calib_velo_to_cam.txt'))
        Tr_velo_to_cam = np.zeros((3,4))
        Tr_velo_to_cam[0:3,0:3] = np.reshape(velo2cam['R'], [3,3])
        Tr_velo_to_cam[:,3] = velo2cam['T']
        data['Tr_velo_to_cam'] = np.reshape(Tr_velo_to_cam, [12])
        data['R0_rect'] = cam2cam['R_rect_00']
        data['P2'] = cam2cam['P_rect_02']
        return data

    def cart2hom(self, pts_3d):
        ''' Input: nx3 points in Cartesian
            Oupput: nx4 points in Homogeneous by pending 1
            make it from n x 3 to n x 4
        '''
        n = pts_3d.shape[0]
        pts_3d_hom = np.hstack((pts_3d, np.ones((n,1))))
        return pts_3d_hom
 
    # =========================== 
    # ------- 3d to 3d ---------- 
    # =========================== 
    def project_velo_to_ref(self, pts_3d_velo):
        pts_3d_velo = self.cart2hom(pts_3d_velo) # nx4
        # return nx3 points, in the reference image plane
        # nx3 output
        return np.dot(pts_3d_velo, np.transpose(self.V2C))

    def project_ref_to_velo(self, pts_3d_ref):
        # nx3 output
        pts_3d_ref = self.cart2hom(pts_3d_ref) # nx4
        return np.dot(pts_3d_ref, np.transpose(self.C2V))

    def project_rect_to_ref(self, pts_3d_rect):
        ''' Input and Output are nx3 points '''
        return np.transpose(np.dot(np.linalg.inv(self.R0), np.transpose(pts_3d_rect)))
    
    def project_ref_to_rect(self, pts_3d_ref):
        ''' Input and Output are nx3 points '''
        return np.transpose(np.dot(self.R0, np.transpose(pts_3d_ref)))
 
    def project_rect_to_velo(self, pts_3d_rect):
        ''' Input: nx3 points in rect camera coord.
            Output: nx3 points in velodyne coord.
        ''' 
        pts_3d_ref = self.project_rect_to_ref(pts_3d_rect)
        return self.project_ref_to_velo(pts_3d_ref)

    def project_velo_to_rect(self, pts_3d_velo):
        pts_3d_ref = self.project_velo_to_ref(pts_3d_velo)
        return self.project_ref_to_rect(pts_3d_ref)

    # =========================== 
    # ------- 3d to 2d ---------- 
    # =========================== 
    def project_rect_to_image(self, pts_3d_rect):
        ''' Input: nx3 points in rect camera coord.
            Output: nx2 points in image2 coord.
        '''
        pts_3d_rect = self.cart2hom(pts_3d_rect)
        pts_2d = np.dot(pts_3d_rect, np.transpose(self.P)) # nx3
        # pts_2d[:, 2], means the depth in the uv-map
        # thus we need to div it 
        pts_2d[:,0] /= pts_2d[:,2]
        pts_2d[:,1] /= pts_2d[:,2]
        return pts_2d[:,0:2]
    
    def project_velo_to_image(self, pts_3d_velo):
        ''' Input: nx3 points in velodyne coord.
            Output: nx2 points in image2 coord.
        '''
        pts_3d_rect = self.project_velo_to_rect(pts_3d_velo)
        return self.project_rect_to_image(pts_3d_rect)

    # =========================== 
    # ------- 2d to 3d ---------- 
    # =========================== 
    def project_image_to_rect(self, uv_depth):
        ''' Input: nx3 first two channels are uv, 3rd channel
                   is depth in rect camera coord.
            Output: nx3 points in rect camera coord.
        '''
        # first we multi the depth(uv_map[:, 2]) back, and get the transform
        n = uv_depth.shape[0]
        x = ((uv_depth[:,0]-self.c_u)*uv_depth[:,2])/self.f_u + self.b_x
        y = ((uv_depth[:,1]-self.c_v)*uv_depth[:,2])/self.f_v + self.b_y
        pts_3d_rect = np.zeros((n,3))
        pts_3d_rect[:,0] = x
        pts_3d_rect[:,1] = y
        pts_3d_rect[:,2] = uv_depth[:,2]
        return pts_3d_rect

    def project_image_to_velo(self, uv_depth):
        pts_3d_rect = self.project_image_to_rect(uv_depth)
        return self.project_rect_to_velo(pts_3d_rect)


def in_hull(p, hull):
    from scipy.spatial import Delaunay
    if not isinstance(hull,Delaunay):
        hull = Delaunay(hull)
    return hull.find_simplex(p)>=0

def extract_pc_in_box3d(pc, box3d):
    ''' pc: (N,3), box3d: (8,3) '''
    box3d_roi_inds = in_hull(pc[:,0:3], box3d)
    box3d_roi_logits = np.zeros([pc.shape[0], 2], np.float32)
    box3d_roi_logits[box3d_roi_inds, 1] = 1
    return box3d_roi_logits 

def extract_pc_in_box2d(pc, box2d):
    ''' pc: (N, 3), box2d: (4, 2) '''
    box2d_roi_inds = in_hull(pc[:,[0, 2]], box2d)
    box2d_roi_logits = np.zeros([pc.shape[0]], np.float32)
    box2d_roi_logits[box2d_roi_inds] = 1
    return box2d_roi_logits


def transform_from_rot_trans(R, t):
    ''' Transforation matrix from rotation matrix and translation vector. '''
    R = R.reshape(3, 3)
    t = t.reshape(3, 1)
    return np.vstack((np.hstack([R, t]), [0, 0, 0, 1]))


def inverse_rigid_trans(Tr):
    ''' Inverse a rigid body transform matrix (3x4 as [R|t])
        [R'|-R't; 0|1]
    '''
    inv_Tr = np.zeros_like(Tr) # 3x4
    inv_Tr[0:3,0:3] = np.transpose(Tr[0:3,0:3])
    inv_Tr[0:3,3] = np.dot(-np.transpose(Tr[0:3,0:3]), Tr[0:3,3])
    return inv_Tr

def read_label(label_filename):
    # first get the line in the gt
    lines = [line.rstrip() for line in open(label_filename)]
    # lines represent the instance object
    objects = [Object3d(line) for line in lines]
    return objects

def get_road_plane(img_idx, planes_dir):
    """Reads the road plane from file

    :param int img_idx : Index of image
    :param str planes_dir : directory containing plane text files

    :return plane : List containing plane equation coefficients
    """

    plane_file = planes_dir + '/%06d.txt' % img_idx

    with open(plane_file, 'r') as input_file:
        lines = input_file.readlines()
        input_file.close()

    # Plane coefficients stored in 4th row
    lines = lines[3].split()

    # Convert str to float
    lines = [float(i) for i in lines]

    plane = np.asarray(lines)

    # Ensure normal is always facing up.
    # In Kitti's frame of reference, +y is down
    if plane[1] > 0:
        plane = -plane

    # Normalize the plane coefficients
    norm = np.linalg.norm(plane[0:3])
    plane = plane / norm

    return plane

def load_image(img_filename):
    return cv2.imread(img_filename)

def load_velo_scan(velo_filename):
    scan = np.fromfile(velo_filename, dtype=np.float32)
    scan = scan.reshape((-1, 4))
    return scan

####################### ProjectToImage ################################
def project_to_image(pts_3d, P):
    ''' Project 3d points to image plane.

    Usage: pts_2d = projectToImage(pts_3d, P)
      input: pts_3d: nx3 matrix
             P:      3x4 projection matrix
      output: pts_2d: nx2 matrix

      P(3x4) dot pts_3d_extended(4xn) = projected_pts_2d(3xn)
      => normalize projected_pts_2d(2xn)

      <=> pts_3d_extended(nx4) dot P'(4x3) = projected_pts_2d(nx3)
          => normalize projected_pts_2d(nx2)
    '''
    n = pts_3d.shape[0]
    pts_3d_extend = np.hstack((pts_3d, np.ones((n,1))))
    pts_2d = np.dot(pts_3d_extend, np.transpose(P)) # nx3
    pts_2d[:,0] /= pts_2d[:,2]
    pts_2d[:,1] /= pts_2d[:,2]
    return pts_2d[:,0:2]

def tf_project_to_image_tensor(pts_3d, P):
    """Projects 3D points to 2D points in image space.

    Args:
        points_3d: a list of float32 tensor of shape [-1, 3]
        cam_p2_matrix: a float32 tensor of shape [3, 4] representing
            the camera matrix.

    Returns:
        points_2d: a list of float32 tensor of shape [-1, 2]
            This is the projected 3D points into 2D .i.e. corresponding
            3D points in image coordinates.
    """
    n = tf.shape(pts_3d)[0]
    pts_3d_extend = tf.concat([pts_3d, tf.ones([n, 1], dtype=pts_3d.dtype)], axis=-1) 
    pts_2d = tf.matmul(pts_3d_extend, tf.transpose(P)) # nx3
    pts_2d_x = pts_2d[:, 0] / pts_2d[:, 2]
    pts_2d_y = pts_2d[:, 1] / pts_2d[:, 2]
    pts_2d = tf.stack([pts_2d_x, pts_2d_y], axis=-1) # nx2

    return pts_2d 
####################### ProjectToImage ################################


def draw_projected_box3d(image, qs, color=(255, 255, 255), thickness=2):
    ''' Draw 3d bounding box in image
        qs: (8,3) array of vertices for the 3d box in following order:
            1 -------- 0
           /|         /|
          2 -------- 3 .
          | |        | |
          . 5 -------- 4
          |/         |/
          6 -------- 7
    '''
    qs = qs.astype(np.int32)
    for k in range(0,4):
       # Ref: http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html
       i,j=k,(k+1)%4
       # use LINE_AA for opencv3
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

       i,j=k+4,(k+1)%4 + 4
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

       i,j=k,k+4
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)
    return image


""" Helper methods for loading and parsing KITTI data.

Author: <NAME>
Date: September 2017
"""
from __future__ import print_function

import tensorflow as tf
import numpy as np
import cv2
import os


class Object3d(object):
    """ 3d object label """
    def __init__(self, label_file_line=None):
        # each object represents a line in the %06d.txt
        # that is each object3D represent a instance object in the image
        # Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 
        # 1.89 0.48 1.20 1.84 1.47 8.41 0.01
        data = label_file_line.split(' ')
        data[1:] = [float(x) for x in data[1:]]

        # extract label, truncation, occlusion
        self.type = data[0] # 'Car', 'Pedestrian', ...
        self.truncation = data[1] # truncated pixel ratio [0..1]
        self.occlusion = int(data[2]) # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown
        self.alpha = data[3] # object observation angle [-pi..pi]

        # extract 2d bounding box in 0-based coordinates
        self.xmin = data[4] # left
        self.ymin = data[5] # top
        self.xmax = data[6] # right
        self.ymax = data[7] # bottom
        self.box2d = np.array([self.xmin,self.ymin,self.xmax,self.ymax])
        
        # extract 3d bounding box information
        self.h = data[8] # box height
        self.w = data[9] # box width
        self.l = data[10] # box length (in meters)
        self.t = (data[11],data[12],data[13]) # location (x,y,z) in camera coord.
        # self.ry is the rotation angle around the y-axis
        self.ry = data[14] # yaw angle (around Y-axis in camera coordinates) [-pi..pi]
        if len(data) == 16:
            self.score = data[15]

    def print_object(self):
        print('Type, truncation, occlusion, alpha: %s, %d, %d, %f' % \
            (self.type, self.truncation, self.occlusion, self.alpha))
        print('2d bbox (x0,y0,x1,y1): %f, %f, %f, %f' % \
            (self.xmin, self.ymin, self.xmax, self.ymax))
        print('3d bbox h,w,l: %f, %f, %f' % \
            (self.h, self.w, self.l))
        print('3d bbox location, ry: (%f, %f, %f), %f' % \
            (self.t[0],self.t[1],self.t[2],self.ry))


class Calibration(object):
    ''' Calibration matrices and utils
        # this is used to multi the 
        3d XYZ in <label>.txt are in rect camera coord.
        2d box xy are in image2 coord
        Points in <lidar>.bin are in Velodyne coord.

        y_image2 = P^2_rect * x_rect
        y_image2 = P^2_rect * R0_rect * Tr_velo_to_cam * x_velo
        x_ref = Tr_velo_to_cam * x_velo
        x_rect = R0_rect * x_ref

        P^2_rect = [f^2_u,  0,      c^2_u,  -f^2_u b^2_x;
                    0,      f^2_v,  c^2_v,  -f^2_v b^2_y;
                    0,      0,      1,      0]
                 = K * [1|t]

        image2 coord:
         ----> x-axis (u)
        |
        |
        v y-axis (v)

        velodyne coord:
        front x, left y, up z

        rect/ref camera coord:
        right x, down y, front z

        Ref (KITTI paper): http://www.cvlibs.net/publications/Geiger2013IJRR.pdf

        TODO(rqi): do matrix multiplication only once for each projection.
    '''
    def __init__(self, calib_filepath, from_video=False, calib=None):
        if calib:
            calibs = calib
        elif from_video:
            calibs = self.read_calib_from_video(calib_filepath)
        else:
            calibs = self.read_calib_file(calib_filepath)
        # Projection matrix from rect camera coord to image2 coord
        # P2 means the left color image in the kitti dataset
        # which is the images we used
        self.P = calibs['P2'] 
        self.P = np.reshape(self.P, [3, 4])
        # Rigid transform from Velodyne coord to reference camera coord
        # the velo_to_cam matrix
        self.V2C = calibs['Tr_velo_to_cam']
        self.V2C = np.reshape(self.V2C, [3,4])
        # the cam_to_velo matrix [3, 4]
        self.C2V = inverse_rigid_trans(self.V2C)
        # Rotation from reference camera coord to rect camera coord
        # rotating the reference camera coord from non-rect coord to rect camera coord
        # thus, we could print the lidar point by multi the rect lidar points
        self.R0 = calibs['R0_rect']
        self.R0 = np.reshape(self.R0,[3,3])

        # Camera intrinsics and extrinsics
        self.c_u = self.P[0,2]
        self.c_v = self.P[1,2]
        self.f_u = self.P[0,0]
        self.f_v = self.P[1,1]
        self.b_x = self.P[0,3]/(-self.f_u) # relative 
        self.b_y = self.P[1,3]/(-self.f_v)

    def read_calib_file(self, filepath):
        ''' Read in a calibration file and parse into a dictionary.
        Ref: https://github.com/utiasSTARS/pykitti/blob/master/pykitti/utils.py
        read_calib_file: return the matrix in the %06d.txt
        '''
        data = {}
        with open(filepath, 'r') as f:
            for line in f.readlines():
                line = line.rstrip()
                if len(line)==0: continue
                key, value = line.split(':', 1)
                # The only non-float values in these files are dates, which
                # we don't care about anyway
                try:
                    data[key] = np.array([float(x) for x in value.split()])
                except ValueError:
                    pass

        return data
    
    def read_calib_from_video(self, calib_root_dir):
        ''' Read calibration for camera 2 from video calib files.
            there are calib_cam_to_cam and calib_velo_to_cam under the calib_root_dir
        '''
        data = {}
        cam2cam = self.read_calib_file(os.path.join(calib_root_dir, 'calib_cam_to_cam.txt'))
        velo2cam = self.read_calib_file(os.path.join(calib_root_dir, 'calib_velo_to_cam.txt'))
        Tr_velo_to_cam = np.zeros((3,4))
        Tr_velo_to_cam[0:3,0:3] = np.reshape(velo2cam['R'], [3,3])
        Tr_velo_to_cam[:,3] = velo2cam['T']
        data['Tr_velo_to_cam'] = np.reshape(Tr_velo_to_cam, [12])
        data['R0_rect'] = cam2cam['R_rect_00']
        data['P2'] = cam2cam['P_rect_02']
        return data

    def cart2hom(self, pts_3d):
        ''' Input: nx3 points in Cartesian
            Oupput: nx4 points in Homogeneous by pending 1
            make it from n x 3 to n x 4
        '''
        n = pts_3d.shape[0]
        pts_3d_hom = np.hstack((pts_3d, np.ones((n,1))))
        return pts_3d_hom
 
    # =========================== 
    # ------- 3d to 3d ---------- 
    # =========================== 
    def project_velo_to_ref(self, pts_3d_velo):
        pts_3d_velo = self.cart2hom(pts_3d_velo) # nx4
        # return nx3 points, in the reference image plane
        # nx3 output
        return np.dot(pts_3d_velo, np.transpose(self.V2C))

    def project_ref_to_velo(self, pts_3d_ref):
        # nx3 output
        pts_3d_ref = self.cart2hom(pts_3d_ref) # nx4
        return np.dot(pts_3d_ref, np.transpose(self.C2V))

    def project_rect_to_ref(self, pts_3d_rect):
        ''' Input and Output are nx3 points '''
        return np.transpose(np.dot(np.linalg.inv(self.R0), np.transpose(pts_3d_rect)))
    
    def project_ref_to_rect(self, pts_3d_ref):
        ''' Input and Output are nx3 points '''
        return np.transpose(np.dot(self.R0, np.transpose(pts_3d_ref)))
 
    def project_rect_to_velo(self, pts_3d_rect):
        ''' Input: nx3 points in rect camera coord.
            Output: nx3 points in velodyne coord.
        ''' 
        pts_3d_ref = self.project_rect_to_ref(pts_3d_rect)
        return self.project_ref_to_velo(pts_3d_ref)

    def project_velo_to_rect(self, pts_3d_velo):
        pts_3d_ref = self.project_velo_to_ref(pts_3d_velo)
        return self.project_ref_to_rect(pts_3d_ref)

    # =========================== 
    # ------- 3d to 2d ---------- 
    # =========================== 
    def project_rect_to_image(self, pts_3d_rect):
        ''' Input: nx3 points in rect camera coord.
            Output: nx2 points in image2 coord.
        '''
        pts_3d_rect = self.cart2hom(pts_3d_rect)
        pts_2d = np.dot(pts_3d_rect, np.transpose(self.P)) # nx3
        # pts_2d[:, 2], means the depth in the uv-map
        # thus we need to div it 
        pts_2d[:,0] /= pts_2d[:,2]
        pts_2d[:,1] /= pts_2d[:,2]
        return pts_2d[:,0:2]
    
    def project_velo_to_image(self, pts_3d_velo):
        ''' Input: nx3 points in velodyne coord.
            Output: nx2 points in image2 coord.
        '''
        pts_3d_rect = self.project_velo_to_rect(pts_3d_velo)
        return self.project_rect_to_image(pts_3d_rect)

    # =========================== 
    # ------- 2d to 3d ---------- 
    # =========================== 
    def project_image_to_rect(self, uv_depth):
        ''' Input: nx3 first two channels are uv, 3rd channel
                   is depth in rect camera coord.
            Output: nx3 points in rect camera coord.
        '''
        # first we multi the depth(uv_map[:, 2]) back, and get the transform
        n = uv_depth.shape[0]
        x = ((uv_depth[:,0]-self.c_u)*uv_depth[:,2])/self.f_u + self.b_x
        y = ((uv_depth[:,1]-self.c_v)*uv_depth[:,2])/self.f_v + self.b_y
        pts_3d_rect = np.zeros((n,3))
        pts_3d_rect[:,0] = x
        pts_3d_rect[:,1] = y
        pts_3d_rect[:,2] = uv_depth[:,2]
        return pts_3d_rect

    def project_image_to_velo(self, uv_depth):
        pts_3d_rect = self.project_image_to_rect(uv_depth)
        return self.project_rect_to_velo(pts_3d_rect)


def in_hull(p, hull):
    from scipy.spatial import Delaunay
    if not isinstance(hull,Delaunay):
        hull = Delaunay(hull)
    return hull.find_simplex(p)>=0

def extract_pc_in_box3d(pc, box3d):
    ''' pc: (N,3), box3d: (8,3) '''
    box3d_roi_inds = in_hull(pc[:,0:3], box3d)
    box3d_roi_logits = np.zeros([pc.shape[0], 2], np.float32)
    box3d_roi_logits[box3d_roi_inds, 1] = 1
    return box3d_roi_logits 

def extract_pc_in_box2d(pc, box2d):
    ''' pc: (N, 3), box2d: (4, 2) '''
    box2d_roi_inds = in_hull(pc[:,[0, 2]], box2d)
    box2d_roi_logits = np.zeros([pc.shape[0]], np.float32)
    box2d_roi_logits[box2d_roi_inds] = 1
    return box2d_roi_logits


def transform_from_rot_trans(R, t):
    ''' Transforation matrix from rotation matrix and translation vector. '''
    R = R.reshape(3, 3)
    t = t.reshape(3, 1)
    return np.vstack((np.hstack([R, t]), [0, 0, 0, 1]))


def inverse_rigid_trans(Tr):
    ''' Inverse a rigid body transform matrix (3x4 as [R|t])
        [R'|-R't; 0|1]
    '''
    inv_Tr = np.zeros_like(Tr) # 3x4
    inv_Tr[0:3,0:3] = np.transpose(Tr[0:3,0:3])
    inv_Tr[0:3,3] = np.dot(-np.transpose(Tr[0:3,0:3]), Tr[0:3,3])
    return inv_Tr

def read_label(label_filename):
    # first get the line in the gt
    lines = [line.rstrip() for line in open(label_filename)]
    # lines represent the instance object
    objects = [Object3d(line) for line in lines]
    return objects

def get_road_plane(img_idx, planes_dir):
    """Reads the road plane from file

    :param int img_idx : Index of image
    :param str planes_dir : directory containing plane text files

    :return plane : List containing plane equation coefficients
    """

    plane_file = planes_dir + '/%06d.txt' % img_idx

    with open(plane_file, 'r') as input_file:
        lines = input_file.readlines()
        input_file.close()

    # Plane coefficients stored in 4th row
    lines = lines[3].split()

    # Convert str to float
    lines = [float(i) for i in lines]

    plane = np.asarray(lines)

    # Ensure normal is always facing up.
    # In Kitti's frame of reference, +y is down
    if plane[1] > 0:
        plane = -plane

    # Normalize the plane coefficients
    norm = np.linalg.norm(plane[0:3])
    plane = plane / norm

    return plane

def load_image(img_filename):
    return cv2.imread(img_filename)

def load_velo_scan(velo_filename):
    scan = np.fromfile(velo_filename, dtype=np.float32)
    scan = scan.reshape((-1, 4))
    return scan

####################### ProjectToImage ################################
def project_to_image(pts_3d, P):
    ''' Project 3d points to image plane.

    Usage: pts_2d = projectToImage(pts_3d, P)
      input: pts_3d: nx3 matrix
             P:      3x4 projection matrix
      output: pts_2d: nx2 matrix

      P(3x4) dot pts_3d_extended(4xn) = projected_pts_2d(3xn)
      => normalize projected_pts_2d(2xn)

      <=> pts_3d_extended(nx4) dot P'(4x3) = projected_pts_2d(nx3)
          => normalize projected_pts_2d(nx2)
    '''
    n = pts_3d.shape[0]
    pts_3d_extend = np.hstack((pts_3d, np.ones((n,1))))
    pts_2d = np.dot(pts_3d_extend, np.transpose(P)) # nx3
    pts_2d[:,0] /= pts_2d[:,2]
    pts_2d[:,1] /= pts_2d[:,2]
    return pts_2d[:,0:2]

def tf_project_to_image_tensor(pts_3d, P):
    """Projects 3D points to 2D points in image space.

    Args:
        points_3d: a list of float32 tensor of shape [-1, 3]
        cam_p2_matrix: a float32 tensor of shape [3, 4] representing
            the camera matrix.

    Returns:
        points_2d: a list of float32 tensor of shape [-1, 2]
            This is the projected 3D points into 2D .i.e. corresponding
            3D points in image coordinates.
    """
    n = tf.shape(pts_3d)[0]
    pts_3d_extend = tf.concat([pts_3d, tf.ones([n, 1], dtype=pts_3d.dtype)], axis=-1) 
    pts_2d = tf.matmul(pts_3d_extend, tf.transpose(P)) # nx3
    pts_2d_x = pts_2d[:, 0] / pts_2d[:, 2]
    pts_2d_y = pts_2d[:, 1] / pts_2d[:, 2]
    pts_2d = tf.stack([pts_2d_x, pts_2d_y], axis=-1) # nx2

    return pts_2d 
####################### ProjectToImage ################################


def draw_projected_box3d(image, qs, color=(255, 255, 255), thickness=2):
    ''' Draw 3d bounding box in image
        qs: (8,3) array of vertices for the 3d box in following order:
            1 -------- 0
           /|         /|
          2 -------- 3 .
          | |        | |
          . 5 -------- 4
          |/         |/
          6 -------- 7
    '''
    qs = qs.astype(np.int32)
    for k in range(0,4):
       # Ref: http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html
       i,j=k,(k+1)%4
       # use LINE_AA for opencv3
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

       i,j=k+4,(k+1)%4 + 4
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

       i,j=k,k+4
       cv2.line(image, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)
    return image


Helper methods for loading and parsing KITTI data.  Author: <NAME> Date: September 2017 3d object label # each object represents a line in the %06d.txt # that is each object3D represent a instance object in the image # Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 # 1.89 0.48 1.20 1.84 1.47 8.41 0.01 # extract label, truncation, occlusion # 'Car', 'Pedestrian', ... # truncated pixel ratio [0..1] # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown # object observation angle [-pi..pi] # extract 2d bounding box in 0-based coordinates # left # top # right # bottom # extract 3d bounding box information # box height # box width # box length (in meters) # location (x,y,z) in camera coord. # self.ry is the rotation angle around the y-axis # yaw angle (around Y-axis in camera coordinates) [-pi..pi] Calibration matrices and utils         # this is used to multi the          3d XYZ in <label>.txt are in rect camera coord.         2d box xy are in image2 coord         Points in <lidar>.bin are in Velodyne coord.          y_image2 = P^2_rect * x_rect         y_image2 = P^2_rect * R0_rect * Tr_velo_to_cam * x_velo         x_ref = Tr_velo_to_cam * x_velo         x_rect = R0_rect * x_ref          P^2_rect = [f^2_u,  0,      c^2_u,  -f^2_u b^2_x;                     0,      f^2_v,  c^2_v,  -f^2_v b^2_y;                     0,      0,      1,      0]                  = K * [1|t]          image2 coord:          ----> x-axis (u)         |         |         v y-axis (v)          velodyne coord:         front x, left y, up z          rect/ref camera coord:         right x, down y, front z          Ref (KITTI paper): http://www.cvlibs.net/publications/Geiger2013IJRR.pdf          TODO(rqi): do matrix multiplication only once for each projection. # Projection matrix from rect camera coord to image2 coord # P2 means the left color image in the kitti dataset # which is the images we used # Rigid transform from Velodyne coord to reference camera coord # the velo_to_cam matrix # the cam_to_velo matrix [3, 4] # Rotation from reference camera coord to rect camera coord # rotating the reference camera coord from non-rect coord to rect camera coord # thus, we could print the lidar point by multi the rect lidar points # Camera intrinsics and extrinsics # relative Read in a calibration file and parse into a dictionary.         Ref: https://github.com/utiasSTARS/pykitti/blob/master/pykitti/utils.py         read_calib_file: return the matrix in the %06d.txt # The only non-float values in these files are dates, which # we don't care about anyway Read calibration for camera 2 from video calib files.             there are calib_cam_to_cam and calib_velo_to_cam under the calib_root_dir Input: nx3 points in Cartesian             Oupput: nx4 points in Homogeneous by pending 1             make it from n x 3 to n x 4 # =========================== # ------- 3d to 3d ---------- # =========================== # nx4 # return nx3 points, in the reference image plane # nx3 output # nx3 output # nx4 Input and Output are nx3 points Input and Output are nx3 points Input: nx3 points in rect camera coord.             Output: nx3 points in velodyne coord. # =========================== # ------- 3d to 2d ---------- # =========================== Input: nx3 points in rect camera coord.             Output: nx2 points in image2 coord. # nx3 # pts_2d[:, 2], means the depth in the uv-map # thus we need to div it Input: nx3 points in velodyne coord.             Output: nx2 points in image2 coord. # =========================== # ------- 2d to 3d ---------- # =========================== Input: nx3 first two channels are uv, 3rd channel                    is depth in rect camera coord.             Output: nx3 points in rect camera coord. # first we multi the depth(uv_map[:, 2]) back, and get the transform pc: (N,3), box3d: (8,3) pc: (N, 3), box2d: (4, 2) Transforation matrix from rotation matrix and translation vector. Inverse a rigid body transform matrix (3x4 as [R|t])         [R'|-R't; 0|1] # 3x4 # first get the line in the gt # lines represent the instance object Reads the road plane from file      :param int img_idx : Index of image     :param str planes_dir : directory containing plane text files      :return plane : List containing plane equation coefficients # Plane coefficients stored in 4th row # Convert str to float # Ensure normal is always facing up. # In Kitti's frame of reference, +y is down # Normalize the plane coefficients ####################### ProjectToImage ################################ Project 3d points to image plane.      Usage: pts_2d = projectToImage(pts_3d, P)       input: pts_3d: nx3 matrix              P:      3x4 projection matrix       output: pts_2d: nx2 matrix        P(3x4) dot pts_3d_extended(4xn) = projected_pts_2d(3xn)       => normalize projected_pts_2d(2xn)        <=> pts_3d_extended(nx4) dot P'(4x3) = projected_pts_2d(nx3)           => normalize projected_pts_2d(nx2) # nx3 Projects 3D points to 2D points in image space.      Args:         points_3d: a list of float32 tensor of shape [-1, 3]         cam_p2_matrix: a float32 tensor of shape [3, 4] representing             the camera matrix.      Returns:         points_2d: a list of float32 tensor of shape [-1, 2]             This is the projected 3D points into 2D .i.e. corresponding             3D points in image coordinates. # nx3 # nx2 ####################### ProjectToImage ################################ Draw 3d bounding box in image         qs: (8,3) array of vertices for the 3d box in following order:             1 -------- 0            /|         /|           2 -------- 3 .           | |        | |           . 5 -------- 4           |/         |/           6 -------- 7 # Ref: http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html # use LINE_AA for opencv3