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

Author: <NAME>
Date: February 2020
"""
from __future__ import print_function

import numpy as np
import cv2

from PIL import Image

transform_mat = {"Scene01": {"15-deg-left": -6/12.0, "30-deg-left":-5/12.0,
                             "15-deg-right":-8/12.0, "30-deg-right":-9/12.0,
                             "clone":-7/12.0, "morning":-7/12.0, "rain":-7/12.0,
                             "fog":-7/12.0, "overcast":-7/12.0, "sunset":-7/12.0},
                 "Scene02": {"15-deg-left": 4/12.0, "30-deg-left": 5/12.0,
                             "15-deg-right":2/12.0, "30-deg-right":1/12.0,
                             "clone":3/12.0, "morning": 3/12.0, "rain": 3/12.0,
                             "fog":3/12.0, "overcast": 3/12.0, "sunset": 3/12.0},
                 "Scene06": {"15-deg-left": 0/12.0, "30-deg-left": 1/12.0,
                             "15-deg-right": -2/12.0, "30-deg-right":-3/12.0,
                             "clone":-1/12.0, "morning": -1/12.0, "rain": -1/12.0,
                             "fog":-1/12.0, "overcast": -1/12.0, "sunset": -1/12.0},
                 "Scene20": {"15-deg-left": -9/12.0, "30-deg-left": -8/12.0,
                             "15-deg-right": -11/12.0, "30-deg-right":-12/12.0,
                             "clone":-10/12.0, "morning": -10/12.0, "rain": -10/12.0,
                             "fog":-10/12.0, "overcast": -10/12.0, "sunset": -10/12.0},
                 "Scene18": {"15-deg-left": -1.5/12.0, "30-deg-left": -0.5/12.0,
                             "15-deg-right": -3.5/12.0, "30-deg-right":-4.5/12.0,
                             "clone": -2.5/12.0, "morning": -2.5/12.0, "rain": -2.5/12.0,
                             "fog": -2.5/12.0, "overcast": -2.5/12.0, "sunset": -2.5/12.0}}


class Object3d(object):
    """ 3d object label """

    def __init__(self, data):
        """
        Process one object instance of a frame
        :param data: Series object for one object instance
        """
        # extract label, truncation, occlusion
        self.type = data["label"]  # 'Car', 'Van', 'Truck', etc ...
        self.truncation = data[
            "truncation_ratio"]  # truncated pixel ratio [0..1]
        self.occlusion = 1 - data[
            "occupancy_ratio"]  # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown
        self.alpha = data["alpha"]  # object observation angle [-pi..pi]

        # extract 2d bounding box in 0-based coordinates
        self.xmin = data["left"]  # left
        self.ymin = data["top"]  # top
        self.xmax = data["right"]  # right
        self.ymax = data["bottom"]  # bottom
        self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])

        # extract 3d bounding box information
        self.h = data["height"]  # box height
        self.w = data["width"]  # box width
        self.l = data["length"]  # box length (in meters)
        self.t = (data["camera_space_X"], data["camera_space_Y"],
                  data["camera_space_Z"])  # location (x,y,z) in camera coord.
        self.ry = data["rotation_camera_space_y"]  # yaw angle
        # (around Y-axis in camera coordinates) [-pi..pi]

    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
        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
    """

    def __init__(self, scene, subscene, intrinsic_values, extrinsic_values):
        # Projection matrix from rect camera coord to image2 coord
        self.P = self.process_intrinsic(intrinsic_values)

        # Rotation from reference camera coord to rect camera coord
        self.extrinsics = self.process_extrinsic(extrinsic_values)
        self.R0 = np.linalg.inv(self.extrinsics[:3, :3])
        # R0= "9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01"
        # self.R0 = np.reshape(np.array([float(x) for x in R0.split()]), [3,3])

        # transform = {"Scene01": -np.pi / 2.0, "Scene02": np.pi/4.0}
        velo = "7.533745000000e-03 -9.999714000000e-01 -6.166020000000e-04 -4.069766000000e-03 1.480249000000e-02 7.280733000000e-04 -9.998902000000e-01 -7.631618000000e-02 9.998621000000e-01 7.523790000000e-03 1.480755000000e-02 -2.717806000000e-01"

        self.V2C = np.array([float(x) for x in velo.split()])
        self.V2C = np.reshape(self.V2C, [3,4])
        self.C2V = inverse_rigid_trans(self.V2C)
        angle = transform_mat[scene][subscene] * np.pi # morning
        self.C2V = np.dot(rotz(angle), inverse_rigid_trans(self.V2C))
        self.V2C = inverse_rigid_trans(self.C2V)


        # 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)

    @staticmethod
    def process_intrinsic(intrinsic_values):
        K = np.zeros((3, 4))
        K[0][0] = intrinsic_values["K[0,0]"]
        K[1][1] = intrinsic_values["K[1,1]"]
        K[1][2] = intrinsic_values["K[1,2]"]
        K[0][2] = intrinsic_values["K[0,2]"]
        K[2][2] = 1
        return K

    @staticmethod
    def process_extrinsic(extrinsic_values):
        extrinsics = np.array(list(extrinsic_values)[2:]).reshape((4, 4))
        return extrinsics

    def cart2hom(self, pts_3d):
        ''' Input: nx3 points in Cartesian
            Oupput: nx4 points in Homogeneous by pending 1
        '''
        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 np.dot(pts_3d_velo, np.transpose(self.V2C))

    def project_ref_to_velo(self, pts_3d_ref):
        pts_3d_ref = self.cart2hom(pts_3d_ref) # nx4
        temp = np.dot(pts_3d_ref, np.transpose(self.C2V))
        # temp[:,0], temp[:, 1] = temp[:, 1], temp[:, 0]
        return temp

    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[:,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.
        '''
        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 rotx(t):
    """ 3D Rotation about the x-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[1, 0, 0],
                     [0, c, -s],
                     [0, s, c]])


def roty(t):
    """ Rotation about the y-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, 0, s],
                     [0, 1, 0],
                     [-s, 0, c]])


def rotz(t):
    """ Rotation about the z-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s, 0],
                     [s, c, 0],
                     [0, 0, 1]])


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 load_depth(depth_filename):
    return np.array(Image.open(depth_filename)).astype(
        np.float64) / 100.0  # convert to meters


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


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))))
    # print(('pts_3d_extend shape: ', pts_3d_extend.shape))
    pts_2d = np.dot(pts_3d_extend, np.transpose(P))  # nx3
    # print(pts_2d)
    pts_2d[:, 0] /= pts_2d[:, 2]
    pts_2d[:, 1] /= pts_2d[:, 2]
    return pts_2d[:, 0:2]


def compute_box_3d(obj, P):
    """ Takes an object and a projection matrix (P) and projects the 3d
        bounding box into the image plane.
        Returns:
            corners_2d: (8,2) array in left image coord.
            corners_3d: (8,3) array in in rect camera coord.
    """
    # compute rotational matrix around yaw axis
    R = roty(obj.ry)

    # 3d bounding box dimensions
    l = obj.l
    w = obj.w
    h = obj.h

    # 3d bounding box corners
    x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]
    y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
    z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]

    # rotate and translate 3d bounding box
    corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    # print corners_3d.shape
    corners_3d[0, :] = corners_3d[0, :] + obj.t[0]
    corners_3d[1, :] = corners_3d[1, :] + obj.t[1]
    corners_3d[2, :] = corners_3d[2, :] + obj.t[2]

    # only draw 3d bounding box for objs in front of the camera
    if np.any(corners_3d[2, :] < 0.1):
        corners_2d = None
        return corners_2d, np.transpose(corners_3d)

    # project the 3d bounding box into the image plane
    corners_2d = project_to_image(np.transpose(corners_3d), P)
    # print 'corners_2d: ', corners_2d
    return corners_2d, np.transpose(corners_3d)


def compute_orientation_3d(obj, P):
    """ Takes an object and a projection matrix (P) and projects the 3d
        object orientation vector into the image plane.
        Returns:
            orientation_2d: (2,2) array in left image coord.
            orientation_3d: (2,3) array in in rect camera coord.
    """

    # compute rotational matrix around yaw axis
    R = roty(obj.ry)

    # orientation in object coordinate system
    orientation_3d = np.array([[0.0, obj.l], [0, 0], [0, 0]])

    # rotate and translate in camera coordinate system, project in image
    orientation_3d = np.dot(R, orientation_3d)
    orientation_3d[0, :] = orientation_3d[0, :] + obj.t[0]
    orientation_3d[1, :] = orientation_3d[1, :] + obj.t[1]
    orientation_3d[2, :] = orientation_3d[2, :] + obj.t[2]

    # vector behind image plane?
    if np.any(orientation_3d[2, :] < 0.1):
        orientation_2d = None
        return orientation_2d, np.transpose(orientation_3d)

    # project orientation into the image plane
    orientation_2d = project_to_image(np.transpose(orientation_3d), P)
    return orientation_2d, np.transpose(orientation_3d)


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
    """
    if qs is not None:
        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 Virtual KITTI data.

Author: <NAME>
Date: February 2020
"""
from __future__ import print_function

import numpy as np
import cv2

from PIL import Image

transform_mat = {"Scene01": {"15-deg-left": -6/12.0, "30-deg-left":-5/12.0,
                             "15-deg-right":-8/12.0, "30-deg-right":-9/12.0,
                             "clone":-7/12.0, "morning":-7/12.0, "rain":-7/12.0,
                             "fog":-7/12.0, "overcast":-7/12.0, "sunset":-7/12.0},
                 "Scene02": {"15-deg-left": 4/12.0, "30-deg-left": 5/12.0,
                             "15-deg-right":2/12.0, "30-deg-right":1/12.0,
                             "clone":3/12.0, "morning": 3/12.0, "rain": 3/12.0,
                             "fog":3/12.0, "overcast": 3/12.0, "sunset": 3/12.0},
                 "Scene06": {"15-deg-left": 0/12.0, "30-deg-left": 1/12.0,
                             "15-deg-right": -2/12.0, "30-deg-right":-3/12.0,
                             "clone":-1/12.0, "morning": -1/12.0, "rain": -1/12.0,
                             "fog":-1/12.0, "overcast": -1/12.0, "sunset": -1/12.0},
                 "Scene20": {"15-deg-left": -9/12.0, "30-deg-left": -8/12.0,
                             "15-deg-right": -11/12.0, "30-deg-right":-12/12.0,
                             "clone":-10/12.0, "morning": -10/12.0, "rain": -10/12.0,
                             "fog":-10/12.0, "overcast": -10/12.0, "sunset": -10/12.0},
                 "Scene18": {"15-deg-left": -1.5/12.0, "30-deg-left": -0.5/12.0,
                             "15-deg-right": -3.5/12.0, "30-deg-right":-4.5/12.0,
                             "clone": -2.5/12.0, "morning": -2.5/12.0, "rain": -2.5/12.0,
                             "fog": -2.5/12.0, "overcast": -2.5/12.0, "sunset": -2.5/12.0}}


class Object3d(object):
    """ 3d object label """

    def __init__(self, data):
        """
        Process one object instance of a frame
        :param data: Series object for one object instance
        """
        # extract label, truncation, occlusion
        self.type = data["label"]  # 'Car', 'Van', 'Truck', etc ...
        self.truncation = data[
            "truncation_ratio"]  # truncated pixel ratio [0..1]
        self.occlusion = 1 - data[
            "occupancy_ratio"]  # 0=visible, 1=partly occluded, 2=fully occluded, 3=unknown
        self.alpha = data["alpha"]  # object observation angle [-pi..pi]

        # extract 2d bounding box in 0-based coordinates
        self.xmin = data["left"]  # left
        self.ymin = data["top"]  # top
        self.xmax = data["right"]  # right
        self.ymax = data["bottom"]  # bottom
        self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])

        # extract 3d bounding box information
        self.h = data["height"]  # box height
        self.w = data["width"]  # box width
        self.l = data["length"]  # box length (in meters)
        self.t = (data["camera_space_X"], data["camera_space_Y"],
                  data["camera_space_Z"])  # location (x,y,z) in camera coord.
        self.ry = data["rotation_camera_space_y"]  # yaw angle
        # (around Y-axis in camera coordinates) [-pi..pi]

    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
        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
    """

    def __init__(self, scene, subscene, intrinsic_values, extrinsic_values):
        # Projection matrix from rect camera coord to image2 coord
        self.P = self.process_intrinsic(intrinsic_values)

        # Rotation from reference camera coord to rect camera coord
        self.extrinsics = self.process_extrinsic(extrinsic_values)
        self.R0 = np.linalg.inv(self.extrinsics[:3, :3])
        # R0= "9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01"
        # self.R0 = np.reshape(np.array([float(x) for x in R0.split()]), [3,3])

        # transform = {"Scene01": -np.pi / 2.0, "Scene02": np.pi/4.0}
        velo = "7.533745000000e-03 -9.999714000000e-01 -6.166020000000e-04 -4.069766000000e-03 1.480249000000e-02 7.280733000000e-04 -9.998902000000e-01 -7.631618000000e-02 9.998621000000e-01 7.523790000000e-03 1.480755000000e-02 -2.717806000000e-01"

        self.V2C = np.array([float(x) for x in velo.split()])
        self.V2C = np.reshape(self.V2C, [3,4])
        self.C2V = inverse_rigid_trans(self.V2C)
        angle = transform_mat[scene][subscene] * np.pi # morning
        self.C2V = np.dot(rotz(angle), inverse_rigid_trans(self.V2C))
        self.V2C = inverse_rigid_trans(self.C2V)


        # 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)

    @staticmethod
    def process_intrinsic(intrinsic_values):
        K = np.zeros((3, 4))
        K[0][0] = intrinsic_values["K[0,0]"]
        K[1][1] = intrinsic_values["K[1,1]"]
        K[1][2] = intrinsic_values["K[1,2]"]
        K[0][2] = intrinsic_values["K[0,2]"]
        K[2][2] = 1
        return K

    @staticmethod
    def process_extrinsic(extrinsic_values):
        extrinsics = np.array(list(extrinsic_values)[2:]).reshape((4, 4))
        return extrinsics

    def cart2hom(self, pts_3d):
        ''' Input: nx3 points in Cartesian
            Oupput: nx4 points in Homogeneous by pending 1
        '''
        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 np.dot(pts_3d_velo, np.transpose(self.V2C))

    def project_ref_to_velo(self, pts_3d_ref):
        pts_3d_ref = self.cart2hom(pts_3d_ref) # nx4
        temp = np.dot(pts_3d_ref, np.transpose(self.C2V))
        # temp[:,0], temp[:, 1] = temp[:, 1], temp[:, 0]
        return temp

    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[:,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.
        '''
        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 rotx(t):
    """ 3D Rotation about the x-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[1, 0, 0],
                     [0, c, -s],
                     [0, s, c]])


def roty(t):
    """ Rotation about the y-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, 0, s],
                     [0, 1, 0],
                     [-s, 0, c]])


def rotz(t):
    """ Rotation about the z-axis. """
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s, 0],
                     [s, c, 0],
                     [0, 0, 1]])


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 load_depth(depth_filename):
    return np.array(Image.open(depth_filename)).astype(
        np.float64) / 100.0  # convert to meters


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


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))))
    # print(('pts_3d_extend shape: ', pts_3d_extend.shape))
    pts_2d = np.dot(pts_3d_extend, np.transpose(P))  # nx3
    # print(pts_2d)
    pts_2d[:, 0] /= pts_2d[:, 2]
    pts_2d[:, 1] /= pts_2d[:, 2]
    return pts_2d[:, 0:2]


def compute_box_3d(obj, P):
    """ Takes an object and a projection matrix (P) and projects the 3d
        bounding box into the image plane.
        Returns:
            corners_2d: (8,2) array in left image coord.
            corners_3d: (8,3) array in in rect camera coord.
    """
    # compute rotational matrix around yaw axis
    R = roty(obj.ry)

    # 3d bounding box dimensions
    l = obj.l
    w = obj.w
    h = obj.h

    # 3d bounding box corners
    x_corners = [l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2]
    y_corners = [0, 0, 0, 0, -h, -h, -h, -h]
    z_corners = [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2]

    # rotate and translate 3d bounding box
    corners_3d = np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    # print corners_3d.shape
    corners_3d[0, :] = corners_3d[0, :] + obj.t[0]
    corners_3d[1, :] = corners_3d[1, :] + obj.t[1]
    corners_3d[2, :] = corners_3d[2, :] + obj.t[2]

    # only draw 3d bounding box for objs in front of the camera
    if np.any(corners_3d[2, :] < 0.1):
        corners_2d = None
        return corners_2d, np.transpose(corners_3d)

    # project the 3d bounding box into the image plane
    corners_2d = project_to_image(np.transpose(corners_3d), P)
    # print 'corners_2d: ', corners_2d
    return corners_2d, np.transpose(corners_3d)


def compute_orientation_3d(obj, P):
    """ Takes an object and a projection matrix (P) and projects the 3d
        object orientation vector into the image plane.
        Returns:
            orientation_2d: (2,2) array in left image coord.
            orientation_3d: (2,3) array in in rect camera coord.
    """

    # compute rotational matrix around yaw axis
    R = roty(obj.ry)

    # orientation in object coordinate system
    orientation_3d = np.array([[0.0, obj.l], [0, 0], [0, 0]])

    # rotate and translate in camera coordinate system, project in image
    orientation_3d = np.dot(R, orientation_3d)
    orientation_3d[0, :] = orientation_3d[0, :] + obj.t[0]
    orientation_3d[1, :] = orientation_3d[1, :] + obj.t[1]
    orientation_3d[2, :] = orientation_3d[2, :] + obj.t[2]

    # vector behind image plane?
    if np.any(orientation_3d[2, :] < 0.1):
        orientation_2d = None
        return orientation_2d, np.transpose(orientation_3d)

    # project orientation into the image plane
    orientation_2d = project_to_image(np.transpose(orientation_3d), P)
    return orientation_2d, np.transpose(orientation_3d)


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
    """
    if qs is not None:
        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 Virtual KITTI data.  Author: <NAME> Date: February 2020 3d object label Process one object instance of a frame         :param data: Series object for one object instance # extract label, truncation, occlusion # 'Car', 'Van', 'Truck', etc ... # 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. # yaw angle # (around Y-axis in camera coordinates) [-pi..pi] Calibration matrices and utils         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 # Projection matrix from rect camera coord to image2 coord # Rotation from reference camera coord to rect camera coord # R0= "9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01" # self.R0 = np.reshape(np.array([float(x) for x in R0.split()]), [3,3]) # transform = {"Scene01": -np.pi / 2.0, "Scene02": np.pi/4.0} # morning # Camera intrinsics and extrinsics # relative Input: nx3 points in Cartesian             Oupput: nx4 points in Homogeneous by pending 1 # =========================== # ------- 3d to 3d ---------- # =========================== # nx4 # nx4 # temp[:,0], temp[:, 1] = temp[:, 1], temp[:, 0] 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 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. 3D Rotation about the x-axis. Rotation about the y-axis. Rotation about the z-axis. Transforation matrix from rotation matrix and translation vector. Inverse a rigid body transform matrix (3x4 as [R|t])         [R'|-R't; 0|1] # 3x4 # convert to meters 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) # print(('pts_3d_extend shape: ', pts_3d_extend.shape)) # nx3 # print(pts_2d) Takes an object and a projection matrix (P) and projects the 3d         bounding box into the image plane.         Returns:             corners_2d: (8,2) array in left image coord.             corners_3d: (8,3) array in in rect camera coord. # compute rotational matrix around yaw axis # 3d bounding box dimensions # 3d bounding box corners # rotate and translate 3d bounding box # print corners_3d.shape # only draw 3d bounding box for objs in front of the camera # project the 3d bounding box into the image plane # print 'corners_2d: ', corners_2d Takes an object and a projection matrix (P) and projects the 3d         object orientation vector into the image plane.         Returns:             orientation_2d: (2,2) array in left image coord.             orientation_3d: (2,3) array in in rect camera coord. # compute rotational matrix around yaw axis # orientation in object coordinate system # rotate and translate in camera coordinate system, project in image # vector behind image plane? # project orientation into the image plane 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