utils.py

"""
Mask R-CNN
Common utility functions and classes.
Copyright (c) 2017 Matterport, Inc.
Licensed under the MIT License (see LICENSE for details)
Written by Waleed Abdulla
"""

import sys
import os
import open3d as o3d
import cv2
import math
import random
import numpy as np
import tensorflow as tf
import scipy.misc
import skimage.color
import _pickle as cPickle
from ctypes import *
import copy
import ICP
import glob
import time
import PIL
from PIL import Image
from pdb import set_trace as bp

from aligning import estimateSimilarityTransform

sys.path.append('./cocoapi/PythonAPI')
from pycocotools.cocoeval import COCOeval
from pycocotools import mask as maskUtils
import matplotlib
import matplotlib.pyplot as plt



############################################################
#  Bounding Boxes
############################################################

def extract_bboxes(mask):
    """Compute bounding boxes from masks.
    mask: [height, width, num_instances]. Mask pixels are either 1 or 0.
    Returns: bbox array [num_instances, (y1, x1, y2, x2)].
    """
    boxes = np.zeros([mask.shape[-1], 4], dtype=np.int32)
    for i in range(mask.shape[-1]):
        m = mask[:, :, i]
        # Bounding box.
        horizontal_indicies = np.where(np.any(m, axis=0))[0]
        vertical_indicies = np.where(np.any(m, axis=1))[0]
        if horizontal_indicies.shape[0]:
            x1, x2 = horizontal_indicies[[0, -1]]
            y1, y2 = vertical_indicies[[0, -1]]
            # x2 and y2 should not be part of the box. Increment by 1.
            x2 += 1
            y2 += 1
        else:
            # No mask for this instance. Might happen due to
            # resizing or cropping. Set bbox to zeros
            x1, x2, y1, y2 = 0, 0, 0, 0
        boxes[i] = np.array([y1, x1, y2, x2])
    return boxes.astype(np.int32)


def compute_iou(box, boxes, box_area, boxes_area):
    """Calculates IoU of the given box with the array of the given boxes.
    box: 1D vector [y1, x1, y2, x2]
    boxes: [boxes_count, (y1, x1, y2, x2)]
    box_area: float. the area of 'box'
    boxes_area: array of length boxes_count.
    Note: the areas are passed in rather than calculated here for
          efficiency. Calculate once in the caller to avoid duplicate work.
    """
    # Calculate intersection areas
    y1 = np.maximum(box[0], boxes[:, 0])
    y2 = np.minimum(box[2], boxes[:, 2])
    x1 = np.maximum(box[1], boxes[:, 1])
    x2 = np.minimum(box[3], boxes[:, 3])
    intersection = np.maximum(x2 - x1, 0) * np.maximum(y2 - y1, 0)
    union = box_area + boxes_area[:] - intersection[:]
    iou = intersection / union
    return iou


def compute_overlaps(boxes1, boxes2):
    """Computes IoU overlaps between two sets of boxes.
    boxes1, boxes2: [N, (y1, x1, y2, x2)].
    For better performance, pass the largest set first and the smaller second.
    """
    # Areas of anchors and GT boxes
    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])

    # Compute overlaps to generate matrix [boxes1 count, boxes2 count]
    # Each cell contains the IoU value.
    overlaps = np.zeros((boxes1.shape[0], boxes2.shape[0]))
    for i in range(overlaps.shape[1]):
        box2 = boxes2[i]
        overlaps[:, i] = compute_iou(box2, boxes1, area2[i], area1)
    return overlaps



def compute_overlaps_masks(masks1, masks2):
    '''Computes IoU overlaps between two sets of masks.
    masks1, masks2: [Height, Width, instances]
    '''
    # flatten masks
    masks1 = np.reshape(masks1 > .5, (-1, masks1.shape[-1])).astype(np.float32)
    masks2 = np.reshape(masks2 > .5, (-1, masks2.shape[-1])).astype(np.float32)
    area1 = np.sum(masks1, axis=0)
    area2 = np.sum(masks2, axis=0)

    # intersections and union
    intersections = np.dot(masks1.T, masks2)
    union = area1[:, None] + area2[None, :] - intersections
    overlaps = intersections / union

    return overlaps


def compute_mean_l1_coord_diff(mask1, mask2, coord1, coord2, synset, cls_id):
    '''Computes IoU overlaps between two sets of masks.
    mask1, mask2: [Height, Width]
    coord1, coord2: [Height, Width, 3]
    
    '''
    # flatten masks
    num_pixels = mask1.shape[0] * mask1.shape[1]
    
    mask1 = np.reshape(mask1 > .5, (-1)).astype(np.float32)
    mask2 = np.reshape(mask2 > .5, (-1)).astype(np.float32)
    coord1 = np.reshape(coord1, (-1, 3)).astype(np.float32)
    coord2 = np.reshape(coord2, (-1, 3)).astype(np.float32)

    # intersections and union
    intersections = np.logical_and(mask1, mask2)
    num_pixel_intersection = len(np.where(intersections)[0])
    
    pts1 = coord1[intersections, :].transpose() - 0.5
    pts2 = coord2[intersections, :].transpose() - 0.5
    
    def rotation_y_matrix(theta):
        rotation_matrix =  \
                np.array([ np.cos(theta), 0,  np.sin(theta),
                             0,           1,  0,
                          -np.sin(theta), 0,  np.cos(theta)])
        rotation_matrix = np.reshape(rotation_matrix, (3, 3))
        return rotation_matrix
      
    #if synset[cls_id] in ['bottle', 'bowl', 'can']:
    if synset[cls_id] in ['stem', 'non-stem']:
        M = 20
        pts1_symmetry = np.zeros(pts1.shape+(M,))  ## shape: (3, N, 6)
        for i in range(M):
            rotated_pts1 = rotation_y_matrix(float(i)*np.float32(2*math.pi/M)) @ pts1
            pts1_symmetry[:, :, i] = rotated_pts1
        
        pts2_reshape = pts2.reshape([3, -1, 1])
        mean_dists = np.mean(np.linalg.norm(pts1_symmetry - pts2_reshape, axis=0), axis=0)
        mean_dist = np.amin(mean_dists)
    #elif synset[cls_id] in ['phone']:
    elif synset[cls_id] in ['box']:
        pts1_symmetry = np.zeros(pts1.shape+(2,))
        
        for i in range(2):
            rotated_pts1 = rotation_y_matrix(float(i)*np.float32(2*math.pi/2)) @ pts1
            #print(rotated_pts1)
            pts1_symmetry[:, :, i] = rotated_pts1
        
        pts2_reshape = pts2.reshape([3, -1, 1])
        mean_dists = np.mean(np.linalg.norm(pts1_symmetry - pts2_reshape, axis=0), axis=0)
        mean_dist = np.amin(mean_dists)
    else:
        #print(synset[cls_id])
        diff = pts1 - pts2
        dist = np.linalg.norm(diff, axis=0)
        assert dist.shape[0] == num_pixel_intersection
        mean_dist = np.mean(dist)
    
    mean_l1_coord_diff = mean_dist
    #print(mean_l1_coord_diff, pts1.shape[0])
    
    return mean_l1_coord_diff




def compute_3d_iou(bbox_3d_1, bbox_3d_2, handle_visibility, class_name_1, class_name_2):
    '''Computes IoU overlaps between two 3d bboxes.
       bbox_3d_1, bbox_3d_1: [3, 8]
    '''
    # flatten masks
    def asymmetric_3d_iou(bbox_3d_1, bbox_3d_2):
        bbox_1_max = np.amax(bbox_3d_1, axis=0)
        bbox_1_min = np.amin(bbox_3d_1, axis=0)
        bbox_2_max = np.amax(bbox_3d_2, axis=0)
        bbox_2_min = np.amin(bbox_3d_2, axis=0)

        overlap_min = np.maximum(bbox_1_min, bbox_2_min)
        overlap_max = np.minimum(bbox_1_max, bbox_2_max)

        # intersections and union
        if np.amin(overlap_max - overlap_min) <0:
        	intersections = 0
        else:
        	intersections = np.prod(overlap_max - overlap_min)
        union = np.prod(bbox_1_max - bbox_1_min) + np.prod(bbox_2_max - bbox_2_min) - intersections
        overlaps = intersections / union
        return overlaps


    if bbox_3d_1 is None or bbox_3d_2 is None:
        return -1

    symmetry_flag = False
    #if class_name_1 in ['bottle', 'bowl', 'can'] and class_name_1 == class_name_2:
    if class_name_1 in ['stem', 'non-stem'] and class_name_1 == class_name_2:
        symmetry_flag = True
    if class_name_1 == 'mug' and class_name_1 == class_name_2 and handle_visibility==0:
        symmetry_flag = True

    if symmetry_flag:
        print('*'*10)
        n = 20
        theta = 2*math.pi/n

        # y_rotation_matrix = np.array([[np.cos(theta), 0, np.sin(theta)],
        #                               [0, 1, 0], 
        #                               [-np.sin(theta), 0, np.cos(theta)]])
        
        z_rotation_matrix = np.array( [ [np.cos(theta), -np.sin(theta), 0],
							            [np.sin(theta),  np.cos(theta), 0],
							            [0,              0,             1]])

        max_iou = 0
        for i in range(n):
            bbox_center = np.mean(bbox_3d_1, -1, keepdims=True)
            bbox_3d_1 = z_rotation_matrix @ (bbox_3d_1 - bbox_center) + bbox_center # we want to rotate about the centre of the object.
            max_iou = max(max_iou, asymmetric_3d_iou(bbox_3d_1, bbox_3d_2))
        return max_iou
    else:
        return asymmetric_3d_iou(bbox_3d_1, bbox_3d_2)



def compute_3d_iou_new(RT_1, RT_2, dimensions_1, dimensions_2, class_name_1, class_name_2):
    '''Computes IoU overlaps between two 3d bboxes.
       bbox_3d_1, bbox_3d_1: [3, 8]
    '''
    # flatten masks
    def asymmetric_3d_iou(RT_1, RT_2, dimensions_1, dimensions_2):
        cube_1 = get_3d_bbox(dimensions_1, 0)
        bbox_3d_1 = transform_coordinates_3d(cube_1, RT_1)

        cube_2 = get_3d_bbox(dimensions_2, 0)
        bbox_3d_2 = transform_coordinates_3d(cube_2, RT_2)

        bbox_1_max = np.amax(bbox_3d_1, axis=0)
        bbox_1_min = np.amin(bbox_3d_1, axis=0)
        bbox_2_max = np.amax(bbox_3d_2, axis=0)
        bbox_2_min = np.amin(bbox_3d_2, axis=0)

        overlap_min = np.maximum(bbox_1_min, bbox_2_min)
        overlap_max = np.minimum(bbox_1_max, bbox_2_max)

        # intersections and union
        if np.amin(overlap_max - overlap_min) < 0:
            intersections = 0
        else:
            intersections = np.prod(overlap_max - overlap_min)
        union = np.prod(bbox_1_max - bbox_1_min) + np.prod(bbox_2_max - bbox_2_min) - intersections
        overlaps = intersections / union
        return overlaps


    if RT_1 is None or RT_2 is None:
        return -1

    symmetry_flag = False
    #if (class_name_1 in ['bottle', 'bowl', 'can'] and class_name_1 == class_name_2) or (class_name_1 == 'mug' and class_name_1 == class_name_2 and handle_visibility==0):
    if (class_name_1 in ['stem', 'non-stem'] and class_name_1 == class_name_2):
        #print('*'*10)
    
        # cube_1 = get_3d_bbox(dimensions_1, 0)
        # cube_2 = get_3d_bbox(dimensions_2, 0)
        # bbox_3d_2 = transform_coordinates_3d(cube_2, RT_2)
        
        def z_rotation_matrix(theta):
            return np.array( [  [np.cos(theta), -np.sin(theta), 0, 0],
                                [np.sin(theta),  np.cos(theta), 0, 0],
                                [0,              0,             1, 0],
                                [0,              0,             0, 1]])
        def y_rotation_matrix(theta):
            return np.array([[np.cos(theta), 0, np.sin(theta), 0],
                             [0, 1, 0 , 0], 
                             [-np.sin(theta), 0, np.cos(theta), 0],
                             [0, 0, 0 , 1]])

        n = 20
        max_iou = 0
        for i in range(n):
            # # before
            # bbox_center = np.mean(bbox_3d_1, -1, keepdims=True)
            # bbox_3d_1 = y_rotation_matrix @ (bbox_3d_1 - bbox_center) + bbox_center
            # max_iou = max(max_iou, asymmetric_3d_iou(bbox_3d_1, bbox_3d_2))
            theta = 2*math.pi*i/float(n)
            # First rotate object around Z-axis for theta degrees, then apply the original RT            
            rotated_RT_1 = RT_1 @ z_rotation_matrix(theta)
            max_iou = max(max_iou, asymmetric_3d_iou(rotated_RT_1, RT_2, dimensions_1, dimensions_2))
    else:
        max_iou = asymmetric_3d_iou(RT_1, RT_2, dimensions_1, dimensions_2)
    
    # Return IoU percentage (0.01 = 1%)
    return max_iou
        



def compute_RT_distances(RT_1, RT_2):
    '''
    :param RT_1: [4, 4]. homogeneous affine transformation
    :param RT_2: [4, 4]. homogeneous affine transformation
    :return: theta: angle difference of R in degree, shift: l2 difference of T in centimeter
    '''
    ## Xavier: INFO
    # np.linalg.det computes the determinant
    # np.linalg.norm computes one of eight matrix norms (TODO: what does this mean?)

    #print(RT_1[3, :], RT_2[3, :])
    
    if RT_1 is None or RT_2 is None:
        return -1

    ## make sure the last row is [0, 0, 0, 1]
    try:
        assert np.array_equal(RT_1[3, :], RT_2[3, :])
        assert np.array_equal(RT_1[3, :], np.array([0, 0, 0, 1]))
    except AssertionError:
        print(RT_1[3, :], RT_2[3, :])

    R1 = RT_1[:3, :3]/np.cbrt(np.linalg.det(RT_1[:3, :3]))
    T1 = RT_1[:3, 3]

    R2 = RT_2[:3, :3]/np.cbrt(np.linalg.det(RT_2[:3, :3]))
    T2 = RT_2[:3, 3]

    R = R1 @ R2.transpose()
    theta = np.arccos((np.trace(R) - 1)/2) * 180/np.pi
    shift = np.linalg.norm(T1-T2) * 100
    # print(theta, shift)

    if theta < 5 and shift < 5:
        return 10 - theta - shift
    else:
        return -1


def compute_RT_degree_cm_symmetry(RT_1, RT_2, class_id, synset_names):
    '''
    :param RT_1: [4, 4]. homogeneous affine transformation
    :param RT_2: [4, 4]. homogeneous affine transformation
    :return: theta: angle difference of R in degree, shift: l2 difference of T in centimeter

    # TODO: make sure my objects are y-up...
    
    synset_names = ['BG',  # 0
                    'bottle',  # 1
                    'bowl',  # 2
                    'camera',  # 3
                    'can',  # 4
                    'cap',  # 5
                    'phone',  # 6
                    'monitor',  # 7
                    'laptop',  # 8
                    'mug'  # 9
                    ]
    
    synset_names = ['BG',  # 0
                    'bottle',  # 1
                    'bowl',  # 2
                    'camera',  # 3
                    'can',  # 4
                    'laptop',  # 5
                    'mug'  # 6
                    ]
    '''

    ## make sure the last row is [0, 0, 0, 1]
    if RT_1 is None or RT_2 is None:
        return -1
    try:
        assert np.array_equal(RT_1[3, :], RT_2[3, :])
        assert np.array_equal(RT_1[3, :], np.array([0, 0, 0, 1]))
    except AssertionError:
        print(RT_1[3, :], RT_2[3, :])
        exit()

    R1 = RT_1[:3, :3] / np.cbrt(np.linalg.det(RT_1[:3, :3]))
    T1 = RT_1[:3, 3]

    R2 = RT_2[:3, :3] / np.cbrt(np.linalg.det(RT_2[:3, :3]))
    T2 = RT_2[:3, 3]

#     try:
#         assert np.abs(np.linalg.det(R1) - 1) < 0.01
#         assert np.abs(np.linalg.det(R2) - 1) < 0.01
#     except AssertionError:
#         print(np.linalg.det(R1), np.linalg.det(R2))

    if synset_names[class_id] in ['stem', 'non-stem']: 

        ### Compute theta for an object that is symmetric when rotating around y-axis 
        
        #y = np.array([0, 1, 0])
        z = np.array([0, 0, 1])
        y1 = R1 @ z
        y2 = R2 @ z
        theta = np.arccos(y1.dot(y2) / (np.linalg.norm(y1) * np.linalg.norm(y2)))
    
    elif synset_names[class_id] in ['box']:
        
        ### Compute theta for an object that is symmetric when rotated 180 degrees around y-axis 
        
        #y_180_RT = np.diag([-1.0, 1.0, -1.0])
        z_180_RT = np.diag([-1.0, -1.0, 1.0])
        
        # X: Step 1a. Compute the difference in rotation between these two matrices
        R = R1 @ R2.transpose() 
        # X: Step 1b. Compute the difference in rotation between these two matrices, but rotated 180 degrees around y-axis
        R_rot = R1 @ z_180_RT @ R2.transpose() 
        
        # X: Step 2. Compute the axis-angle (ω, θ) representation of R and R_rot using the following formula
        # We take the minimum rotational error, because we want the loss of a box to be equal to the same box but rotation 180 degrees around y-axis.
        theta = min(np.arccos((np.trace(R) - 1) / 2),
                    np.arccos((np.trace(R_rot) - 1) / 2))
    else:
        
        ### Compute theta for an object that has no symmetry
        
        R = R1 @ R2.transpose()
        theta = np.arccos((np.trace(R) - 1) / 2)

    theta *= 180 / np.pi # Radian to degrees
    shift = np.linalg.norm(T1 - T2) / 10 # /10 because millimeter to centimeter
    result = np.array([theta, shift])

    return result


def compute_RT_projection_2d_symmetry(RT_1, RT_2, class_id, handle_visibility, mesh_vertices, intrinsics, synset_names, num_rotation=20):
    '''
    :param RT_1: [4, 4]. homogeneous affine transformation
    :param RT_2: [4, 4]. homogeneous affine transformation
    :param vertices: [3, N].
    :param intrinsics: [4, 4]
    :return: mean 2d projection distance in pixel

    synset_names = ['BG',  # 0
                    'bottle',  # 1
                    'bowl',  # 2
                    'camera',  # 3
                    'can',  # 4
                    'laptop',  # 5
                    'mug'  # 6
                    ]
    '''

    ## make sure the last row is [0, 0, 0, 1]
    if RT_1 is None or RT_2 is None:
        return -1
    try:
        assert np.array_equal(RT_1[3, :], RT_2[3, :])
        assert np.array_equal(RT_1[3, :], np.array([0, 0, 0, 1]))
    except AssertionError:
        print(RT_1[3, :], RT_2[3, :])
        exit()

    RT_1[:3, :3] = RT_1[:3, :3]/np.cbrt(np.linalg.det(RT_1[:3, :3]))
    R1 = RT_1[:3, :3]
    #T1 = RT_1[:3, 3]

    RT_2[:3, :3] = RT_2[:3, :3]/np.cbrt(np.linalg.det(RT_2[:3, :3]))
    R2 = RT_2[:3, :3]
    #T2 = RT_2[:3, 3]
    try:
        assert np.abs(np.linalg.det(R1) - 1) < 0.01
        assert np.abs(np.linalg.det(R2) - 1) < 0.01
    except AssertionError:
        print(np.linalg.det(R1), np.linalg.det(R2))
    
    # check the vertices are in meter unit
    vertices = np.copy(mesh_vertices)/1000
    assert np.amax(vertices) < 0.5, np.amax(vertices) 
    assert np.amax(vertices) > 0, np.amax(vertices)
    assert np.amin(vertices) < 0, np.amin(vertices)
    assert np.amin(vertices) > -0.5, np.amin(vertices)
 
    assert vertices.shape[0] == 3
    num_vertices = vertices.shape[1]
    
    coords_3d_1 = transform_coordinates_3d(vertices, RT_1)
    projected_1 = calculate_2d_projections(coords_3d_1, intrinsics)
    coords_3d_2 = transform_coordinates_3d(vertices, RT_2)
    projected_2 = calculate_2d_projections(coords_3d_2, intrinsics)
    
    # calculate reprojection 2d error
    dists = np.linalg.norm(projected_1 - projected_2, axis=1)
    assert len(dists) == num_vertices
    min_mean_dist = np.mean(dists)


    ## take care of symmetry categories

    # freely rotate around y axis

    if (synset_names[class_id] in ['stem', 'non-stem']):  
        
        def y_rotation_matrix(theta):
            return np.array([[np.cos(theta), 0, np.sin(theta)],
                             [0, 1, 0], 
                             [-np.sin(theta), 0, np.cos(theta)]])
        for i in range(1, num_rotation):
            theta = 2*math.pi*i/float(num_rotation)
            coords_3d_2 = transform_coordinates_3d(y_rotation_matrix(theta)@vertices, RT_2)
            projected_2 = calculate_2d_projections(coords_3d_2, intrinsics)

            dists = np.linalg.norm(projected_1 - projected_2, axis=1) #axis=1
            assert len(dists) == num_vertices
            min_mean_dist = min(min_mean_dist, np.mean(dists))

    # rotate 180 around y axis
    elif synset_names[class_id] in ['box']:
        y_180_RT = np.diag([-1.0, 1.0, -1.0])
        
        coords_3d_2 = transform_coordinates_3d(y_180_RT@vertices, RT_2)
        projected_2 = calculate_2d_projections(coords_3d_2, intrinsics)
        dists = np.linalg.norm(projected_1 - projected_2, axis=1) #axis=1
        assert len(dists) == num_vertices
        min_mean_dist = min(min_mean_dist, np.mean(dists))
        
    # rotate 180 around z axis
    elif synset_names[class_id] in ['eggbox', 'glue']:
        z_180_RT = np.diag([-1.0, -1.0, 1.0])
        coords_3d_2 = transform_coordinates_3d(z_180_RT@vertices, RT_2)
        projected_2 = calculate_2d_projections(coords_3d_2, intrinsics)
        dists = np.linalg.norm(projected_1 - projected_2, axis=1)
        assert len(dists) == num_vertices
        min_mean_dist = min(min_mean_dist, np.mean(dists))

    else: ## normal asymmetric objects        
        min_mean_dist = min_mean_dist

    
    return min_mean_dist


def non_max_suppression(boxes, scores, threshold):
    """Performs non-maximum suppression and returns indices of kept boxes.
    boxes: [N, (y1, x1, y2, x2)]. Notice that (y2, x2) lays outside the box.
    scores: 1-D array of box scores.
    threshold: Float. IoU threshold to use for filtering.
    """
    assert boxes.shape[0] > 0
    if boxes.dtype.kind != "f":
        boxes = boxes.astype(np.float32)

    # Compute box areas
    y1 = boxes[:, 0]
    x1 = boxes[:, 1]
    y2 = boxes[:, 2]
    x2 = boxes[:, 3]
    area = (y2 - y1) * (x2 - x1)

    # Get indicies of boxes sorted by scores (highest first)
    ixs = scores.argsort()[::-1]

    pick = []
    while len(ixs) > 0:
        # Pick top box and add its index to the list
        i = ixs[0]
        pick.append(i)
        # Compute IoU of the picked box with the rest
        iou = compute_iou(boxes[i], boxes[ixs[1:]], area[i], area[ixs[1:]])
        # Identify boxes with IoU over the threshold. This
        # returns indicies into ixs[1:], so add 1 to get
        # indicies into ixs.
        remove_ixs = np.where(iou > threshold)[0] + 1
        # Remove indicies of the picked and overlapped boxes.
        ixs = np.delete(ixs, remove_ixs)
        ixs = np.delete(ixs, 0)
    return np.array(pick, dtype=np.int32)


def apply_box_deltas(boxes, deltas):
    """Applies the given deltas to the given boxes.
    boxes: [N, (y1, x1, y2, x2)]. Note that (y2, x2) is outside the box.
    deltas: [N, (dy, dx, log(dh), log(dw))]
    """
    boxes = boxes.astype(np.float32)
    # Convert to y, x, h, w
    height = boxes[:, 2] - boxes[:, 0]
    width = boxes[:, 3] - boxes[:, 1]
    center_y = boxes[:, 0] + 0.5 * height
    center_x = boxes[:, 1] + 0.5 * width
    # Apply deltas
    center_y += deltas[:, 0] * height
    center_x += deltas[:, 1] * width
    height *= np.exp(deltas[:, 2])
    width *= np.exp(deltas[:, 3])
    # Convert back to y1, x1, y2, x2
    y1 = center_y - 0.5 * height
    x1 = center_x - 0.5 * width
    y2 = y1 + height
    x2 = x1 + width
    return np.stack([y1, x1, y2, x2], axis=1)


def box_refinement_graph(box, gt_box):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (y1, x1, y2, x2)]
    """
    box = tf.cast(box, tf.float32)
    gt_box = tf.cast(gt_box, tf.float32)

    height = box[:, 2] - box[:, 0]
    width = box[:, 3] - box[:, 1]
    center_y = box[:, 0] + 0.5 * height
    center_x = box[:, 1] + 0.5 * width

    gt_height = gt_box[:, 2] - gt_box[:, 0]
    gt_width = gt_box[:, 3] - gt_box[:, 1]
    gt_center_y = gt_box[:, 0] + 0.5 * gt_height
    gt_center_x = gt_box[:, 1] + 0.5 * gt_width

    dy = (gt_center_y - center_y) / height
    dx = (gt_center_x - center_x) / width
    dh = tf.log(gt_height / height)
    dw = tf.log(gt_width / width)

    result = tf.stack([dy, dx, dh, dw], axis=1)
    return result


def box_refinement(box, gt_box):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (y1, x1, y2, x2)]. (y2, x2) is
    assumed to be outside the box.
    """
    box = box.astype(np.float32)
    gt_box = gt_box.astype(np.float32)

    height = box[:, 2] - box[:, 0]
    width = box[:, 3] - box[:, 1]
    center_y = box[:, 0] + 0.5 * height
    center_x = box[:, 1] + 0.5 * width

    gt_height = gt_box[:, 2] - gt_box[:, 0]
    gt_width = gt_box[:, 3] - gt_box[:, 1]
    gt_center_y = gt_box[:, 0] + 0.5 * gt_height
    gt_center_x = gt_box[:, 1] + 0.5 * gt_width

    dy = (gt_center_y - center_y) / height
    dx = (gt_center_x - center_x) / width
    dh = np.log(gt_height / height)
    dw = np.log(gt_width / width)

    return np.stack([dy, dx, dh, dw], axis=1)


def get_3d_bbox(scale, shift = 0):
    """
    Input: 
        scale: [3] or scalar
        shift: [3] or scalar
    Return 
        bbox_3d: [3, N]

    """
    if hasattr(scale, "__iter__"):
        bbox_3d = np.array([
                  [ scale[0] / 2, +scale[1] / 2,  scale[2] / 2],
                  [ scale[0] / 2, +scale[1] / 2, -scale[2] / 2],
                  [-scale[0] / 2, +scale[1] / 2,  scale[2] / 2],
                  [-scale[0] / 2, +scale[1] / 2, -scale[2] / 2],
                  [+scale[0] / 2, -scale[1] / 2,  scale[2] / 2],
                  [+scale[0] / 2, -scale[1] / 2, -scale[2] / 2],
                  [-scale[0] / 2, -scale[1] / 2,  scale[2] / 2],
                  [-scale[0] / 2, -scale[1] / 2, -scale[2] / 2]]) + shift
    else:
        bbox_3d = np.array(
                 [[ scale / 2,  +scale / 2,  scale / 2], #(x,y,z)
                  [ scale / 2,  +scale / 2, -scale / 2], #(x,y,-z)
                  [-scale / 2,  +scale / 2,  scale / 2], #(-x,y,z)
                  [-scale / 2,  +scale / 2, -scale / 2], #(-x,y,-z)
                  [+scale / 2,  -scale / 2,  scale / 2], #(x,-y,z)
                  [+scale / 2,  -scale / 2, -scale / 2], #(x,-y,-z)
                  [-scale / 2,  -scale / 2,  scale / 2], #(-x,-y,z)
                  [-scale / 2,  -scale / 2, -scale / 2]]) +shift #(-x,-y,-z)

    bbox_3d = bbox_3d.transpose()
    return bbox_3d



def transform_coordinates_3d(coordinates, RT):
    """
    Input: 
        coordinates: [3, N]
        RT: [4, 4]
    Return 
        new_coordinates: [3, N]

    """
    assert coordinates.shape[0] == 3
    coordinates = np.vstack([coordinates, np.ones((1, coordinates.shape[1]), dtype=np.float32)])
    new_coordinates = RT @ coordinates
    new_coordinates = new_coordinates[:3, :]/new_coordinates[3, :]
    return new_coordinates


def calculate_2d_projections(coordinates_3d, intrinsics):
    """
    Input: 
        coordinates: [3, N]
        intrinsics: [3, 3]
    Return 
        projected_coordinates: [N, 2]
    """
    projected_coordinates = intrinsics @ coordinates_3d
    projected_coordinates = projected_coordinates[:2, :] / projected_coordinates[2, :]
    projected_coordinates = projected_coordinates.transpose()
    projected_coordinates = np.array(projected_coordinates, dtype=np.int32)

    return projected_coordinates

 

############################################################
#  IMAGE AUGMENTATION
############################################################

def calculate_rotation(image_size, angle):
    image_center = tuple(np.array(image_size) / 2)

    # Convert the OpenCV 3x2 rotation matrix to 3x3
    rot_mat = np.vstack(
        [cv2.getRotationMatrix2D(image_center, angle, 1.0), [0, 0, 1]]
    )

    rot_mat_notranslate = np.matrix(rot_mat[0:2, 0:2])

    # Shorthand for below calcs
    image_w2 = image_size[0] * 0.5
    image_h2 = image_size[1] * 0.5

    # Obtain the rotated coordinates of the image corners
    rotated_coords = [
        (np.array([-image_w2, image_h2]) * rot_mat_notranslate).A[0],
        (np.array([image_w2, image_h2]) * rot_mat_notranslate).A[0],
        (np.array([-image_w2, -image_h2]) * rot_mat_notranslate).A[0],
        (np.array([image_w2, -image_h2]) * rot_mat_notranslate).A[0]
    ]

    # Find the size of the new image
    x_coords = [pt[0] for pt in rotated_coords]
    x_pos = [x for x in x_coords if x > 0]
    x_neg = [x for x in x_coords if x < 0]

    y_coords = [pt[1] for pt in rotated_coords]
    y_pos = [y for y in y_coords if y > 0]
    y_neg = [y for y in y_coords if y < 0]

    right_bound = max(x_pos)
    left_bound = min(x_neg)
    top_bound = max(y_pos)
    bot_bound = min(y_neg)

    new_w = int(abs(right_bound - left_bound))
    new_h = int(abs(top_bound - bot_bound))

    # We require a translation matrix to keep the image centred
    trans_mat = np.matrix([
        [1, 0, int(new_w * 0.5 - image_w2)],
        [0, 1, int(new_h * 0.5 - image_h2)],
        [0, 0, 1]
    ])

    # Compute the tranform for the combined rotation and translation
    affine_mat = (np.matrix(trans_mat) * np.matrix(rot_mat))[0:2, :]

    return new_w, new_h, affine_mat



def rotate_image(image, new_w, new_h, affine_mat, interpolation=cv2.INTER_LINEAR):
    """
    Rotates an OpenCV 2 / NumPy image about it's centre by the given angle
    (in degrees). The returned image will be large enough to hold the entire
    new image, with a black background
    """

    # Get the image size
    # No that's not an error - NumPy stores image matricies backwards

    # Apply the transform
    result = cv2.warpAffine(
        image,
        affine_mat,
        (new_w, new_h),
        flags=interpolation
    )

    return result


def largest_rotated_rect(w, h, angle):
    """
    Given a rectangle of size wxh that has been rotated by 'angle' (in
    radians), computes the width and height of the largest possible
    axis-aligned rectangle within the rotated rectangle.

    Original JS code by 'Andri' and Magnus Hoff from Stack Overflow

    Converted to Python by Aaron Snoswell
    """

    quadrant = int(math.floor(angle / (math.pi / 2))) & 3
    sign_alpha = angle if ((quadrant & 1) == 0) else math.pi - angle
    alpha = (sign_alpha % math.pi + math.pi) % math.pi

    bb_w = w * math.cos(alpha) + h * math.sin(alpha)
    bb_h = w * math.sin(alpha) + h * math.cos(alpha)

    gamma = math.atan2(bb_w, bb_w) if (w < h) else math.atan2(bb_w, bb_w)

    delta = math.pi - alpha - gamma

    length = h if (w < h) else w

    d = length * math.cos(alpha)
    a = d * math.sin(alpha) / math.sin(delta)

    y = a * math.cos(gamma)
    x = y * math.tan(gamma)

    return (
        bb_w - 2 * x,
        bb_h - 2 * y
    )


def crop_around_center(image, width, height):
    """
    Given a NumPy / OpenCV 2 image, crops it to the given width and height,
    around it's centre point
    """

    image_size = (image.shape[1], image.shape[0])
    image_center = (int(image_size[0] * 0.5), int(image_size[1] * 0.5))

    if(width > image_size[0]):
        width = image_size[0]

    if(height > image_size[1]):
        height = image_size[1]

    x1 = int(image_center[0] - width * 0.5)
    x2 = int(image_center[0] + width * 0.5)
    y1 = int(image_center[1] - height * 0.5)
    y2 = int(image_center[1] + height * 0.5)

    return image[y1:y2, x1:x2]

def rotate_and_crop(image, rotate_degree, interpolation):
    image_height, image_width = image.shape[0:2]


    new_w, new_h, affine_mat = calculate_rotation(image.shape[0:2][::-1], rotate_degree)
    image_rotated = rotate_image(image, new_w, new_h, affine_mat, interpolation)
    image_rotated_cropped = crop_around_center(
        image_rotated,
        *largest_rotated_rect(
            image_width,
            image_height,
            math.radians(rotate_degree)
        )
    )

    return image_rotated_cropped


def rotate_and_crop_images(image, masks, coords, rotate_degree):

    image_height, image_width = image.shape[0:2]
    new_w, new_h, affine_mat = calculate_rotation(image.shape[0:2][::-1], rotate_degree)

    image_rotated = rotate_image(image, new_w, new_h, affine_mat, cv2.INTER_LINEAR)
    mask_rotated = rotate_image(masks, new_w, new_h, affine_mat, cv2.INTER_NEAREST)
    
    rect = largest_rotated_rect(
            image_width,
            image_height,
            math.radians(rotate_degree)
        )

    image_rotated_cropped = crop_around_center(image_rotated, *rect)
    mask_rotated_cropped = crop_around_center(mask_rotated, *rect)

    image_rotated_cropped = cv2.resize(image_rotated_cropped, (image_width, image_height),interpolation=cv2.INTER_LINEAR)
    mask_rotated_cropped = cv2.resize(mask_rotated_cropped, (image_width, image_height), interpolation=cv2.INTER_NEAREST)

    if coords is not None:
        coord_rotated = rotate_image(coords, new_w, new_h, affine_mat, cv2.INTER_NEAREST)
        coord_rotated_cropped = crop_around_center(coord_rotated, *rect)
        coord_rotated_cropped = cv2.resize(coord_rotated_cropped, (image_width, image_height), interpolation=cv2.INTER_NEAREST)

        return image_rotated_cropped, mask_rotated_cropped, coord_rotated_cropped
    else:
        return image_rotated_cropped, mask_rotated_cropped


############################################################
#  Dataset
############################################################

class Dataset(object):
    """The base class for dataset classes.
    To use it, create a new class that adds functions specific to the dataset
    you want to use. For example:
    class CatsAndDogsDataset(Dataset):
        def load_cats_and_dogs(self):
            ...
        def load_mask(self, image_id):
            ...
        def image_reference(self, image_id):
            ...
    See COCODataset and ShapesDataset as examples.
    """
    def __init__(self, class_map=None):
        self._image_ids = []
        self.image_info = []
        # Background is always the first class
        self.class_info = [{"source": "", "id": 0, "name": "BG"}]
        self.source_class_ids = {}

    def add_class(self, source, class_id, class_name):
        assert "." not in source, "Source name cannot contain a dot"
        # Does the class exist already?
        for info in self.class_info:
            if info['source'] == source and info["id"] == class_id:
                # source.class_id combination already available, skip
                return
        # Add the class
        self.class_info.append({
            "source": source,
            "id": class_id,
            "name": class_name,
        })

    def add_image(self, source, image_id, rgbpath, depthpath, nocspath, maskpath,  **kwargs):
        image_info = {
            "id": image_id,
            "source": source,
            "rgb_path": rgbpath,
            "depth_path": depthpath,
            "nocs_path": nocspath,
            "mask_path": maskpath,
        }
        image_info.update(kwargs) #x: add additional args
        self.image_info.append(image_info)

    def image_reference(self, image_id):
        """Return a link to the image in its source Website or details about
        the image that help looking it up or debugging it.
        Override for your dataset, but pass to this function
        if you encounter images not in your dataset.
        """
        return ""

    def prepare(self, class_map=None):
        """Prepares the Dataset class for use.d
        """
        def clean_name(name):
            """Returns a shorter version of object names for cleaner display."""
            return ",".join(name.split(",")[:1])

        # Build (or rebuild) everything else from the info dicts.
        #self.num_classes = len(self.class_info)
        self.num_classes = 0

        #self.class_ids = np.arange(self.num_classes)
        self.class_ids = []

        #self.class_names = [clean_name(c["name"]) for c in self.class_info]
        self.class_names = []


        #self.class_from_source_map = {"{}.{}".format(info['source'], info['id']): id
        #                              for info, id in zip(self.class_info, self.class_ids)}
        self.class_from_source_map = {}

        for cls_info in self.class_info:
            source = cls_info["source"]
            if source == 'coco':
                map_key = "{}.{}".format(cls_info['source'], cls_info['id'])
                print("map_key,", map_key)
                self.class_from_source_map[map_key] = self.class_names.index(class_map[cls_info["name"]])
            else:
                self.class_ids.append(self.num_classes)
                self.num_classes += 1
                self.class_names.append(cls_info["name"])

                map_key = "{}.{}".format(cls_info['source'], cls_info['id'])
                self.class_from_source_map[map_key] = self.class_ids[-1]


        self.num_images = len(self.image_info)
        self._image_ids = np.arange(self.num_images)


        # Mapping from source class and image IDs to internal IDs
        self.image_from_source_map = {"{}.{}".format(info['source'], info['id']): id
                                      for info, id in zip(self.image_info, self.image_ids)}

        # Map sources to class_ids they support
        self.sources = list(set([i['source'] for i in self.class_info]))


        '''
        self.source_class_ids = {}
        # Loop over datasets
        for source in self.sources:
            self.source_class_ids[source] = []
            # Find classes that belong to this dataset
            for i, info in enumerate(self.class_info):
                # Include BG class in all datasets
                if i == 0 or source == info['source']:
                    self.source_class_ids[source].append(i)
        '''

        #print(self.class_names)
        #print(self.class_from_source_map)
        #print(self.sources)
        #print(self.source_class_ids)



    def map_source_class_id(self, source_class_id):
        """Takes a source class ID and returns the int class ID assigned to it.
        For example:
        dataset.map_source_class_id("coco.12") -> 23
        """
        return self.class_from_source_map[source_class_id] if source_class_id in self.class_from_source_map else None

    def get_source_class_id(self, class_id, source):
        """Map an internal class ID to the corresponding class ID in the source dataset."""
        info = self.class_info[class_id]
        assert info['source'] == source
        return info['id']

    def append_data(self, class_info, image_info):
        self.external_to_class_id = {}
        for i, c in enumerate(self.class_info):
            for ds, id in c["map"]:
                self.external_to_class_id[ds + str(id)] = i

        # Map external image IDs to internal ones.
        self.external_to_image_id = {}
        for i, info in enumerate(self.image_info):
            self.external_to_image_id[info["ds"] + str(info["id"])] = i

    @property
    def image_ids(self):
        return self._image_ids

    def source_image_link(self, image_id):
        """Returns the path or URL to the image.
        Override this to return a URL to the image if it's availble online for easy
        debugging.
        """
        return self.image_info[image_id]["path"]

    def load_image(self, image_id):
        """Load the specified image and return a [H,W,3] Numpy array.
        """
        # Load image
        image = scipy.misc.imread(self.image_info[image_id]['path'])
        # If grayscale. Convert to RGB for consistency.
        if image.ndim != 3:
            image = skimage.color.gray2rgb(image)
        return image

    def load_mask(self, image_id):
        """Load instance masks for the given image.
        Different datasets use different ways to store masks. Override this
        method to load instance masks and return them in the form of am
        array of binary masks of shape [height, width, instances].
        Returns:
            masks: A bool array of shape [height, width, instance count] with
                a binary mask per instance.
            class_ids: a 1D array of class IDs of the instance masks.
        """
        # Override this function to load a mask from your dataset.
        # Otherwise, it returns an empty mask.
        mask = np.empty([0, 0, 0])
        class_ids = np.empty([0], np.int32)
        return mask, class_ids


def resize_image(image, min_dim=None, max_dim=None, padding=False):
    """
    Resizes an image keeping the aspect ratio.
    min_dim: if provided, resizes the image such that it's smaller
        dimension == min_dim
    max_dim: if provided, ensures that the image longest side doesn't
        exceed this value.
    padding: If true, pads image with zeros so it's size is max_dim x max_dim
    Returns:
    image: the resized image
    window: (y1, x1, y2, x2). If max_dim is provided, padding might
        be inserted in the returned image. If so, this window is the
        coordinates of the image part of the full image (excluding
        the padding). The x2, y2 pixels are not included.
    scale: The scale factor used to resize the image
    padding: Padding added to the image [(top, bottom), (left, right), (0, 0)]
    """
    # Default window (y1, x1, y2, x2) and default scale == 1.
    h, w = image.shape[:2]
    window = (0, 0, h, w)
    scale = 1

    # Scale?
    if min_dim:
        # Scale up but not down
        scale = max(1, min_dim / min(h, w))
    # Does it exceed max dim?
    if max_dim:
        image_max = max(h, w)
        if round(image_max * scale) > max_dim:
            scale = max_dim / image_max
    # Resize image and mask
    if scale != 1:

        # scipy depracted use PIL, careful pil is (width, height)
        # image = scipy.misc.imresize(
        #     image, (round(h * scale), round(w * scale)))

        image = np.array( Image.fromarray(image).resize( ( round(w * scale), round(h * scale) ) ) )
    # Need padding?
    if padding:
        # Get new height and width
        h, w = image.shape[:2]
        top_pad = (max_dim - h) // 2
        bottom_pad = max_dim - h - top_pad
        left_pad = (max_dim - w) // 2
        right_pad = max_dim - w - left_pad
        padding = [(top_pad, bottom_pad), (left_pad, right_pad), (0, 0)]
        image = np.pad(image, padding, mode='constant', constant_values=0)
        window = (top_pad, left_pad, h + top_pad, w + left_pad)
    return image, window, scale, padding


def resize_mask(mask, scale, padding):
    """Resizes a mask using the given scale and padding.
    Typically, you get the scale and padding from resize_image() to
    ensure both, the image, the mask, and the coordinate map are resized consistently.
    scale: mask scaling factor
    padding: Padding to add to the mask in the form
            [(top, bottom), (left, right), (0, 0)]
    """
    h, w = mask.shape[:2]
    # for instance mask
    if len(mask.shape) == 3:
        mask = scipy.ndimage.zoom(mask, zoom=[scale, scale, 1], order=0)
        new_padding = padding
    # for coordinate map
    elif len(mask.shape) == 4:
        mask = scipy.ndimage.zoom(mask, zoom=[scale, scale, 1, 1], order=0)
        new_padding = padding + [(0, 0)]
    else:
        assert False

    mask = np.pad(mask, new_padding, mode='constant', constant_values=0)

    return mask


def minimize_mask(bbox, mask, mini_shape):
    """Resize masks to a smaller version to cut memory load.
    Mini-masks can then resized back to image scale using expand_masks()
    See inspect_data.ipynb notebook for more details.
    """
    # for instance mask
    if len(mask.shape)==3:
        mini_mask = np.zeros(mini_shape + (mask.shape[-1],), dtype=bool)
        for i in range(mask.shape[-1]):
            m = mask[:, :, i]
            y1, x1, y2, x2 = bbox[i][:4]
            m = m[y1:y2, x1:x2]*255
            # m = scipy.misc.imresize(m.astype(np.uint8), mini_shape, interp='nearest') DEPRECATED USE PILLOW
            m = np.array(Image.fromarray(m.astype(np.uint8)).resize(size=mini_shape, resample=PIL.Image.NEAREST))
            mini_mask[:, :, i] = np.where(m >= 128, 1, 0)

    # for coordinate map
    elif len(mask.shape)==4:
        assert mask.shape[-1] == 3 ## coordinate map

        mini_mask = np.zeros(mini_shape + mask.shape[-2:], dtype=np.float32)
        for i in range(mask.shape[-2]):
            m = mask[:, :, i, :]
            y1, x1, y2, x2 = bbox[i][:4]
            m = m[y1:y2, x1:x2, :]*255
            # m = scipy.misc.imresize(m.astype(np.uint8), mini_shape+(mask.shape[-1],), interp='nearest') DEPRECATED USE PILLOW
            m = np.array(Image.fromarray(m.astype(np.uint8)).resize(size=mini_shape, resample=PIL.Image.NEAREST))
            #print(m.shape)
            mini_mask[:, :, i, :] = m.astype(float)/255

    else:
        assert False
    return mini_mask


def expand_mask(bbox, mini_mask, image_shape):
    """Resizes mini masks back to image size. Reverses the change
    of minimize_mask().
    See inspect_data.ipynb notebook for more details.
    """
    # for instance mask
    if len(mini_mask.shape) == 3:
        mask = np.zeros(image_shape[:2] + (mini_mask.shape[-1],), dtype=bool)
        for i in range(mask.shape[-1]):
            m = mini_mask[:, :, i]
            y1, x1, y2, x2 = bbox[i][:4]
            h = y2 - y1
            w = x2 - x1
            m = scipy.misc.imresize(m.astype(float), (h, w), interp='bilinear')
            mask[y1:y2, x1:x2, i] = np.where(m >= 128, 1, 0)
    elif len(mini_mask.shape) == 4:
        assert mini_mask.shape[-1] == 3  ## coordinate map
        mask = np.zeros(image_shape[:2] + mini_mask.shape[-2:], dtype=np.float32)
        for i in range(mask.shape[-2]):
            m = mini_mask[:, :, i, :]
            y1, x1, y2, x2 = bbox[i][:4]
            h = y2 - y1
            w = x2 - x1
            m = scipy.misc.imresize(m.astype(float), (h, w, mini_mask.shape[-1]), interp='nearest')
            mask[y1:y2, x1:x2, i, :] = m

    return mask


# TODO: Build and use this function to reduce code duplication
def mold_mask(mask, config):
    pass


def unmold_mask(mask, bbox, image_shape):
    """Converts a mask generated by the neural network into a format similar
    to it's original shape.
    mask: [height, width] of type float. A small, typically 28x28 mask.
    bbox: [y1, x1, y2, x2]. The box to fit the mask in.
    Returns a binary mask with the same size as the original image.
    """
    threshold = 0.5
    y1, x1, y2, x2 = bbox

    # IMRESIZE DEPRECATED USE PIL: numpy.array(Image.fromarray(original_image).resize(newsize))
    # mask = scipy.misc.imresize(
    #     mask, (y2 - y1, x2 - x1), interp='bilinear').astype(np.float32) / 255.0

    # PIL USES WIDTH,HEIGHT
    #mask = np.array(Image.fromarray(mask).resize(size=(y2 - y1, x2 - x1), resample=PIL.Image.BILINEAR)).astype(np.float32) / 255.0
    mask = np.array(Image.fromarray(mask).resize(size=(x2 - x1, y2 - y1), resample=PIL.Image.BILINEAR)).astype(np.float32)
    mask = np.where(mask >= threshold, 1, 0).astype(np.uint8)

    # Put the mask in the right location.
    full_mask = np.zeros(image_shape[:2], dtype=np.uint8)
    full_mask[y1:y2, x1:x2] = mask # this was og
    #full_mask[x1:x2, y1:y2] = mask
    return full_mask

def unmold_coord(coord, bbox, image_shape):
    """Converts a mask generated by the neural network into a format similar
    to it's original shape.
    coord: [height, width, 3] of type float. A small, typically 28x28 mask.
    bbox: [y1, x1, y2, x2]. The box to fit the mask in.
    Returns a coordinate map with the same size as the original image.
    """
    y1, x1, y2, x2 = bbox

    #max_coord_x = np.amax(coord[:, :, 0])
    #max_coord_y = np.amax(coord[:, :, 1])
    #max_coord_z = np.amax(coord[:, :, 2])

    #print('before resize:')
    #print(max_coord_x, max_coord_y, max_coord_z)

    #coord = scipy.misc.imresize(
    #    coord, (y2 - y1, x2 - x1, 3), interp='nearest').astype(np.float32)/ 255.0
    #    #coord, (y2 - y1, x2 - x1, 3), interp='bilinear').astype(np.uint8)
    coord = cv2.resize(coord, (x2 - x1, y2 - y1), interpolation=cv2.INTER_LINEAR)

    #max_coord_x_resize = np.amax(coord[:, :, 0])
    #max_coord_y_resize = np.amax(coord[:, :, 1])
    #max_coord_z_resize = np.amax(coord[:, :, 2])

    #print('after resize:')
    #print(max_coord_x_resize, max_coord_y_resize, max_coord_z_resize)


    # Put the mask in the right location.
    full_coord= np.zeros(image_shape, dtype=np.float32)
    full_coord[y1:y2, x1:x2, :] = coord
    return full_coord

## for COCO
def annToRLE(ann, height, width):
    """
    Convert annotation which can be polygons, uncompressed RLE to RLE.
    :return: binary mask (numpy 2D array)
    """
    segm = ann['segmentation']
    if isinstance(segm, list):
        # polygon -- a single object might consist of multiple parts
        # we merge all parts into one mask rle code
        rles = maskUtils.frPyObjects(segm, height, width)
        rle = maskUtils.merge(rles)
    elif isinstance(segm['counts'], list):
        # uncompressed RLE
        rle = maskUtils.frPyObjects(segm, height, width)
    else:
        # rle
        rle = ann['segmentation']
    return rle

def annToMask(ann, height, width):
    """
    Convert annotation which can be polygons, uncompressed RLE, or RLE to binary mask.
    :return: binary mask (numpy 2D array)
    """
    rle = annToRLE(ann, height, width)
    m = maskUtils.decode(rle)
    return m


############################################################
#  Anchors
############################################################

def generate_anchors(scales, ratios, shape, feature_stride, anchor_stride):
    """
    scales: 1D array of anchor sizes in pixels. Example: [32, 64, 128]
    ratios: 1D array of anchor ratios of width/height. Example: [0.5, 1, 2]
    shape: [height, width] spatial shape of the feature map over which
            to generate anchors.
    feature_stride: Stride of the feature map relative to the image in pixels.
    anchor_stride: Stride of anchors on the feature map. For example, if the
        value is 2 then generate anchors for every other feature map pixel.
    """
    # Get all combinations of scales and ratios
    scales, ratios = np.meshgrid(np.array(scales), np.array(ratios))
    scales = scales.flatten()
    ratios = ratios.flatten()

    # Enumerate heights and widths from scales and ratios
    heights = scales / np.sqrt(ratios)
    widths = scales * np.sqrt(ratios)

    # Enumerate shifts in feature space
    shifts_y = np.arange(0, shape[0], anchor_stride) * feature_stride
    shifts_x = np.arange(0, shape[1], anchor_stride) * feature_stride
    shifts_x, shifts_y = np.meshgrid(shifts_x, shifts_y)

    # Enumerate combinations of shifts, widths, and heights
    box_widths, box_centers_x = np.meshgrid(widths, shifts_x)
    box_heights, box_centers_y = np.meshgrid(heights, shifts_y)

    # Reshape to get a list of (y, x) and a list of (h, w)
    box_centers = np.stack(
        [box_centers_y, box_centers_x], axis=2).reshape([-1, 2])
    box_sizes = np.stack([box_heights, box_widths], axis=2).reshape([-1, 2])

    # Convert to corner coordinates (y1, x1, y2, x2)
    boxes = np.concatenate([box_centers - 0.5 * box_sizes,
                            box_centers + 0.5 * box_sizes], axis=1)
    return boxes


def generate_pyramid_anchors(scales, ratios, feature_shapes, feature_strides,
                             anchor_stride):
    """Generate anchors at different levels of a feature pyramid. Each scale
    is associated with a level of the pyramid, but each ratio is used in
    all levels of the pyramid.
    Returns:
    anchors: [N, (y1, x1, y2, x2)]. All generated anchors in one array. Sorted
        with the same order of the given scales. So, anchors of scale[0] come
        first, then anchors of scale[1], and so on.
    """
    # Anchors
    # [anchor_count, (y1, x1, y2, x2)]
    anchors = []
    for i in range(len(scales)):
        anchors.append(generate_anchors(scales[i], ratios, feature_shapes[i],
                                        feature_strides[i], anchor_stride))
    return np.concatenate(anchors, axis=0)


############################################################
#  Evaluation
############################################################

def trim_zeros(x):
    """It's common to have tensors larger than the available data and
    pad with zeros. This function removes rows that are all zeros.
    x: [rows, columns].
    """

    pre_shape = x.shape
    assert len(x.shape) == 2, x.shape
    new_x = x[~np.all(x == 0, axis=1)]
    post_shape = new_x.shape
    assert pre_shape[0] == post_shape[0]
    assert pre_shape[1] == post_shape[1]

    return new_x


def compute_matches(gt_boxes, gt_class_ids, gt_masks,
                    pred_boxes, pred_class_ids, pred_scores, pred_masks,
                    iou_threshold=0.5, score_threshold=0.0):
    """Finds matches between prediction and ground truth instances.
    Returns:
        gt_match: 1-D array. For each GT box it has the index of the matched
                  predicted box.
        pred_match: 1-D array. For each predicted box, it has the index of
                    the matched ground truth box.
        overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # Trim zero padding
    # TODO: cleaner to do zero unpadding upstream
    if len(gt_class_ids)==0 or len(pred_class_ids) == 0:
        return -1 * np.ones([len(gt_class_ids)]), -1 * np.ones([len(pred_class_ids)]), None, np.zeros([0])

    pre_len = len(gt_boxes)
    gt_boxes = trim_zeros(gt_boxes)
    after_len = len(gt_boxes)
    assert pre_len == after_len
    gt_masks = gt_masks[..., :gt_boxes.shape[0]]

    pre_len = len(pred_boxes)
    pred_boxes = trim_zeros(pred_boxes)
    after_len = len(pred_boxes)
    assert pre_len == after_len
    pred_scores = pred_scores[:pred_boxes.shape[0]]

    # Sort predictions by score from high to low
    indices = np.argsort(pred_scores)[::-1]
    
    pred_boxes = pred_boxes[indices]
    pred_class_ids = pred_class_ids[indices]
    pred_scores = pred_scores[indices]
    pred_masks = pred_masks[..., indices]

    # Compute IoU overlaps [pred_masks, gt_masks]
    overlaps = compute_overlaps_masks(pred_masks, gt_masks)

    # Loop through predictions and find matching ground truth boxes
    match_count = 0
    pred_match = -1 * np.ones([pred_boxes.shape[0]])
    gt_match = -1 * np.ones([gt_boxes.shape[0]])
    
    
    for i in range(len(pred_boxes)):
        # Find best matching ground truth box
        # 1. Sort matches by score
        sorted_ixs = np.argsort(overlaps[i])[::-1]
        # 2. Remove low scores
        low_score_idx = np.where(overlaps[i, sorted_ixs] < score_threshold)[0]
        if low_score_idx.size > 0:
            sorted_ixs = sorted_ixs[:low_score_idx[0]]
        # 3. Find the match
        for j in sorted_ixs:
            # If ground truth box is already matched, go to next one
            if gt_match[j] > -1:
                continue
            # If we reach IoU smaller than the threshold, end the loop
            iou = overlaps[i, j]
            if iou < iou_threshold:
                break
            # Do we have a match?
            if pred_class_ids[i] == gt_class_ids[j]:
                match_count += 1
                gt_match[j] = i
                pred_match[i] = j
                break

    return gt_match, pred_match, overlaps, indices


def compute_3d_matches(gt_class_ids, gt_RTs, gt_scales, synset_names,
                       pred_boxes, pred_class_ids, pred_scores, pred_RTs, pred_scales,
                       iou_3d_thresholds, score_threshold=0):
    """Finds matches between prediction and ground truth instances.
    Returns:
        gt_matches: 2-D array. For each GT box it has the index of the matched
                  predicted box.
        pred_matches: 2-D array. For each predicted box, it has the index of
                    the matched ground truth box.
        overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # Trim zero padding
    # TODO: cleaner to do zero unpadding upstream
    num_pred = len(pred_class_ids)
    num_gt = len(gt_class_ids)
    indices = np.zeros(0)
    
    if num_pred:
        pred_boxes = trim_zeros(pred_boxes).copy()
        pred_scores = pred_scores[:pred_boxes.shape[0]].copy()

        # Sort predictions by score from high to low
        indices = np.argsort(pred_scores)[::-1]
        
        pred_boxes = pred_boxes[indices].copy()
        pred_class_ids = pred_class_ids[indices].copy()
        pred_scores = pred_scores[indices].copy()
        pred_scales = pred_scales[indices].copy()
        pred_RTs = pred_RTs[indices].copy()

        
    # pred_3d_bboxs = []
    # for i in range(num_pred):
    #     noc_cube = get_3d_bbox(pred_scales[i, :], 0)
    #     pred_bbox_3d = transform_coordinates_3d(noc_cube, pred_RTs[i])
    #     pred_3d_bboxs.append(pred_bbox_3d)

    # # compute 3d bbox for ground truths
    # # print('Compute gt bboxes...')
    # gt_3d_bboxs = []
    # for j in range(num_gt):
    #     noc_cube = get_3d_bbox(gt_scales[j], 0)
    #     gt_3d_bbox = transform_coordinates_3d(noc_cube, gt_RTs[j])
    #     gt_3d_bboxs.append(gt_3d_bbox)

    # Compute IoU overlaps [pred_bboxs gt_bboxs]
    #overlaps = [[0 for j in range(num_gt)] for i in range(num_pred)]
    overlaps = np.zeros((num_pred, num_gt), dtype=np.float32)
    for i in range(num_pred):
        for j in range(num_gt):
            #overlaps[i, j] = compute_3d_iou(pred_3d_bboxs[i], gt_3d_bboxs[j], gt_handle_visibility[j], 
            #    synset_names[pred_class_ids[i]], synset_names[gt_class_ids[j]])
            overlaps[i, j] = compute_3d_iou_new(pred_RTs[i], gt_RTs[j], pred_scales[i, :], gt_scales[j], \
                                                      synset_names[pred_class_ids[i]], \
                                                      synset_names[gt_class_ids[j]])

    # Loop through predictions and find matching ground truth boxes
    num_iou_3d_thres = len(iou_3d_thresholds)
    pred_matches = -1 * np.ones([num_iou_3d_thres, num_pred])
    gt_matches = -1 * np.ones([num_iou_3d_thres, num_gt])

    for s, iou_thres in enumerate(iou_3d_thresholds):
        for i in range(len(pred_boxes)):
            # Find best matching ground truth box
            # 1. Sort matches by score
            sorted_ixs = np.argsort(overlaps[i])[::-1]
            # 2. Remove low scores
            low_score_idx = np.where(overlaps[i, sorted_ixs] < score_threshold)[0]
            if low_score_idx.size > 0:
                sorted_ixs = sorted_ixs[:low_score_idx[0]]
            # 3. Find the match
            for j in sorted_ixs:
                # If ground truth box is already matched, go to next one
                #print('gt_match: ', gt_match[j])
                if gt_matches[s, j] > -1:
                    continue
                # If we reach IoU smaller than the threshold, end the loop
                iou = overlaps[i, j]
                #print('iou: ', iou)
                if iou < iou_thres:
                    break
                # Do we have a match?
                if not pred_class_ids[i] == gt_class_ids[j]:
                    continue

                if iou > iou_thres:
                    gt_matches[s, j] = i
                    pred_matches[s, i] = j
                    break

    return gt_matches, pred_matches, overlaps, indices


def compute_ap_from_matches_scores(pred_match, pred_scores, gt_match):
    # sort the scores from high to low
    # print(pred_match.shape, pred_scores.shape)
    assert pred_match.shape[0] == pred_scores.shape[0]

    # input("here")


    score_indices = np.argsort(pred_scores)[::-1] # sort descendingly on IoU (returns indices)
    pred_scores = pred_scores[score_indices]
    pred_match  = pred_match[score_indices]

    precisions = np.cumsum(pred_match > -1) / (np.arange(len(pred_match)) + 1)
    recalls = np.cumsum(pred_match > -1).astype(np.float32) / len(gt_match)

    # Pad with start and end values to simplify the math
    precisions = np.concatenate([[0], precisions, [0]])
    recalls = np.concatenate([[0], recalls, [1]])

    #print("precision:", precisions)
    #print("recalls:", recalls)

    # Ensure precision values decrease but don't increase. This way, the
    # precision value at each recall threshold is the maximum it can be
    # for all following recall thresholds, as specified by the VOC paper.
    for i in range(len(precisions) - 2, -1, -1):
        precisions[i] = np.maximum(precisions[i], precisions[i + 1])
    
    #print("precision_CAP:", precisions)
    
    # Compute mean AP over recall range
    indices = np.where(recalls[:-1] != recalls[1:])[0] + 1 #[:-1] all but last element. #[1:] all but first element
    #print("indices:", indices)
    ap = np.sum((recalls[indices] - recalls[indices - 1]) * precisions[indices])
    #print("ap:", ap)

    if math.isnan(ap):
        ap = 0.0

    return ap


def compute_RT_overlaps(gt_class_ids, gt_RTs, 
                        pred_class_ids, pred_RTs, 
                        synset_names):
    """Finds overlaps between prediction and ground truth instances.
    Returns:
        overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # print('num of gt instances: {}, num of pred instances: {}'.format(len(gt_class_ids), len(gt_class_ids)))
    num_pred = len(pred_class_ids)
    num_gt = len(gt_class_ids)

    # Compute IoU overlaps [pred_bboxs gt_bboxs]
    #overlaps = [[0 for j in range(num_gt)] for i in range(num_pred)]
    overlaps = np.zeros((num_pred, num_gt, 2))

    for i in range(num_pred):
        for j in range(num_gt):
            overlaps[i, j, :] = compute_RT_degree_cm_symmetry(pred_RTs[i], 
                                                              gt_RTs[j], 
                                                              gt_class_ids[j],
                                                              synset_names)
            
    return overlaps


def compute_RT_projection_2d_overlaps(gt_class_ids, gt_RTs, gt_handle_visibility,
                                      pred_class_ids, pred_RTs, 
                                      meshes, intrinsics, synset_names):
    """Finds overlaps between prediction and ground truth instances.
    Returns:
        overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # print('num of gt instances: {}, num of pred instances: {}'.format(len(gt_class_ids), len(gt_class_ids)))
    num_pred = len(pred_class_ids)
    num_gt = len(gt_class_ids)

    # Compute IoU overlaps [pred_bboxs gt_bboxs]
    #overlaps = [[0 for j in range(num_gt)] for i in range(num_pred)]
    overlaps = np.zeros((num_pred, num_gt))

    for i in range(num_pred):
        for j in range(num_gt):
            overlaps[i, j] = compute_RT_projection_2d_symmetry(pred_RTs[i], 
                                                               gt_RTs[j], 
                                                               gt_class_ids[j], 
                                                               gt_handle_visibility[j],
                                                               meshes[gt_class_ids[j]],
                                                               intrinsics,
                                                               synset_names)
            
    return overlaps


def compute_match_from_projection_2d_dist(overlaps, pred_class_ids, gt_class_ids, projection_thres_list):
    num_projection_thres = len(projection_thres_list)

    num_pred = len(pred_class_ids)
    num_gt = len(gt_class_ids)

    pred_matches = -1 * np.ones((num_projection_thres, num_pred))
    gt_matches = -1 * np.ones((num_projection_thres, num_gt))

    if num_pred == 0 or num_gt == 0:
        return gt_matches, pred_matches


    assert num_pred == overlaps.shape[0]
    assert num_gt == overlaps.shape[1]
    

                 
    for p, projection_thres in enumerate(projection_thres_list):
        for i in range(num_pred):
            # Find best matching ground truth box
            # 1. Sort matches by scores from low to high
            sorted_ixs = np.argsort(overlaps[i, :])
            # 2. Remove low scores
            # low_score_idx = np.where(sum_degree_shift >= 100)[0]
            # if low_score_idx.size > 0:
            #     sorted_ixs = sorted_ixs[:low_score_idx[0]]
            # 3. Find the match
            for j in sorted_ixs:
                # If ground truth box is already matched, go to next one
                #print(j, len(gt_match), len(pred_class_ids), len(gt_class_ids))
                if gt_matches[p, j] > -1 or pred_class_ids[i] != gt_class_ids[j]:
                    continue
                # If we reach IoU smaller than the threshold, end the loop
                if overlaps[i, j] > projection_thres:
                    continue

                gt_matches[p, j] = i
                pred_matches[p, i] = j
                break

    return gt_matches, pred_matches


def compute_match_from_degree_cm(overlaps, pred_class_ids, gt_class_ids, degree_thres_list, shift_thres_list):
    num_degree_thres = len(degree_thres_list)
    num_shift_thres = len(shift_thres_list)


    num_pred = len(pred_class_ids)
    num_gt = len(gt_class_ids)

    pred_matches = -1 * np.ones((num_degree_thres, num_shift_thres, num_pred))
    gt_matches = -1 * np.ones((num_degree_thres, num_shift_thres, num_gt))

    if num_pred == 0 or num_gt == 0:
        return gt_matches, pred_matches


    assert num_pred == overlaps.shape[0]
    assert num_gt == overlaps.shape[1]
    assert overlaps.shape[2] == 2
    

    for d, degree_thres in enumerate(degree_thres_list):                
        for s, shift_thres in enumerate(shift_thres_list):
            for i in range(num_pred):
                
                ### Find best matching ground truth box
                
                # 1. Sort matches by scores from low to high
                sum_degree_shift = np.sum(overlaps[i, :, :], axis=-1)
                sorted_ixs = np.argsort(sum_degree_shift)
                
                # 2. Remove low scores
                # low_score_idx = np.where(sum_degree_shift >= 100)[0]
                # if low_score_idx.size > 0:
                #     sorted_ixs = sorted_ixs[:low_score_idx[0]]
                
                # 3. Find the match
                for j in sorted_ixs:
                    # If ground truth box is already matched, go to next one
                    #print(j, len(gt_match), len(pred_class_ids), len(gt_class_ids))
                    if gt_matches[d, s, j] > -1 or pred_class_ids[i] != gt_class_ids[j]:
                        continue
                    # If we reach DEGREE or CM error bigger than the threshold, end the loop
                    if overlaps[i, j, 0] > degree_thres or overlaps[i, j, 1] > shift_thres:
                        continue

                    gt_matches[d, s, j] = i
                    pred_matches[d, s, i] = j
                    break

    return gt_matches, pred_matches



def compute_degree_cm_mAP(final_results, synset_names, log_dir, degree_thresholds=[360], shift_thresholds=[100], iou_3d_thresholds=[0.1], iou_pose_thres=0.1, use_matches_for_pose=False):
    """Compute Average Precision at a set IoU threshold (default 0.5).
    Returns:
    mAP: Mean Average Precision
    precisions: List of precisions at different class score thresholds.
    recalls: List of recall values at different class score thresholds.
    overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    
    font = {'family' : 'normal',
            'weight' : 'bold',
            'size'   : 35}

    matplotlib.rc('font', **font)

    num_classes = len(synset_names)
    degree_thres_list = list(degree_thresholds) + [360]
    #degree_thres_list = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15, 360]
    num_degree_thres = len(degree_thres_list)

    shift_thres_list = list(shift_thresholds) + [100]
    #shift_thres_list = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15, 100]
    num_shift_thres = len(shift_thres_list)

    iou_thres_list = list(iou_3d_thresholds)
    num_iou_thres = len(iou_thres_list)

    if use_matches_for_pose:
        assert iou_pose_thres in iou_thres_list

    iou_3d_aps = np.zeros((num_classes + 1, num_iou_thres))
    iou_pred_matches_all = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    iou_pred_scores_all  = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    iou_gt_matches_all   = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    
    pose_aps = np.zeros((num_classes + 1, num_degree_thres, num_shift_thres))
    pose_pred_matches_all = [np.zeros((num_degree_thres, num_shift_thres, 0)) for _  in range(num_classes)]
    pose_gt_matches_all  = [np.zeros((num_degree_thres, num_shift_thres, 0)) for _  in range(num_classes)]
    pose_pred_scores_all = [np.zeros((num_degree_thres, num_shift_thres, 0)) for _  in range(num_classes)]

    # loop over results to gather pred matches and gt matches for iou and pose metrics
    progress = 0
    for progress, result in enumerate(final_results):
        print(progress, len(final_results))
        print(final_results[progress]['image_path'])
        
        # Get GROUND TRUTHS
        gt_class_ids = result['gt_class_ids'].astype(np.int32)
        gt_RTs = np.array(result['gt_RTs'])
        gt_scales = np.array(result['gt_scales'])
        #gt_handle_visibility = result['gt_handle_visibility']

        # Get PREDICTIONS
        pred_bboxes = np.array(result['pred_bboxes'])
        pred_class_ids = result['pred_class_ids']
        pred_scales = result['pred_scales']
        pred_scores = result['pred_scores']
        pred_RTs = np.array(result['pred_RTs'])
        #print(pred_bboxes.shape[0], pred_class_ids.shape[0], pred_scores.shape[0], pred_RTs.shape[0])

        if len(gt_class_ids) == 0 and len(pred_class_ids) == 0:
            continue


        # Get GTs and PREDs for each
        for cls_id in range(1, num_classes):
            
            cls_gt_class_ids = gt_class_ids[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros(0)
            cls_gt_scales = gt_scales[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros((0, 3))
            cls_gt_RTs = gt_RTs[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros((0, 4, 4))

            cls_pred_class_ids = pred_class_ids[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros(0)
            cls_pred_bboxes =  pred_bboxes[pred_class_ids==cls_id, :] if len(pred_class_ids) else np.zeros((0, 4))
            cls_pred_scores = pred_scores[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros(0)
            cls_pred_RTs = pred_RTs[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros((0, 4, 4))
            cls_pred_scales = pred_scales[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros((0, 3))


            # if len(cls_gt_class_ids) == 0 and len(cls_pred_class_ids) == 0:
            #     continue

            # elif len(cls_gt_class_ids) > 0 and len(cls_pred_class_ids) == 0:
            #     iou_gt_matches_all[cls_id] = np.concatenate((iou_gt_matches_all[cls_id], -1*np.ones((num_iou_thres, len(cls_gt_class_ids)))), axis=-1)
            #     if not use_matches_for_pose:
            #         pose_gt_matches_all[cls_id] = np.concatenate((pose_gt_matches_all[cls_id], -1*np.ones((num_degree_thres, num_shift_thres, len(cls_gt_class_ids)))), axis=-1)
            #     continue

            # elif len(cls_pred_class_ids)>0 and len(cls_gt_class_ids)==0:
            #     assert iou_pred_matches_all[cls_id].shape[1] == iou_pred_scores_all[cls_id].shape[1]

            #     iou_pred_matches_all[cls_id] = np.concatenate((iou_pred_matches_all[cls_id], -1*np.ones((num_iou_thres, len(cls_pred_class_ids)))),  axis=-1)
            #     cls_pred_scores_tile = np.tile(cls_pred_scores, (num_iou_thres, 1))
            #     iou_pred_scores_all[cls_id] = np.concatenate((iou_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)

            #     assert iou_pred_matches_all[cls_id].shape[1] == iou_pred_scores_all[cls_id].shape[1]

            #     if not use_matches_for_pose:
            #         pose_pred_matches_all[cls_id] = np.concatenate((pose_pred_matches_all[cls_id], -1*np.ones((num_degree_thres, num_shift_thres, len(cls_pred_class_ids)))), axis=-1)
            #         cls_pred_scores_tile = np.tile(cls_pred_scores, (num_degree_thres, num_shift_thres, 1))
            #         pose_pred_scores_all[cls_id] = np.concatenate((pose_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)
            #     continue



            # calculate the overlap between each gt instance and pred instance
            # if synset_names[cls_id] != 'mug':
            #     cls_gt_handle_visibility = np.ones_like(cls_gt_class_ids)
            # else:
            #     cls_gt_handle_visibility = gt_handle_visibility[gt_class_ids==cls_id] if len(gt_class_ids) else np.ones(0)


            iou_cls_gt_match, iou_cls_pred_match, _, iou_pred_indices = compute_3d_matches(cls_gt_class_ids, cls_gt_RTs, cls_gt_scales, synset_names,
                                                                                           cls_pred_bboxes, cls_pred_class_ids, cls_pred_scores, cls_pred_RTs, cls_pred_scales,
                                                                                           iou_thres_list)
            if len(iou_pred_indices):
                cls_pred_class_ids = cls_pred_class_ids[iou_pred_indices]
                cls_pred_RTs = cls_pred_RTs[iou_pred_indices]
                cls_pred_scores = cls_pred_scores[iou_pred_indices]
                cls_pred_bboxes = cls_pred_bboxes[iou_pred_indices]


            iou_pred_matches_all[cls_id] = np.concatenate((iou_pred_matches_all[cls_id], iou_cls_pred_match), axis=-1)
            cls_pred_scores_tile = np.tile(cls_pred_scores, (num_iou_thres, 1))
            iou_pred_scores_all[cls_id] = np.concatenate((iou_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)
            assert iou_pred_matches_all[cls_id].shape[1] == iou_pred_scores_all[cls_id].shape[1]
            iou_gt_matches_all[cls_id] = np.concatenate((iou_gt_matches_all[cls_id], iou_cls_gt_match), axis=-1)

            if use_matches_for_pose:
                thres_ind = list(iou_thres_list).index(iou_pose_thres)

                iou_thres_pred_match = iou_cls_pred_match[thres_ind, :]

                
                cls_pred_class_ids = cls_pred_class_ids[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros(0)
                cls_pred_RTs = cls_pred_RTs[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros((0, 4, 4))
                cls_pred_scores = cls_pred_scores[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros(0)
                cls_pred_bboxes = cls_pred_bboxes[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros((0, 4))


                iou_thres_gt_match = iou_cls_gt_match[thres_ind, :]
                cls_gt_class_ids = cls_gt_class_ids[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros(0)
                cls_gt_RTs = cls_gt_RTs[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros((0, 4, 4))
                #cls_gt_handle_visibility = cls_gt_handle_visibility[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros(0)



            RT_overlaps = compute_RT_overlaps(cls_gt_class_ids, cls_gt_RTs,
                                              cls_pred_class_ids, cls_pred_RTs,
                                              synset_names)


            pose_cls_gt_match, pose_cls_pred_match = compute_match_from_degree_cm(RT_overlaps, 
                                                                                  cls_pred_class_ids, 
                                                                                  cls_gt_class_ids, 
                                                                                  degree_thres_list, 
                                                                                  shift_thres_list)
            

            pose_pred_matches_all[cls_id] = np.concatenate((pose_pred_matches_all[cls_id], pose_cls_pred_match), axis=-1)
            
            cls_pred_scores_tile = np.tile(cls_pred_scores, (num_degree_thres, num_shift_thres, 1))
            pose_pred_scores_all[cls_id]  = np.concatenate((pose_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)
            assert pose_pred_scores_all[cls_id].shape[2] == pose_pred_matches_all[cls_id].shape[2], '{} vs. {}'.format(pose_pred_scores_all[cls_id].shape, pose_pred_matches_all[cls_id].shape)
            pose_gt_matches_all[cls_id] = np.concatenate((pose_gt_matches_all[cls_id], pose_cls_gt_match), axis=-1)

    
    ###### 3D IOU PLOT
    ######

    # Set figure parameters
    fig_iou = plt.figure(figsize=(9, 13.5)) 
    ax_iou = plt.subplot(111)
    plt.title('3D IoU mAP', fontweight="bold")
    #plt.ylabel('mAP (%)', fontweight="bold")
    plt.ylim((0, 105))
    plt.xlabel('3D IoU (%) thresholds', fontweight="bold")
    plt.xticks([25,50,75,100])
    plt.grid(True)
    plt.margins(x=0)
    plt.tight_layout()

    # output paths
    iou_output_path = os.path.join(log_dir, 'IoU_3D_AP_{}-{}.png'.format(iou_thres_list[0], iou_thres_list[-1]))
    iou_dict_pkl_path = os.path.join(log_dir, 'IoU_3D_AP_{}-{}.pkl'.format(iou_thres_list[0], iou_thres_list[-1]))

    # get IoU thresholds
    iou_dict = {}
    iou_dict['thres_list'] = iou_thres_list
    my_iou_thres_list = np.multiply(iou_dict['thres_list'],100)
    print(iou_thres_list)
    
    # Compute APs for all classes
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]

        #print(class_name)
        for s, iou_thres in enumerate(my_iou_thres_list):
            iou_3d_aps[cls_id, s] = compute_ap_from_matches_scores(iou_pred_matches_all[cls_id][s, :],
                                                                   iou_pred_scores_all[cls_id][s, :],
                                                                   iou_gt_matches_all[cls_id][s, :])    
            #print("({:.1f},{:.1f})".format(iou_thres, iou_3d_aps[cls_id, s]*100), end='')
        #print("\n")
        ax_iou.plot(my_iou_thres_list, iou_3d_aps[cls_id, :]*100, label=class_name, linewidth=6)
    
    # Copy the matrix that has category-information (in other words: the non-mean AP)
    iou_3d_aps_cats = iou_3d_aps.copy()

    # Compute the mean (over the classes) AP
    iou_3d_aps[-1, :] = np.mean(iou_3d_aps[1:-1, :], axis=0)

    # Plot it
    ax_iou.plot(my_iou_thres_list, iou_3d_aps[-1, :]*100, label='mean', linewidth=6)

    # Draw the legend
    #ax_iou.legend()

    # Save and close
    fig_iou.savefig(iou_output_path)
    plt.close(fig_iou)

    # Save to python dict
    iou_dict['aps'] = iou_3d_aps

    # Save to pkl file
    with open(iou_dict_pkl_path, 'wb') as f:
        cPickle.dump(iou_dict, f)
    

    ###### IRRELEVANT draw pose AP vs. thresholds
    ######
    if use_matches_for_pose:
        prefix='Pose_Only_'
    else:
        prefix='Pose_Detection_'


    pose_dict_pkl_path = os.path.join(log_dir, prefix+'AP_{}-{}degree_{}-{}cm.pkl'.format(degree_thres_list[0], degree_thres_list[-2], 
                                                                                          shift_thres_list[0], shift_thres_list[-2]))
    pose_dict = {}
    pose_dict['degree_thres'] = degree_thres_list
    pose_dict['shift_thres_list'] = shift_thres_list

    for i, degree_thres in enumerate(degree_thres_list):                
        for j, shift_thres in enumerate(shift_thres_list):
            # print(i, j)
            for cls_id in range(1, num_classes):
                cls_pose_pred_matches_all = pose_pred_matches_all[cls_id][i, j, :]
                cls_pose_gt_matches_all = pose_gt_matches_all[cls_id][i, j, :]
                cls_pose_pred_scores_all = pose_pred_scores_all[cls_id][i, j, :]

                pose_aps[cls_id, i, j] = compute_ap_from_matches_scores(cls_pose_pred_matches_all, 
                                                                        cls_pose_pred_scores_all, 
                                                                        cls_pose_gt_matches_all)

            pose_aps[-1, i, j] = np.mean(pose_aps[1:-1, i, j])
    
    pose_dict['aps'] = pose_aps
    with open(pose_dict_pkl_path, 'wb') as f:
        cPickle.dump(pose_dict, f)


    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]
        print(class_name)
        # print(np.amin(aps[i, :, :]), np.amax(aps[i, :, :]))
    
        #ap_image = cv2.resize(pose_aps[cls_id, :, :]*255, (320, 320), interpolation = cv2.INTER_LINEAR)
        fig_iou = plt.figure()
        ax_iou = plt.subplot(111)
        plt.ylabel('Rotation thresholds/degree')
        plt.ylim((degree_thres_list[0], degree_thres_list[-2]))
        plt.xlabel('translation/cm')
        plt.xlim((shift_thres_list[0], shift_thres_list[-2]))
        plt.imshow(pose_aps[cls_id, :-1, :-1], cmap='jet', interpolation='bilinear')

        output_path = os.path.join(log_dir, prefix+'AP_{}_{}-{}degree_{}-{}cm.png'.format(class_name, 
                                                                                   degree_thres_list[0], degree_thres_list[-2], 
                                                                                   shift_thres_list[0], shift_thres_list[-2]))
        plt.colorbar()
        plt.savefig(output_path)
        plt.close(fig_iou)        
    
    #ap_mean_image = cv2.resize(pose_aps[-1, :, :]*255, (320, 320), interpolation = cv2.INTER_LINEAR) 
    
    fig_pose = plt.figure()
    ax_pose = plt.subplot(111)
    plt.ylabel('Rotation thresholds/degree')
    plt.ylim((degree_thres_list[0], degree_thres_list[-2]))
    plt.xlabel('translation/cm')
    plt.xlim((shift_thres_list[0], shift_thres_list[-2]))
    plt.imshow(pose_aps[-1, :-1, :-1], cmap='jet', interpolation='bilinear')
    output_path = os.path.join(log_dir, prefix+'mAP_{}-{}degree_{}-{}cm.png'.format(degree_thres_list[0], degree_thres_list[-2], 
                                                                             shift_thres_list[0], shift_thres_list[-2]))
    plt.colorbar()
    plt.savefig(output_path)
    plt.close(fig_pose)

    ###### TRANSLATION PLOT
    ######
    fig_rot = plt.figure(figsize=(9, 13.5))
    ax_rot = plt.subplot(111)
    #plt.ylabel('AP')
    plt.title('Translation mAP', fontweight="bold")
    plt.ylim((0, 105))
    plt.xticks([5,10,15])
    plt.xlabel('Translation threshold (cm)', fontweight="bold")
    plt.grid(True)
    plt.margins(x=0)
    plt.tight_layout()
    
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]
        
        ### MY STUFF
        print("Class (trans):", class_name)
        poses_tr = np.around(pose_aps[cls_id, -1, :-1], 3)
        for my_idx in range(0, len(shift_thres_list)-1):
            print("({:.1f},{:.1f})".format(shift_thres_list[my_idx], poses_tr[my_idx]*100), end='')
        print("\n")
        ###
        ax_rot.plot(shift_thres_list[:-1], pose_aps[cls_id, -1, :-1]*100, label=class_name, linewidth=6)

    ax_rot.plot(shift_thres_list[:-1], pose_aps[-1, -1, :-1]*100, label='mean', linewidth=6)
    output_path = os.path.join(log_dir, prefix+'mAP_{}-{}cm.png'.format(shift_thres_list[0], shift_thres_list[-2]))
    ax_rot.legend(prop={'size': 38})
    
    fig_rot.savefig(output_path)
    
    plt.close(fig_rot)


    ###### ROTATION PLOT
    ######
    fig_trans = plt.figure(figsize=(9, 13.5))
    ax_trans = plt.subplot(111)
    plt.title('Rotation mAP', fontweight="bold")
    #plt.ylabel('AP')
    plt.ylim((0, 105))
    plt.xticks([5,10,15])
    plt.xlabel('Rotation threshold (degree)', fontweight="bold")
    plt.grid(True)
    plt.margins(x=0)
    plt.tight_layout()
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]
        
        ### my stuff
        # print(class_name)
        # poses_ro = np.around(pose_aps[cls_id, :-1, shift_thres_list.index(5)],3)
        # for my_idx in range(0, len(degree_thres_list)-1):
        #     print("({:.1f},{:.1f})".format(degree_thres_list[my_idx], poses_ro[my_idx]*100), end='')
        #print("\n")
        ####
        ax_trans.plot(degree_thres_list[:-1], pose_aps[cls_id, :-1, -1]*100, label=class_name, linewidth=6)

    ax_trans.plot(degree_thres_list[:-1], pose_aps[-1, :-1, -1]*100, label='mean', linewidth=6)
    output_path = os.path.join(log_dir, prefix+'mAP_{}-{}degree.png'.format(degree_thres_list[0], degree_thres_list[-2]))
    
    #ax_trans.legend()
    fig_trans.savefig(output_path)
    plt.close(fig_trans)

    iou_aps = iou_3d_aps
    
    mean_poses_tr = np.around(pose_aps[-1, -1, :-1], 3)
    print("\nMean translation")
    for my_idx in range(0, len(shift_thres_list)-1):
            print("({:.1f},{:.1f})".format(shift_thres_list[my_idx], mean_poses_tr[my_idx]*100), end='')

    # mean_poses_ro = np.around(pose_aps[-1, :-1, shift_thres_list.index(5)], 3)
    # print("\n\nMean rotations")
    # for my_idx in range(0, len(degree_thres_list)-1):
    #         print("({:.1f},{:.1f})".format(degree_thres_list[my_idx], mean_poses_ro[my_idx]*100), end='')

    # print("\n\nMean 3D IoU")
    # for my_idx in range(0, len(iou_thres_list)):
    #         print("({:.1f},{:.1f})".format(iou_thres_list[my_idx]*100, iou_3d_aps[-1, iou_thres_list.index(iou_thres_list[my_idx])]*100), end='')


    ### Print the mean AP measurements
    print('3D IoU at 25: {:.1f}'.format(iou_aps[-1, iou_thres_list.index(0.25)] * 100))
    print('3D IoU at 50: {:.1f}'.format(iou_aps[-1, iou_thres_list.index(0.5)] * 100))
    print('5 degree, 5cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(5),shift_thres_list.index(5)] * 100))
    print('5 degree, 100cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(5),shift_thres_list.index(100)] * 100))
    print('10 degree, 5cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(10),shift_thres_list.index(5)] * 100))
    print('10 degree, 10cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(10),shift_thres_list.index(10)] * 100))
    print('15 degree, 5cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(15),shift_thres_list.index(5)] * 100))
    print('15 degree, 15cm: {:.1f}'.format(pose_aps[-1, degree_thres_list.index(15),shift_thres_list.index(15)] * 100))
    print('Copy into latex:{:.1f} & {:.1f} & {:.1f} & {:.1f} & {:.1f}\n\n'.format(  
                iou_aps[-1, iou_thres_list.index(0.25)] * 100,
                iou_aps[-1, iou_thres_list.index(0.5)] * 100,
                pose_aps[-1, degree_thres_list.index(5),shift_thres_list.index(5)] * 100,
                pose_aps[-1, degree_thres_list.index(10),shift_thres_list.index(5)] * 100,
                pose_aps[-1, degree_thres_list.index(10),shift_thres_list.index(10)] * 100 ))

    
    ### Print the AP measurements per class
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]

        # 3D IoU
        print('3D IoU at 25: {:.1f} for {}'.format(iou_3d_aps_cats[cls_id, iou_thres_list.index(0.25)] * 100, class_name))
        print('3D IoU at 50: {:.1f} for {}'.format(iou_3d_aps_cats[cls_id, iou_thres_list.index(0.5)] * 100, class_name))

        # Pose
        print('5 degree, 5cm: {:.1f} for {}'.format(pose_aps[cls_id, degree_thres_list.index(5),shift_thres_list.index(5)] * 100, class_name))
        print('10 degree, 5cm: {:.1f} for {}'.format(pose_aps[cls_id, degree_thres_list.index(10),shift_thres_list.index(5)] * 100, class_name))
        print('10 degree, 10cm: {:.1f} for {}'.format(pose_aps[cls_id, degree_thres_list.index(10),shift_thres_list.index(10)] * 100, class_name))

        print('Copy into latex:{:.1f} & {:.1f} & {:.1f} & {:.1f} & {:.1f}'.format(  
                iou_3d_aps_cats[cls_id, iou_thres_list.index(0.25)] * 100,
                iou_3d_aps_cats[cls_id, iou_thres_list.index(0.5)] * 100,
                pose_aps[cls_id, degree_thres_list.index(5),shift_thres_list.index(5)] * 100,
                pose_aps[cls_id, degree_thres_list.index(10),shift_thres_list.index(5)] * 100,
                pose_aps[cls_id, degree_thres_list.index(10),shift_thres_list.index(10)] * 100) )
        print("\n")
        
    return iou_3d_aps, pose_aps


def compute_2D_projection_mAP(final_results, synset_names, log_dir, meshes, intrinsics, iou_3d_thresholds=[0.1], iou_pose_thres=0.1, projection_2d_thresholds=[5], use_matches_for_pose=False):
    """Compute Average Precision at a set IoU threshold (default 0.5).
    Returns:
    mAP: Mean Average Precision
    precisions: List of precisions at different class score thresholds.
    recalls: List of recall values at different class score thresholds.
    overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    
    num_classes = len(synset_names)
    
    iou_thres_list = list(iou_3d_thresholds)
    num_iou_thres = len(iou_thres_list)

    projection_thres_list = list(projection_2d_thresholds)
    num_projection_thres = len(projection_thres_list)

    if use_matches_for_pose:
        assert iou_pose_thres in iou_thres_list

    iou_3d_aps = np.zeros((num_classes + 1, num_iou_thres))
    iou_pred_matches_all = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    iou_pred_scores_all  = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    iou_gt_matches_all   = [np.zeros((num_iou_thres, 0)) for _ in range(num_classes)]
    
    
    projection_2d_aps = np.zeros((num_classes + 1, num_projection_thres))
    projection_pred_matches_all = [np.zeros((num_projection_thres, 0)) for _ in range(num_classes)]
    projection_pred_scores_all  = [np.zeros((num_projection_thres, 0)) for _ in range(num_classes)]
    projection_gt_matches_all   = [np.zeros((num_projection_thres, 0)) for _ in range(num_classes)]


    # loop over results to gather pred matches and gt matches for iou and pose metrics
    progress = 0
    for progress, result in enumerate(final_results):
        print(progress, len(final_results))
        gt_class_ids = result['gt_class_ids'].astype(np.int32)
        gt_RTs = np.array(result['gt_RTs'])
        gt_scales = np.array(result['gt_scales'])
        gt_handle_visibility = result['gt_handle_visibility']
    
        pred_bboxes = np.array(result['pred_bboxes'])
        pred_class_ids = result['pred_class_ids']
        pred_scales = result['pred_scales']
        pred_scores = result['pred_scores']
        pred_RTs = np.array(result['pred_RTs'])
        #print(pred_bboxes.shape[0], pred_class_ids.shape[0], pred_scores.shape[0], pred_RTs.shape[0])

        if len(gt_class_ids) == 0 and len(pred_class_ids) == 0:
            continue


        for cls_id in range(1, num_classes):
            # get gt and predictions in this class
            cls_gt_class_ids = gt_class_ids[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros(0)
            cls_gt_scales = gt_scales[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros((0, 3))
            cls_gt_RTs = gt_RTs[gt_class_ids==cls_id] if len(gt_class_ids) else np.zeros((0, 4, 4))

            cls_pred_class_ids = pred_class_ids[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros(0)
            cls_pred_bboxes =  pred_bboxes[pred_class_ids==cls_id, :] if len(pred_class_ids) else np.zeros((0, 4))
            cls_pred_scores = pred_scores[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros(0)
            cls_pred_RTs = pred_RTs[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros((0, 4, 4))
            cls_pred_scales = pred_scales[pred_class_ids==cls_id] if len(pred_class_ids) else np.zeros((0, 3))


            # calculate the overlap between each gt instance and pred instance
            if synset_names[cls_id] != 'mug':
                cls_gt_handle_visibility = np.ones_like(cls_gt_class_ids)
            else:
                cls_gt_handle_visibility = gt_handle_visibility[gt_class_ids==cls_id] if len(gt_class_ids) else np.ones(0)


            iou_cls_gt_match, iou_cls_pred_match, _, iou_pred_indices = compute_3d_matches(cls_gt_class_ids, cls_gt_RTs, cls_gt_scales, cls_gt_handle_visibility, synset_names,
                                                                                           cls_pred_bboxes, cls_pred_class_ids, cls_pred_scores, cls_pred_RTs, cls_pred_scales,
                                                                                           iou_thres_list)
            if len(iou_pred_indices):
                cls_pred_class_ids = cls_pred_class_ids[iou_pred_indices]
                cls_pred_RTs = cls_pred_RTs[iou_pred_indices]
                cls_pred_scores = cls_pred_scores[iou_pred_indices]
                cls_pred_bboxes = cls_pred_bboxes[iou_pred_indices]


            iou_pred_matches_all[cls_id] = np.concatenate((iou_pred_matches_all[cls_id], iou_cls_pred_match), axis=-1)
            cls_pred_scores_tile = np.tile(cls_pred_scores, (num_iou_thres, 1))
            iou_pred_scores_all[cls_id] = np.concatenate((iou_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)
            assert iou_pred_matches_all[cls_id].shape[1] == iou_pred_scores_all[cls_id].shape[1]
            iou_gt_matches_all[cls_id] = np.concatenate((iou_gt_matches_all[cls_id], iou_cls_gt_match), axis=-1)

            if use_matches_for_pose:
                thres_ind = list(iou_thres_list).index(iou_pose_thres)

                iou_thres_pred_match = iou_cls_pred_match[thres_ind, :]

                
                cls_pred_class_ids = cls_pred_class_ids[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros(0)
                cls_pred_RTs = cls_pred_RTs[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros((0, 4, 4))
                cls_pred_scores = cls_pred_scores[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros(0)
                cls_pred_bboxes = cls_pred_bboxes[iou_thres_pred_match > -1] if len(iou_thres_pred_match) > 0 else np.zeros((0, 4))


                iou_thres_gt_match = iou_cls_gt_match[thres_ind, :]
                cls_gt_class_ids = cls_gt_class_ids[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros(0)
                cls_gt_RTs = cls_gt_RTs[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros((0, 4, 4))
                cls_gt_handle_visibility = cls_gt_handle_visibility[iou_thres_gt_match > -1] if len(iou_thres_gt_match) > 0 else np.zeros(0)



            RT_projection_2d_overlaps = compute_RT_projection_2d_overlaps(cls_gt_class_ids, cls_gt_RTs, cls_gt_handle_visibility,
                                                                          cls_pred_class_ids, cls_pred_RTs,
                                                                          meshes, intrinsics, synset_names)


            projection_cls_gt_match, projection_cls_pred_match = compute_match_from_projection_2d_dist(RT_projection_2d_overlaps, 
                                                                                                       cls_pred_class_ids, 
                                                                                                       cls_gt_class_ids, 
                                                                                                       projection_thres_list)
                            

            projection_pred_matches_all[cls_id] = np.concatenate((projection_pred_matches_all[cls_id], projection_cls_pred_match), axis=-1)
            
            cls_pred_scores_tile = np.tile(cls_pred_scores, (num_projection_thres, 1))
            projection_pred_scores_all[cls_id]  = np.concatenate((projection_pred_scores_all[cls_id], cls_pred_scores_tile), axis=-1)
            assert projection_pred_scores_all[cls_id].shape[1] == projection_pred_matches_all[cls_id].shape[1], '{} vs. {}'.format(projection_pred_scores_all[cls_id].shape, projection_pred_matches_all[cls_id].shape)
            projection_gt_matches_all[cls_id] = np.concatenate((projection_gt_matches_all[cls_id], projection_cls_gt_match), axis=-1)

    
    
    # draw iou 3d AP vs. iou thresholds
    fig_iou = plt.figure()
    ax_iou = plt.subplot(111)
    plt.ylabel('AP')
    plt.ylim((0, 1))
    plt.xlabel('3D IoU thresholds')
    iou_output_path = os.path.join(log_dir, 'IoU_3D_AP_{}-{}.png'.format(iou_thres_list[0], iou_thres_list[-1]))
    iou_dict_pkl_path = os.path.join(log_dir, 'IoU_3D_AP_{}-{}.pkl'.format(iou_thres_list[0], iou_thres_list[-1]))

    iou_dict = {}
    iou_dict['thres_list'] = iou_thres_list
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]
        print(class_name)
        for s, iou_thres in enumerate(iou_thres_list):
            iou_3d_aps[cls_id, s] = compute_ap_from_matches_scores(iou_pred_matches_all[cls_id][s, :],
                                                                   iou_pred_scores_all[cls_id][s, :],
                                                                   iou_gt_matches_all[cls_id][s, :])    
        ax_iou.plot(iou_thres_list, iou_3d_aps[cls_id, :], label=class_name)
        
    iou_3d_aps[-1, :] = np.mean(iou_3d_aps[1:-1, :], axis=0)
    ax_iou.plot(iou_thres_list, iou_3d_aps[-1, :], label='mean')
    ax_iou.legend()
    fig_iou.savefig(iou_output_path)
    plt.close(fig_iou)

    iou_dict['aps'] = iou_3d_aps
    with open(iou_dict_pkl_path, 'wb') as f:
        cPickle.dump(iou_dict, f)
    

    # draw pose AP vs. thresholds
    if use_matches_for_pose:
        prefix='Pose_Only_'
    else:
        prefix='Pose_Detection_'


    LINE_WIDTH=3.0
    AXES_TITLE_SIZE=28
    XTICK_SIZE=28
    YTICK_SIZE=28
    LEGEND_SIZE=24
    TITLE_SIZE=32

    # plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
    plt.rc('axes', titlesize=TITLE_SIZE)     # fontsize of the axes title
    plt.rc('axes', labelsize=AXES_TITLE_SIZE)    # fontsize of the x and y labels
    plt.rc('xtick', labelsize=XTICK_SIZE)    # fontsize of the tick labels
    plt.rc('ytick', labelsize=YTICK_SIZE)    # fontsize of the tick labels
    plt.rc('legend', fontsize=LEGEND_SIZE)    # legend fontsize
    plt.rc('figure', titlesize=TITLE_SIZE)  # fontsize of the figure title

    plt.rcParams['font.family'] = 'Bahnschrift'


    projection_dict_pkl_path = os.path.join(log_dir, prefix+'Projection_2D_AP_{}-{}pixel.pkl'.format(projection_thres_list[0], projection_thres_list[-1]))
    pose_dict = {}
    pose_dict['projection_thres_list'] = projection_thres_list

    for i, projection_thres in enumerate(projection_thres_list):                
        for cls_id in range(1, num_classes):
            projection_2d_aps[cls_id, i] = compute_ap_from_matches_scores(projection_pred_matches_all[cls_id][i, :], 
                                                                             projection_pred_scores_all[cls_id][i, :], 
                                                                             projection_gt_matches_all[cls_id][i, :])


        projection_2d_aps[-1, i] = np.mean(projection_2d_aps[1:-1, i])
    
    pose_dict['aps'] = projection_2d_aps
    with open(projection_dict_pkl_path, 'wb') as f:
        cPickle.dump(pose_dict, f)

    plot_data_dir = 'data/linemod/plots'
    f, axarr = plt.subplots(2, 4, figsize=(20, 8), dpi=50)
    for i, cls_id in enumerate([1, 3, 4, 5, 6, 7, 8, 9]):
        
        class_name = synset_names[cls_id]

        axarr[i].set_title(class_name)
        axarr[i].set(xlabel='Pixel threshold', ylabel='AP %')
        axarr[i].set_xlim([0, 25])
        axarr[i].set_ylim([0, 101])
        axarr[i].plot(projection_thres_list, projection_2d_aps[cls_id, :]*100, label='Ours', linewidth=LINE_WIDTH)

        others_data_dir = plot_data_dir + '/' + class_name
        curve_data = glob.glob(others_data_dir, '*.txt')
        for curve_path in curve_data:
            curve_name = os.path.basename(curve_path)[:-4]
            if curve_name == 'bb8':
                curve_name = 'BB8'
            elif curve_name == 'posecnn':
                curve_name = 'PoseCNN'
            elif curve_name == 'tekin':
                curve_name = 'Tekin et al.'
            else:
                assert False, curve_name
            curve = np.loadtxt(curve_path)
            axarr[i].plot(curve[:, 0], curve[:, 1], label=curve_name, linewidth=LINE_WIDTH)

        axarr[i].legend(loc='lower right')



    for ax in axarr.flat:
        ax.label_outer()
        ax.set(ylim=(0, 101))
        ax.grid(True)       


    print(output_dir)
    plt.savefig(os.path.join(output_dir, 'result.png'))
    print('saved.')

    fig_pose = plt.figure()
    ax_pose = plt.subplot(111)
    plt.ylabel('AP')
    plt.ylim((0, 1.02))
    plt.xlim((projection_thres_list[0], projection_thres_list[-1]))
    plt.xlabel('Reprojection error thresholds/pixel')
    
    for cls_id in range(1, num_classes):
        class_name = synset_names[cls_id]
        print(class_name)
        ax_pose.plot(projection_thres_list, projection_2d_aps[cls_id, :], label=class_name)
    ax_pose.plot(projection_thres_list, projection_2d_aps[-1, :], label='mean')
    
    output_path = os.path.join(log_dir, prefix+'Projection_2D_AP_{}-{}pixel.png'.format(projection_thres_list[0], projection_thres_list[-1]))
    fig_pose.legend()
    fig_pose.savefig(output_path)
    plt.close(fig_pose)


    print('3D IoU at 25: {:.1f}'.format(iou_3d_aps[-1, iou_thres_list.index(0.25)] * 100))
    print('3D IoU at 50: {:.1f}'.format(iou_3d_aps[-1, iou_thres_list.index(0.5)] * 100))
    print('2D Projection at 2.5 pixel: {:.1f}'.format(projection_2d_aps[-1, projection_thres_list.index(2.5)] * 100))
    print('2D Projection at 5 pixel: {:.1f}'.format(projection_2d_aps[-1, projection_thres_list.index(5)] * 100))
    print('2D Projection at 10 pixel: {:.1f}'.format(projection_2d_aps[-1, projection_thres_list.index(10)] * 100))
    print('2D Projection at 15 pixel: {:.1f}'.format(projection_2d_aps[-1, projection_thres_list.index(15)] * 100))
    print('2D Projection at 20 pixel: {:.1f}'.format(projection_2d_aps[-1, projection_thres_list.index(20)] * 100))


    return iou_3d_aps, projection_2d_aps


def compute_coords_aps(final_results, synset_names, iou_thresholds, coord_thresholds):
    """Compute Average Precision at a set IoU threshold (default 0.5).
    Returns:
    mAP: Mean Average Precision
    precisions: List of precisions at different class score thresholds.
    recalls: List of recall values at different class score thresholds.
    overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    
    num_classes = len(synset_names)
    num_iou_thres = len(iou_thresholds)
    num_coord_thres = len(coord_thresholds)
    
    mean_coord_dist_cls = {}
    
#     pred_matches_all = {}
#     pred_scores_all = {}
#     gt_matches_all = {}

    
    for cls_id in range(1, num_classes):
        mean_coord_dist_cls[cls_id] = []
#         pred_matches_all[cls_id] = [[[] for _ in range(num_shift_thres)] for _ in range(num_degree_thres)]
#         gt_matches_all[cls_id] = [[[] for _ in range(num_shift_thres)] for _ in range(num_degree_thres)]
#         pred_scores_all[cls_id] = [[[] for _ in range(num_shift_thres)] for _ in range(num_degree_thres)]

    progress = 0
    for progress, result in enumerate(final_results):
        print(progress, len(final_results))
        
        gt_class_ids = result['gt_class_ids'].astype(np.int32)
        gt_bboxes = np.array(result['gt_bboxes'])
        gt_masks = np.array(result['gt_masks'])
        gt_coords = np.array(result['gt_coords'])
        
        #print(gt_class_ids.shape, gt_bboxes.shape, gt_masks.shape, gt_coords.shape)
        #gt_scales = result['gt_scale']
        
        pred_class_ids = result['pred_class_ids'].astype(np.int32)
        pred_bboxes = np.array(result['pred_bboxes'])
        pred_masks = np.array(result['pred_masks'])
        pred_coords = np.array(result['pred_coords'])
        pred_scores = result['pred_scores']
        #print(pred_class_ids.shape, pred_bboxes.shape, pred_masks.shape, pred_coords.shape)
        

        
        
        for cls_id in range(1, num_classes):
            # get gt and predictions in this class
            if(len(gt_class_ids)):
                cls_gt_class_ids = gt_class_ids[gt_class_ids==cls_id] 
                cls_gt_bboxes = gt_bboxes[gt_class_ids==cls_id] 
                cls_gt_masks = gt_masks[..., gt_class_ids==cls_id]
                cls_gt_coords = gt_coords[..., gt_class_ids==cls_id, :]
            else:
                cls_gt_class_ids = []
                cls_gt_bboxes = []
                cls_gt_masks = []
                cls_gt_coords = []

            if(len(pred_class_ids)):
                cls_pred_class_ids = pred_class_ids[pred_class_ids==cls_id] 
                cls_pred_bboxes =  pred_bboxes[pred_class_ids==cls_id] 
                cls_pred_scores = pred_scores[pred_class_ids==cls_id]
                cls_pred_masks = pred_masks[:, :, pred_class_ids==cls_id]
                cls_pred_coords = pred_coords[:, :, pred_class_ids==cls_id, :]
            else:
                cls_pred_class_ids = []
                cls_pred_bboxes = []
                cls_pred_scores = []
                cls_pred_masks = []
                cls_pred_coords = []
            
            # calculate the overlap between each gt instance and pred instance
            
            gt_match, pred_match, overlaps, pred_indices = compute_matches(cls_gt_bboxes, cls_gt_class_ids, cls_gt_masks,
                                                             cls_pred_bboxes, cls_pred_class_ids, cls_pred_scores, 
                                                             cls_pred_masks, 0.5)
            
            if len(gt_match) and len(pred_match):
            
                cls_pred_masks_sorted = cls_pred_masks[..., pred_indices]
                cls_pred_coords_sorted = cls_pred_coords[..., pred_indices, :]

                for i in range(len(pred_match)):
                    if pred_match[i] > -1:
                        j = int(pred_match[i])
                        mean_coord_dist = compute_mean_l1_coord_diff(cls_pred_masks_sorted[..., i], 
                                                                     cls_gt_masks[..., j],
                                                                     cls_pred_coords_sorted[..., i, :], 
                                                                     cls_gt_coords[..., j, :],
                                                                     synset_names,
                                                                     cls_id)
                            #print(mean_coord_dist)

                        mean_coord_dist_cls[cls_id].append(mean_coord_dist)
                            #print(mean_coord_dist_cls[cls_id])

            

    for cls_id in range(1, num_classes):
        mean_coord_dist_cls[cls_id] = np.array(mean_coord_dist_cls[cls_id])
        print('mean coord dist of {} class: {}'.format(synset_names[cls_id], np.mean(mean_coord_dist_cls[cls_id])))


def compute_3d_matches_degree_cm(gt_class_ids, gt_RTs, gt_scales, 
                                 pred_boxes, pred_class_ids, pred_scores, pred_RTs, 
                                 synset_names,
                                 thres_shift = 5, 
                                 thres_degree=5, 
                                 bound_threshold=5):
    """Finds matches between prediction and ground truth instances.
    Returns:
        gt_match: 1-D array. For each GT box it has the index of the matched
                  predicted box.
        pred_match: 1-D array. For each predicted box, it has the index of
                    the matched ground truth box.
        overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # Trim zero padding
    # TODO: cleaner to do zero unpadding upstream
    if len(pred_boxes):
        pre_len = len(pred_boxes)
        pred_boxes = trim_zeros(pred_boxes)
        after_len = len(pred_boxes)
        assert pre_len == after_len
        pred_scores = pred_scores[:pred_boxes.shape[0]]
        # Sort predictions by score from high to low
        indices = np.argsort(pred_scores)[::-1]
        pred_boxes = pred_boxes[indices]
        pred_class_ids = pred_class_ids[indices]
        pred_scores = pred_scores[indices]

    # print('num of gt instances: {}, num of pred instances: {}'.format(len(gt_class_ids), len(gt_class_ids)))
    num_pred = len(pred_boxes)
    num_gt = len(gt_class_ids)

    # Compute IoU overlaps [pred_bboxs gt_bboxs]
    #overlaps = [[0 for j in range(num_gt)] for i in range(num_pred)]
    overlaps = np.zeros((num_pred, num_gt), dtype=np.float32)

    for i in range(num_pred):
        for j in range(num_gt):
            overlaps[i, j] = compute_RT_upper_bound_symmetry(pred_RTs[i], 
                                                             gt_RTs[j], 
                                                             gt_class_ids[j], 
                                                             synset_names, 
                                                             thres_shift=thres_shift, 
                                                             thres_theta=thres_degree)

    # Loop through predictions and find matching ground truth boxes
    # multi-level: 5, 10, 15, 20
    pred_match = -1 * np.ones(num_pred)
    gt_match = -1 * np.ones(num_gt)

    for i in range(len(pred_boxes)):
        # Find best matching ground truth box
        # 1. Sort matches by bound from low to high
        sorted_ixs = np.argsort(overlaps[i])
        # 2. Remove low scores
        # low_score_idx = np.where(overlaps[i, sorted_ixs] >= bound_threshold)[0]
        # if low_score_idx.size > 0:
        #     sorted_ixs = sorted_ixs[:low_score_idx[0]]
        # 3. Find the match
        for j in sorted_ixs:
            # If ground truth box is already matched, go to next one
            if gt_match[j] > -1 or pred_class_ids[i] != gt_class_ids[j]:
                continue
            # If we reach IoU smaller than the threshold, end the loop
            min_bound  = overlaps[i, j]
            if min_bound > 1:
                break
            elif min_bound ==1.0:
                gt_match[j] = i
                pred_match[i] = j
                break
            else:
                assert False, 'Invalid overlaps {}'.format(min_bound)

    return gt_match, pred_match, pred_class_ids, pred_scores


def compute_ap(gt_boxes, gt_class_ids, gt_masks,
               pred_boxes, pred_class_ids, pred_scores, pred_masks,
               iou_threshold=0.5):
    """Compute Average Precision at a set IoU threshold (default 0.5).
    Returns:
    mAP: Mean Average Precision
    precisions: List of precisions at different class score thresholds.
    recalls: List of recall values at different class score thresholds.
    overlaps: [pred_boxes, gt_boxes] IoU overlaps.
    """
    # Get matches and overlaps
    gt_match, pred_match, overlaps = compute_matches(
        gt_boxes, gt_class_ids, gt_masks,
        pred_boxes, pred_class_ids, pred_scores, pred_masks,
        iou_threshold)

    # Compute precision and recall at each prediction box step
    precisions = np.cumsum(pred_match > -1) / (np.arange(len(pred_match)) + 1) # TP / (TP + FP)
    recalls = np.cumsum(pred_match > -1).astype(np.float32) / len(gt_match) # TP / (TP + FN)

    # Pad with start and end values to simplify the math
    precisions = np.concatenate([[0], precisions, [0]])
    recalls = np.concatenate([[0], recalls, [1]])

    # Ensure precision values decrease but don't increase. This way, the
    # precision value at each recall threshold is the maximum it can be
    # for all following recall thresholds, as specified by the VOC paper.
    for i in range(len(precisions) - 2, -1, -1):
        precisions[i] = np.maximum(precisions[i], precisions[i + 1])

    # Compute mean AP over recall range
    indices = np.where(recalls[:-1] != recalls[1:])[0] + 1
    mAP = np.sum((recalls[indices] - recalls[indices - 1]) *
                 precisions[indices])

    return mAP, precisions, recalls, overlaps



def compute_ap_range(gt_box, gt_class_id, gt_mask,
                     pred_box, pred_class_id, pred_score, pred_mask,
                     iou_thresholds=None, verbose=1):
    """Compute AP over a range or IoU thresholds. Default range is 0.5-0.95."""
    # Default is 0.5 to 0.95 with increments of 0.05
    iou_thresholds = iou_thresholds or np.arange(0.5, 1.0, 0.05)

    # Compute AP over range of IoU thresholds
    AP = []
    for iou_threshold in iou_thresholds:
        ap, precisions, recalls, overlaps = \
            compute_ap(gt_box, gt_class_id, gt_mask,
                       pred_box, pred_class_id, pred_score, pred_mask,
                       iou_threshold=iou_threshold)
        if verbose:
            print("AP @{:.2f}:\t {:.3f}".format(iou_threshold, ap))
        AP.append(ap)
    AP = np.array(AP).mean()
    if verbose:
        print("AP @{:.2f}-{:.2f}:\t {:.3f}".format(
            iou_thresholds[0], iou_thresholds[-1], AP))
    return AP


def compute_recall(pred_boxes, gt_boxes, iou):
    """Compute the recall at the given IoU threshold. It's an indication
    of how many GT boxes were found by the given prediction boxes.
    pred_boxes: [N, (y1, x1, y2, x2)] in image coordinates
    gt_boxes: [N, (y1, x1, y2, x2)] in image coordinates
    """
    # Measure overlaps
    overlaps = compute_overlaps(pred_boxes, gt_boxes)
    iou_max = np.max(overlaps, axis=1)
    iou_argmax = np.argmax(overlaps, axis=1)
    positive_ids = np.where(iou_max >= iou)[0]
    matched_gt_boxes = iou_argmax[positive_ids]

    recall = len(set(matched_gt_boxes)) / gt_boxes.shape[0]
    return recall, positive_ids


# ## Batch Slicing
# Some custom layers support a batch size of 1 only, and require a lot of work
# to support batches greater than 1. This function slices an input tensor
# across the batch dimension and feeds batches of size 1. Effectively,
# an easy way to support batches > 1 quickly with little code modification.
# In the long run, it's more efficient to modify the code to support large
# batches and getting rid of this function. Consider this a temporary solution
def batch_slice(inputs, graph_fn, batch_size, names=None):
    """Splits inputs into slices and feeds each slice to a copy of the given
    computation graph and then combines the results. It allows you to run a
    graph on a batch of inputs even if the graph is written to support one
    instance only.
    inputs: list of tensors. All must have the same first dimension length
    graph_fn: A function that returns a TF tensor that's part of a graph.
    batch_size: number of slices to divide the data into.
    names: If provided, assigns names to the resulting tensors.
    """
    if not isinstance(inputs, list):
        inputs = [inputs]

    outputs = []
    for i in range(batch_size):
        inputs_slice = [x[i] for x in inputs]
        output_slice = graph_fn(*inputs_slice)
        if not isinstance(output_slice, (tuple, list)):
            output_slice = [output_slice]
        outputs.append(output_slice)
    # Change outputs from a list of slices where each is
    # a list of outputs to a list of outputs and each has
    # a list of slices
    outputs = list(zip(*outputs))

    if names is None:
        names = [None] * len(outputs)

    result = [tf.stack(o, axis=0, name=n)
              for o, n in zip(outputs, names)]
    if len(result) == 1:
        result = result[0]

    return result


############################################################
#  Miscellaneous
############################################################
def load_mesh(mesh_path, is_save=False, is_normalized=False, is_flipped=False):
    with open(mesh_path, 'r') as f:
        lines = f.readlines()

    # mesh = []
    # for line in lines:
    #     if line[:2] == 'v ':
    #         vertex = line[2:].split(' ')
    #         vertex = vertex[:3]
    #         mesh += [[float(x) for x in vertex]]
    # mesh = np.array(mesh, dtype=np.float32)

    vertices = []
    faces = []
    for l in lines:
        l = l.strip()
        words = l.split(' ')
        if words[0] == 'v':
            vertices.append([float(words[1]), float(words[2]), float(words[3])])
        if words[0] == 'f':
            face_words = [x.split('/')[0] for x in words]
            faces.append([int(face_words[1])-1, int(face_words[2])-1, int(face_words[3])-1])


    vertices = np.array(vertices, dtype=np.float64)
    # flip mesh to unity rendering
    if is_flipped:
        vertices[:, 2] = -vertices[:, 2] 
    faces = np.array(faces, dtype=np.int32)
    
    if is_normalized:
        maxs = np.amax(vertices, axis=0)
        mins = np.amin(vertices, axis=0)
        diffs = maxs - mins
        assert diffs.shape[0] == 3
        vertices = vertices/np.linalg.norm(diffs)
    
    if is_save:
        np.savetxt(mesh_path.replace('.obj', '_vertices.txt'), X = vertices)

    return vertices, faces


def draw(img, imgpts, axes, color):
    imgpts = np.int32(imgpts).reshape(-1, 2)


    # draw ground layer in darker color
    color_ground = (int(color[0] * 0.3), int(color[1] * 0.3), int(color[2] * 0.3))
    for i, j in zip([4, 5, 6, 7],[5, 7, 4, 6]):
        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color_ground, 3)


    # draw pillars in blue color
    color_pillar = (int(color[0]*0.6), int(color[1]*0.6), int(color[2]*0.6))
    for i, j in zip(range(4),range(4,8)):
        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color_pillar, 3)

    
    # finally, draw top layer in color
    for i, j in zip([0, 1, 2, 3],[1, 3, 0, 2]):
        img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]), color, 3)


    # draw axes
    img = cv2.line(img, tuple(axes[0]), tuple(axes[1]), (0, 0, 255), 3) # BLUE
    img = cv2.line(img, tuple(axes[0]), tuple(axes[3]), (255, 0, 0), 3) # RED
    img = cv2.line(img, tuple(axes[0]), tuple(axes[2]), (0, 255, 0), 3) ## y last GREEN


    return img



def draw_coord_mask(image, class_ids, RTs, Vs, Fs, intrinsics, output_path, is_overlay=True):
    Rasterizer = cdll.LoadLibrary('./libRasterizer.so')
    # coords = np.copy(image).astype(np.float64)
    # image = cv2.imread(image_path)[:, :, :3]
    # image_1 = copy.deepcopy(image)
    # image_2 = copy.deepcopy(image[:, :, ::-1])

    # print(type(image_1), image_1.dtype)
    # print(type(image_2), image_2.dtype)

    # print(np.unique(np.equal(image_1, image_2)))

    if is_overlay:
        #coords = copy.deepcopy(image[:, :, (0, 1, 2)]).astype(np.float64)
        coords = copy.deepcopy(image[:,:,::-1]).astype(np.float64)
    else:
        coords = 255*np.ones_like(image, dtype=np.float64)
    masks = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    masks[:,:] = 255

    Rasterizer.InitializeZBuffer(image.shape[1], image.shape[0], image.shape[2])

    for ind, RT in enumerate(RTs):
        cls_id = class_ids[ind]
        vertices = Vs[ind].transpose()
        number_of_vertex = vertices.shape[1]
        vertices_homogeneous = np.vstack([vertices, np.ones((1, number_of_vertex))])

        transformed_vertices = RTs[ind] @ vertices_homogeneous
        transformed_vertices = transformed_vertices[:3, :] / transformed_vertices[3, :]
        projected_pixels = intrinsics @ transformed_vertices
        projected_pixels[2, :] = 1.0 / projected_pixels[2, :]
        projected_pixels[:2, :] = projected_pixels[:2, :] * projected_pixels[2, :]
        projected_pixels = projected_pixels.transpose().astype(np.float64)  # shape: N*2


        # for pixel in projected_pixels:
        #     c_x = int(pixel[0])
        #     c_y = int(pixel[1])
        #     cv2.circle(image, (c_x, c_y), 1, (0, 0, 255))

        # post processing
        coordinate = Vs[ind].copy()
        coordinate += 0.5 # Range 0-1
        coordinate[:, 2] = 1 - coordinate[:, 2] # left coordinate system in coord map

        # cv2 imwrite is in order of BGR
        coordinate = coordinate[:, (2, 1, 0)]
        coordinate = coordinate * 255

        F = Fs[ind];
        for j in range(F.shape[0]):
            #print(F[j, :])
            p0 = projected_pixels[F[j, 0]].copy()
            p1 = projected_pixels[F[j, 1]].copy()
            p2 = projected_pixels[F[j, 2]].copy()
            prop0 = coordinate[F[j, 0]].copy()
            prop1 = coordinate[F[j, 1]].copy()
            prop2 = coordinate[F[j, 2]].copy()

            Rasterizer.DrawTriangle(
                c_void_p(p0.ctypes.data),
                c_void_p(p1.ctypes.data),
                c_void_p(p2.ctypes.data),
                c_void_p(prop0.ctypes.data),
                c_void_p(prop1.ctypes.data),
                c_void_p(prop2.ctypes.data),
                c_void_p(coords.ctypes.data),
                ind + 1,
                c_void_p(masks.ctypes.data))

    mask_output_path = output_path+'_mask.png'
    cv2.imwrite(mask_output_path, masks)
    coord_output_path = output_path+'_coord.jpg'
    cv2.imwrite(coord_output_path, coords)


def draw_text(draw_image, bbox, text, draw_box=False):
    fontFace = cv2.FONT_HERSHEY_TRIPLEX
    fontScale = 1
    thickness = 1

    retval, baseline = cv2.getTextSize(text, fontFace, fontScale, thickness)
    
    bbox_margin = 10
    text_margin = 10
    
    text_box_pos_tl = (min(bbox[1] + bbox_margin, 635 - retval[0] - 2* text_margin) , min(bbox[2] + bbox_margin, 475 - retval[1] - 2* text_margin)) 
    text_box_pos_br = (text_box_pos_tl[0] + retval[0] + 2* text_margin,  text_box_pos_tl[1] + retval[1] + 2* text_margin)

    # text_pose is the bottom-left corner of the text
    text_pos = (text_box_pos_tl[0] + text_margin, text_box_pos_br[1] - text_margin - 3)
    
    if draw_box:
        cv2.rectangle(draw_image, 
                      (bbox[1], bbox[0]),
                      (bbox[3], bbox[2]),
                      (255, 0, 0), 2)

    cv2.rectangle(draw_image, 
                  text_box_pos_tl,
                  text_box_pos_br,
                  (255,0,0), -1)
    
    cv2.rectangle(draw_image, 
                  text_box_pos_tl,
                  text_box_pos_br,
                  (0,0,0), 1)

    cv2.putText(draw_image, text, text_pos,
                fontFace, fontScale, (255,255,255), thickness)

    return draw_image


def draw_detections(depth, image, save_dir, data_name, image_id, intrinsics, synset_names, draw_rgb_coord,
                    gt_bbox, gt_class_ids, gt_mask, gt_coord, gt_RTs, gt_scales, gt_handle_visibility,
                    pred_bbox, pred_class_ids, pred_mask, pred_coord, pred_RTs, pred_scores, pred_scales,
                    draw_gt=True, draw_pred=True, draw_tag=True):

    alpha = 0.5

    if draw_gt:
        output_path = os.path.join(save_dir, '{}_{}_coord_gt.png'.format(data_name, image_id))
        draw_image = image.copy()
        num_gt_instances = len(gt_class_ids)

        for i in range(num_gt_instances):
            mask = gt_mask[:, :, i]
            
            #mask = mask[:, :, np.newaxis]
            #mask = np.repeat(mask, 3, axis=-1)
            cind, rind = np.where(mask == 1)
            coord_data = gt_coord[:, :, i, :].copy()
            #coord_data[:, :, 2] = 1 - coord_data[:, :, 2] # undo the z axis flipping to match original data        
            draw_image[cind, rind] = coord_data[cind, rind] * 255
            
        if draw_tag:
            for i in range(num_gt_instances):
                overlay = draw_image.copy()
                overlay = draw_text(overlay, gt_bbox[i], synset_names[gt_class_ids[i]], draw_box=True)
                cv2.addWeighted(overlay, alpha, draw_image, 1 - alpha, 0, draw_image)
            
        # #if draw_tag:
        # for i in range(num_gt_instances):
        #     print('a', synset_names[gt_class_ids[i]])
        #     if synset_names[gt_class_ids[i]] == 'camera':
        #         overlay = draw_image.copy()
        #         cv2.rectangle(overlay, 
        #               (gt_bbox[i][1], gt_bbox[i][0]),
        #               (gt_bbox[i][3], gt_bbox[i][2]),
        #               (255, 0, 0), 2)

        #         cv2.addWeighted(overlay, alpha, draw_image, 1 - alpha, 0, draw_image)

        
        cv2.imwrite(output_path, draw_image[:, :, ::-1])

        output_path = os.path.join(save_dir, '{}_{}_bbox_gt.png'.format(data_name, image_id))
        draw_image_bbox = image.copy()

        if gt_RTs is not None:
            print("gt_RTS is not None.")
            for ind, RT in enumerate(gt_RTs):
                cls_id = gt_class_ids[ind]

                xyz_axis = 0.3*np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]]).transpose()
                transformed_axes = transform_coordinates_3d(xyz_axis, RT)
                projected_axes = calculate_2d_projections(transformed_axes, intrinsics)


                bbox_3d = get_3d_bbox(gt_scales[ind], 0)
                transformed_bbox_3d = transform_coordinates_3d(bbox_3d, RT)
                projected_bbox = calculate_2d_projections(transformed_bbox_3d, intrinsics)
                draw_image_bbox = draw(draw_image_bbox, projected_bbox, projected_axes, (255, 0, 0))

        cv2.imwrite(output_path, draw_image_bbox[:, :, ::-1])


    if draw_pred:

        # Should I draw human chair segmentation
        human_segm_flag = False

        print('a'*50)
        # Vs, Fs = dataset.load_objs(image_id, is_normalized=True) ## scale is estimated in RT
        output_path   = os.path.join(save_dir, '{}_{}_coord_pred.png'.format(data_name, image_id))
        output_path_r = os.path.join(save_dir, '{}_{}_coord_pred_r.png'.format(data_name, image_id))
        output_path_g = os.path.join(save_dir, '{}_{}_coord_pred_g.png'.format(data_name, image_id))
        output_path_b = os.path.join(save_dir, '{}_{}_coord_pred_b.png'.format(data_name, image_id))
        # utils.draw_coord_mask(image, r['class_ids'], pred_RTs, Vs, Fs, intrinsics, output_path)
        draw_image = image.copy()
        
        if draw_rgb_coord:
            r_image = image.copy()
            g_image = image.copy()
            b_image = image.copy()

        # Number of predictions
        num_pred_instances = len(pred_class_ids)    
        
        # Loop over them
        for i in range(num_pred_instances):
            
            # don't draw person or chair
            cls_id = pred_class_ids[i]
            if human_segm_flag == False and (cls_id == 4 or cls_id == 5):
                continue
            
            mask = pred_mask[:, :, i]
            #mask = mask[:, :, np.newaxis]
            #mask = np.repeat(mask, 3, axis=-1)
            cind, rind = np.where(mask == 1)
            coord_data = pred_coord[:, :, i, :].copy()
            #coord_data[:, :, 2] = 1 - coord_data[:, :, 2] # undo the z axis flipping to match original data
            draw_image[cind, rind] = coord_data[cind, rind] * 255
            if draw_rgb_coord:
                b_image[cind, rind, 2] = coord_data[cind, rind, 2] * 255
                b_image[cind, rind, 0:2] = 0

                g_image[cind, rind, 1] = coord_data[cind, rind, 1] * 255
                g_image[cind, rind, 0] = 0
                g_image[cind, rind, 2] = 0

                r_image[cind, rind, 0] = coord_data[cind, rind, 0] * 255
                r_image[cind, rind, 1:3] = 0

        if draw_tag:
            for i in range(num_pred_instances):
                
                cls_id = pred_class_ids[i]
                # don't draw person or chair
                if human_segm_flag == False and (cls_id == 4 or cls_id == 5):
                    continue

                overlay = draw_image.copy()
                text = synset_names[pred_class_ids[i]]+'({:.2f})'.format(pred_scores[i]) #og
                overlay = draw_text(overlay, pred_bbox[i], text, draw_box=True)
                cv2.addWeighted(overlay, alpha, draw_image, 1 - alpha, 0, draw_image)

        cv2.imwrite(output_path, draw_image[:, :, ::-1])


        if draw_rgb_coord:
            cv2.imwrite(output_path_r, r_image[:, :, ::-1])
            cv2.imwrite(output_path_g, g_image[:, :, ::-1])
            cv2.imwrite(output_path_b, b_image[:, :, ::-1])
                        
        
        output_path = os.path.join(save_dir, '{}_{}_bbox_pred.png'.format(data_name, image_id))
        draw_image_bbox = image.copy()

        if gt_class_ids is not None:
            gt_match, pred_match, _, pred_indices = compute_matches(gt_bbox, gt_class_ids, gt_mask,
                                                                    pred_bbox, pred_class_ids, pred_scores, pred_mask,
                                                                    0.5)

            if len(pred_indices):
                pred_class_ids = pred_class_ids[pred_indices]
                pred_scores = pred_scores[pred_indices]        
                pred_RTs = pred_RTs[pred_indices]

        
        for ind in range(num_pred_instances):
            RT = pred_RTs[ind]
            cls_id = pred_class_ids[ind]
            
            # don't draw person or chair
            if human_segm_flag == False and (cls_id == 4 or cls_id == 5):
                continue
            
            # if gt_class_ids is not None:## if gt exists, skip instances that fail to match
            #     gt_ind = int(pred_match[ind])
            #     if gt_ind == -1:
            #         continue

            ### DRAW THE 3 AXES
            xyz_axis = 0.3*np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]]).transpose()
            transformed_axes = transform_coordinates_3d(xyz_axis, RT)
            projected_axes = calculate_2d_projections(transformed_axes, intrinsics)


            ### DRAW THE BOUNDING BOX
            bbox_3d = get_3d_bbox(pred_scales[ind, :], 0)
            transformed_bbox_3d = transform_coordinates_3d(bbox_3d, RT)
            projected_bbox = calculate_2d_projections(transformed_bbox_3d, intrinsics)
            draw_image_bbox = draw(draw_image_bbox, projected_bbox, projected_axes, (255, 0, 0))

            ### WRITE (DRAW) THE DIMENSIONS X, Y, Z
            if draw_tag:
                xx,yy,zz = get_dimensions_from_3Dpoints(transformed_bbox_3d*1000)
                overlay = draw_image_bbox.copy()
                text = 'Dim:{:.1f},{:.1f}, {:.1f}, {})'.format(xx,yy,zz, synset_names[pred_class_ids[ind]])
                overlay = draw_text(overlay, pred_bbox[ind], text)
                cv2.addWeighted(overlay, alpha, draw_image_bbox, 1 - alpha, 0, draw_image_bbox)


        if draw_tag:
            if gt_class_ids is not None: ## if gt exists, draw rotation and translation error
                for ind in range(num_pred_instances):
                    gt_ind = int(pred_match[ind])
                    if gt_ind == -1:
                        continue

                    overlay = draw_image_bbox.copy()
                    RT = pred_RTs[ind]
                    gt_RT = gt_RTs[gt_ind]
                    cls_id = pred_class_ids[ind]
                    
                    # don't draw person or chair
                    if human_segm_flag == False and (cls_id == 4 or cls_id == 5):
                        continue

                    degree, cm = compute_RT_degree_cm_symmetry(RT, gt_RT, cls_id, synset_names)
                    text = '{}({:.1f}, {:.1f})'.format(synset_names[cls_id], degree, cm)
                    
                    ####
                    text = '{}'.format(synset_names[cls_id])
                    ####
                    overlay = draw_text(overlay, pred_bbox[ind], text)
                    cv2.addWeighted(overlay, alpha, draw_image_bbox, 1 - alpha, 0, draw_image_bbox)



        cv2.imwrite(output_path, draw_image_bbox[:, :, ::-1])

def get_dimensions_from_3Dpoints(points3D, verbose=False):

    # double number means that it is minus, e.g.: xx means minus x
    xyz = points3D.transpose()[0]
    xyzz = points3D.transpose()[1]
    xxyz = points3D.transpose()[2]
    xxyzz = points3D.transpose()[3]

    xyyz = points3D.transpose()[4]
    xyyzz = points3D.transpose()[5]
    xxyyz = points3D.transpose()[6]
    xxyyzz = points3D.transpose()[7]

    
    # All heights (y) bot means bottom, 
    height1 = math.sqrt( math.pow( abs( xyyz[0] - xyz[0] ), 2) + 
                         math.pow( abs( xyyz[1] - xyz[1] ), 2) + 
                         math.pow( abs( xyyz[2] - xyz[2] ), 2) ) 

    height2 = math.sqrt( math.pow( abs( xyzz[0] - xyyzz[0] ), 2) + 
                         math.pow( abs( xyzz[1] - xyyzz[1] ), 2) + 
                         math.pow( abs( xyzz[2] - xyyzz[2] ), 2) )

    height3 = math.sqrt( math.pow( abs( xxyz[0] - xxyyz[0] ), 2) + 
                         math.pow( abs( xxyz[1] - xxyyz[1] ), 2) + 
                         math.pow( abs( xxyz[2] - xxyyz[2] ), 2) )

    height4 = math.sqrt( math.pow( abs( xxyzz[0] - xxyyzz[0] ), 2) + 
                         math.pow( abs( xxyzz[1] - xxyyzz[1] ), 2) + 
                         math.pow( abs( xxyzz[2] - xxyyzz[2] ), 2) )
    

    # All widths (x)
    width1 = math.sqrt( math.pow( abs( xxyz[0] - xyz[0] ), 2) + 
                        math.pow( abs( xxyz[1] - xyz[1] ), 2) + 
                        math.pow( abs( xxyz[2] - xyz[2] ), 2) ) 
    
    width2 = math.sqrt( math.pow( abs( xxyzz[0] - xyzz[0] ), 2) + 
                        math.pow( abs( xxyzz[1] - xyzz[1] ), 2) + 
                        math.pow( abs( xxyzz[2] - xyzz[2] ), 2) )                         

    width3 = math.sqrt( math.pow( abs( xxyyz[0] - xyyz[0] ), 2) + 
                        math.pow( abs( xxyyz[1] - xyyz[1] ), 2) + 
                        math.pow( abs( xxyyz[2] - xyyz[2] ), 2) ) 

    width4 = math.sqrt( math.pow( abs( xxyyzz[0] - xyyzz[0] ), 2) + 
                        math.pow( abs( xxyyzz[1] - xyyzz[1] ), 2) + 
                        math.pow( abs( xxyyzz[2] - xyyzz[2] ), 2) )                         
    

    # All lengths (z)
    length1 = math.sqrt( math.pow( abs( xyzz[0] - xyz[0] ), 2) + 
                         math.pow( abs( xyzz[1] - xyz[1] ), 2) + 
                         math.pow( abs( xyzz[2] - xyz[2] ), 2) ) 

    length2 = math.sqrt( math.pow( abs( xxyzz[0] - xxyz[0] ), 2) + 
                         math.pow( abs( xxyzz[1] - xxyz[1] ), 2) + 
                         math.pow( abs( xxyzz[2] - xxyz[2] ), 2) )

    length3 = math.sqrt( math.pow( abs( xyyzz[0] - xyyz[0] ), 2) + 
                         math.pow( abs( xyyzz[1] - xyyz[1] ), 2) + 
                         math.pow( abs( xyyzz[2] - xyyz[2] ), 2) )

    length4 = math.sqrt( math.pow( abs( xxyyzz[0] - xxyyz[0] ), 2) + 
                         math.pow( abs( xxyyzz[1] - xxyyz[1] ), 2) + 
                         math.pow( abs( xxyyzz[2] - xxyyz[2] ), 2) )     


    if verbose:
        print("3D points:")
        print('( x, y, z) : {}'.format(points3D.transpose()[0]))
        print('( x, y,-z) : {}'.format(points3D.transpose()[1]))
        print('(-x, y, z) : {}'.format(points3D.transpose()[2]))
        print('(-x, y,-z) : {}'.format(points3D.transpose()[3]))

        print('( x,-y, z) : {}'.format(points3D.transpose()[4]))
        print('( x,-y,-z) : {}'.format(points3D.transpose()[5]))
        print('(-x,-y, z) : {}'.format(points3D.transpose()[6]))
        print('(-x,-y,-z) : {}'.format(points3D.transpose()[7]))

        print("heights: {:.2f} | {:.2f} | {:.2f} | {:.2f}".format(height1, 
                                                                height2, 
                                                                height3,
                                                                height4))
        
        print("widths: {:.2f} | {:.2f} | {:.2f} | {:.2f}".format(width1, 
                                                                width2, 
                                                                width3,
                                                                width4))

        print("lengths: {:.2f} | {:.2f} | {:.2f} | {:.2f}".format(length1, 
                                                                length2, 
                                                                length3,
                                                                length4))  
    
    # TODO: change these assertations to have one or two decimals accuracy
    # assert height1 == height2 == height3 == height4
    # assert width1 == width2 == width3 == width4
    # assert length1 == length2 == length3 == length4

    return width1, height1, length1


def draw_coco_detections(image, save_dir, data_name, image_id, synset_names, draw_rgb_coord,
                        gt_bbox, gt_class_ids, gt_mask, 
                        pred_bbox, pred_class_ids, pred_mask, pred_coord, pred_scores):

    output_path = os.path.join(save_dir, '{}_{:04d}_image.png'.format(data_name, image_id))
    cv2.imwrite(output_path, image[:, :, ::-1])

    
    output_path = os.path.join(save_dir, '{}_{:04d}_mask_gt.png'.format(data_name, image_id))
    draw_image = image.copy()
    for i in range(len(gt_class_ids)):
        mask = gt_mask[:, :, i]
        #mask = mask[:, :, np.newaxis]
        #mask = np.repeat(mask, 3, axis=-1)
        cind, rind = np.where(mask == 1)
        draw_image[cind, rind] = np.random.rand(3) * 255
        draw_image = draw_text(draw_image, gt_bbox[i], synset_names[gt_class_ids[i]], draw_box=True)
    cv2.imwrite(output_path, draw_image[:, :, ::-1])


    # Vs, Fs = dataset.load_objs(image_id, is_normalized=True) ## scale is estimated in RT
    output_path   = os.path.join(save_dir, '{}_{:04d}_coord_pred.png'.format(data_name, image_id))
    output_path_r = os.path.join(save_dir, '{}_{:04d}_coord_pred_r.png'.format(data_name, image_id))
    output_path_g = os.path.join(save_dir, '{}_{:04d}_coord_pred_g.png'.format(data_name, image_id))
    output_path_b = os.path.join(save_dir, '{}_{:04d}_coord_pred_b.png'.format(data_name, image_id))
    # utils.draw_coord_mask(image, r['class_ids'], pred_RTs, Vs, Fs, intrinsics, output_path)
    draw_image = image.copy()
    if draw_rgb_coord:
        r_image = image.copy()
        g_image = image.copy()
        b_image = image.copy()
    
    for i in range(len(pred_class_ids)):
        
        mask = pred_mask[:, :, i]
        #mask = mask[:, :, np.newaxis]
        #mask = np.repeat(mask, 3, axis=-1)
        cind, rind = np.where(mask == 1)
        coord_data = pred_coord[:, :, i, :].copy()
        #coord_data[:, :, 2] = 1 - coord_data[:, :, 2] # undo the z axis flipping to match original data
        draw_image[cind, rind] = coord_data[cind, rind] * 255
        draw_image = draw_text(draw_image, pred_bbox[i], synset_names[pred_class_ids[i]], draw_box=True)
        if draw_rgb_coord:
            b_image[cind, rind, 2] = coord_data[cind, rind, 2] * 255
            b_image[cind, rind, 0:2] = 0

            g_image[cind, rind, 1] = coord_data[cind, rind, 1] * 255
            g_image[cind, rind, 0] = 0
            g_image[cind, rind, 2] = 0

            r_image[cind, rind, 0] = coord_data[cind, rind, 0] * 255
            r_image[cind, rind, 1:3] = 0

    cv2.imwrite(output_path, draw_image[:, :, ::-1])
    if draw_rgb_coord:
        cv2.imwrite(output_path_r, r_image[:, :, ::-1])
        cv2.imwrite(output_path_g, g_image[:, :, ::-1])
        cv2.imwrite(output_path_b, b_image[:, :, ::-1])
                    
    
    

def backproject(depth, intrinsics, instance_mask):
    """Backprojecting the object points: 2D pixels + depth + intrinsics --> 3D coordinates
    
    We apply the instance mask to the 2D depth image, because we only want the depth points of the object of interest.
    We then backproject these 2D pixels (U,V) to 3D coordinates (X,Y,Z) using the camera intrinsics and depth (Z). 
    """

    # Compute the (multiplicative) inverse of a matrix.
    intrinsics_inv = np.linalg.inv(intrinsics)

    # Get shape of the depth image
    image_shape = depth.shape
    width = image_shape[1]
    height = image_shape[0]

    # Returns evenly spaced values, default = 1. This case: return x = [0,1,2,...,width]
    x = np.arange(width) 
    y = np.arange(height)

    # Get binary mask where values are positive if both depth and instance mask are 1
    #non_zero_mask = np.logical_and(depth > 0, depth < 5000)
    non_zero_mask = (depth > 0)
    final_instance_mask = np.logical_and(instance_mask, non_zero_mask) 
    
    # Get all coordinates of this mask where values are positive
    idxs = np.where(final_instance_mask)
    grid = np.array([idxs[1], idxs[0]]) # Shape = (2,N) where N is number of points

    # shape: height * width
    # mesh_grid = np.meshgrid(x, y) #[height, width, 2]
    # mesh_grid = np.reshape(mesh_grid, [2, -1])

    # Add a Z-coordinate (all 1s)
    N = grid.shape[1]
    ones = np.ones([1, N])
    uv_grid = np.concatenate((grid, ones), axis=0) # Shape = (3,N) where N is number of points

    # 
    xyz = intrinsics_inv @ uv_grid # (3, N)
    xyz = np.transpose(xyz) # (N, 3)

    z = depth[idxs[0], idxs[1]] 

    # print(np.amax(z), np.amin(z))
    pts = xyz * z[:, np.newaxis] / xyz[:, -1:]
    
    # To OpenGL
    # pts[:, 2] = -pts[:, 2]
    # pts[:, 1] = -pts[:, 1]

    # To Xavier's weird system
    pts[:, 0] = -pts[:, 0]
    pts[:, 1] = -pts[:, 1]

    return pts, idxs #x: 3d points, 2d pixels


def align(img, class_ids, masks, coords, depth, intrinsics, synset_names, image_path, save_path=None, if_norm=False, with_scale=True, verbose=True):
    """Xavier: uses coord+depth to predict RTs
    """
    num_instances = len(class_ids)
    error_messages = ''
    elapses = []
    if num_instances == 0:
        return np.zeros((0, 4, 4)), np.ones((0, 3)), error_messages, elapses

    RTs = np.zeros((num_instances, 4, 4))
    bbox_scales = np.ones((num_instances, 3))
    
    for i in range(num_instances):
        
        # get class string
        class_name = synset_names[class_ids[i]]
        
        # only align predictions for our classes
        if class_name not in ["box", "non-stem", "stem"]:
            elapses.append(0)
            RTs[i,:,:] = np.zeros([4,4], dtype=np.float32)
            bbox_scales[i, :] = [0,0,0]
        else:

            # get class number
            class_id = class_ids[i]

            # get corresponding mask and coords
            mask = masks[:, :, i]
            coord = coords[:, :, i, :]

            # Calculate the absolute value element-wise.
            abs_coord_pts = np.abs(coord[mask==1] - 0.5) 

            # Return the maximum of an array or maximum along an axis.
            bbox_scales[i, :] = 2*np.amax(abs_coord_pts, axis=0) 

            #Get 3d points and 2d pixels
            pts, idxs = backproject(depth, intrinsics, mask)
            coord_pts = coord[idxs[0], idxs[1], :] - 0.5
            
            # Show depth points in Open3D
            if False:
                # Only Object Depth
                pcl = o3d.geometry.PointCloud()
                pcl.points = o3d.utility.Vector3dVector(pts)
                print(img[idxs])
                print(img[idxs].shape)
                pcl.colors = o3d.utility.Vector3dVector(img[idxs]/255.0) # color it
                o3d.visualization.draw_geometries([pcl])

            if if_norm:
                scale = np.linalg.norm(bbox_scales[i, :])
                bbox_scales[i, :] /= scale
                coord_pts /= scale

            try:
                start = time.time()
                
                scale_factors, rotation, translation, outtransform = estimateSimilarityTransform(coord_pts, pts, False)

                aligned_RT = np.zeros((4, 4), dtype=np.float32) 
                if with_scale:
                    aligned_RT[:3, :3] = np.diag(scale_factors) / 1000 @ rotation.transpose() # @ = matrix multiplication
                else:
                    aligned_RT[:3, :3] = rotation.transpose()
                aligned_RT[:3, 3] = translation / 1000
                aligned_RT[3, 3] = 1
                
                if save_path is not None:
                    coord_pts_rotated = aligned_RT[:3, :3] @ coord_pts.transpose() + aligned_RT[:3, 3:]
                    coord_pts_rotated = coord_pts_rotated.transpose()
                    np.savetxt(save_path+'_{}_{}_depth_pts.txt'.format(i, class_name), pts)
                    np.savetxt(save_path+'_{}_{}_coord_pts.txt'.format(i, class_name), coord_pts)
                    np.savetxt(save_path+'_{}_{}_coord_pts_aligned.txt'.format(i, class_name), coord_pts_rotated)

                if verbose:
                    print('Mask ID: ', i)
                    print('Scale: ', scale_factors/1000)
                    print('Rotation: ', rotation.transpose())
                    print('Translation: ', translation/1000)

                elapsed = time.time() - start
                #print('elapsed: ', elapsed)
                elapses.append(elapsed)
            

            except Exception as e:
                message = '[ Error ] aligning instance {} in {} fails. Message: {}.'.format(synset_names[class_id], image_path, str(e))
                print(message)
                error_messages += message + '\n'
                aligned_RT = np.identity(4, dtype=np.float32) 

            # print('Estimation takes {:03f}s.'.format(time.time() - start))
            
            # from camera world to computer vision frame
            z_180_RT = np.zeros((4, 4), dtype=np.float32)
            z_180_RT[:3, :3] = np.diag([-1, -1, 1])
            z_180_RT[3, 3] = 1

            RTs[i, :, :] = z_180_RT @ aligned_RT 

    return RTs, bbox_scales, error_messages, elapses


def align_ICP(class_ids, masks, depth, intrinsics, synset_names, image_path, save_path=None, if_norm=False, verbose=False):
    num_instances = len(class_ids)
    error_messages = ''

    print(intrinsics)
    if num_instances == 0:
        return np.zeros((0, 4, 4)), np.ones((0, 3)), error_messages

    RTs = np.zeros((num_instances, 4, 4))
    bbox_scales = np.ones((num_instances, 3))

    def rotation_y_matrix(theta):
            rotation_matrix =  \
                    np.array([ np.cos(theta), 0,  np.sin(theta),
                                 0,           1,  0,
                              -np.sin(theta), 0,  np.cos(theta)])
            rotation_matrix = np.reshape(rotation_matrix, (3, 3))
            return rotation_matrix

    for i in range(num_instances):
        #class_name = synset_names[class_ids[i]]
        class_id = class_ids[i]
        mask = masks[:, :, i]
        # coord = coords[:, :, i, :]
        # abs_coord_pts = np.abs(coord[mask==1] - 0.5)
        # bbox_scales[i, :] = 2*np.amax(abs_coord_pts, axis=0)

        depth_pts, idxs = backproject(depth, intrinsics, mask)
        # coord_pts = coord[idxs[0], idxs[1], :] - 0.5

        # if if_norm:
        #     scale = np.linalg.norm(bbox_scales[i, :])
        #     bbox_scales[i, :] /= scale
        #     coord_pts /= scale
        
        glob_paths = '/home/hewang/Projects/CoordRCNN/data/pts/real_test/{}*.txt'.format(synset_names[class_ids[i]])
        print(glob_paths)
        shape_paths = glob.glob(glob_paths)
        print(shape_paths)


        random.shuffle(shape_paths)
        shape_pts = np.loadtxt(shape_paths[0])
        print(shape_pts.shape)
        bbox_scales[i, :] = np.array([1, 1, 1])
        


        # depth_pts = rotation_y_matrix(0.01)@shape_pts.transpose() + 0.01
        # depth_pts = depth_pts.transpose()

        #try:
        if True:
            scales, rotation, translation = ICP.doICP(shape_pts,
                                                      depth_pts, 
                                                      threshold=5, isViz=False)

            scales = np.amax(shape_pts, axis=0) - np.amin(shape_pts, axis=0) 
            # scales, rotation, translation, outtransform = estimateSimilarityTransform(coord_pts, pts, False)
            #scales = np.diag(bbox_scales[i, :])
            aligned_RT = np.zeros((4, 4), dtype=np.float32) 
            aligned_RT[:3, :3] = np.diag(scales)@rotation#.transpose()
            aligned_RT[:3, 3] = translation / 1000
            aligned_RT[3, 3] = 1
            
            #bbox_scales[i, :] = scales@bbox_scales[i, :]/1000

            if save_path is not None:
                coord_pts_rotated = aligned_RT[:3, :3] @ coord_pts.transpose() + aligned_RT[:3, 3:]
                coord_pts_rotated = coord_pts_rotated.transpose()
                np.savetxt(save_path+'_{}_{}_depth_pts.txt'.format(i, class_name), pts)
                np.savetxt(save_path+'_{}_{}_coord_pts.txt'.format(i, class_name), coord_pts)
                np.savetxt(save_path+'_{}_{}_coord_pts_aligned.txt'.format(i, class_name), coord_pts_rotated)

            if verbose:
                print('Mask ID: ', i)
                print('Scale: ', scales)
                print('Rotation: ', rotation)
                print('Translation: ', translation)
        

        # except Exception as e:
        #     message = '[ Error ] aligning instance {} in {} fails. Message: {}.'.format(synset_names[class_id], image_path, str(e))
        #     print(message)
        #     error_messages += message + '\n'
        #     aligned_RT = np.identity(4, dtype=np.float32) 

        # print('Estimation takes {:03f}s.'.format(time.time() - start))
        # from camera world to computer vision frame
        z_180_RT = np.zeros((4, 4), dtype=np.float32)
        z_180_RT[:3, :3] = np.diag([-1, -1, 1])
        z_180_RT[3, 3] = 1

        RTs[i, :, :] = z_180_RT @ aligned_RT 

    return RTs, bbox_scales, error_messages

def nms(bounding_boxes, confidence_scores, classIDs, maskz, coordz, threshold, verbose=False):
    
    # If no bounding boxes, return empty list
    if len(bounding_boxes) == 0:
        return [], []

    # Bounding boxes
    boxes = np.array(bounding_boxes)

    # coordinates of bounding boxes
    start_y = boxes[:, 0]
    start_x = boxes[:, 1]
    end_y = boxes[:, 2]
    end_x = boxes[:, 3]

    # Confidence scores of bounding boxes
    score = np.array(confidence_scores)

    # Picked bounding boxes
    picked_indices = []

    # Compute areas of bounding boxes
    areas = (end_x - start_x + 1) * (end_y - start_y + 1)

    # Sort by confidence score of bounding boxes
    order = np.argsort(score)

    # Iterate bounding boxes
    while order.size > 0:
        # The index of largest confidence score
        index = order[-1]

        # Pick the bounding box with largest confidence score
        picked_indices.append(index)

        # Compute ordinates of intersection-over-union(IOU)
        x1 = np.maximum(start_x[index], start_x[order[:-1]])
        x2 = np.minimum(end_x[index], end_x[order[:-1]])
        y1 = np.maximum(start_y[index], start_y[order[:-1]])
        y2 = np.minimum(end_y[index], end_y[order[:-1]])

        # Compute areas of intersection-over-union
        w = np.maximum(0.0, x2 - x1 + 1)
        h = np.maximum(0.0, y2 - y1 + 1)
        intersection = w * h

        # Compute the ratio between intersection and union
        ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)

        left = np.where(ratio < threshold)
        order = order[left]

    allindices = np.arange(0,len(classIDs))

    indices2delete = np.setdiff1d(allindices,picked_indices)
    
    if verbose:
        print('allindices: {} | indices2keep: {} | indices2delete: {}'.format(allindices, picked_indices, indices2delete))

    return indices2delete

def nocs2scales(class_ids, masks, coords):
    """Extracts dimensions from a NOCS map"""
    num_instances = len(class_ids)
    bbox_scales = np.ones((num_instances, 3))

    # Loop over the predicted instances in the scene
    for i in range(num_instances):

        coord = coords[:, :, i, :]
        mask = masks[:, :, i]

        # Recover the object dimensions in the NOCSpace
        NOCS_points = coord[mask == 1] - 0.5
        abs_coord_pts = np.abs(NOCS_points)
        bbox_scales[i, :] = 2*np.amax(abs_coord_pts, axis=0)

    return bbox_scales

def EfficientPnP(depth, class_ids, masks, coords, intrinsics, synset_names, with_scale=True):
    """Runs EPnP on NOCS maps, to return Rotation/Translations"""
    
    # Number of predicted instances
    num_instances = len(class_ids)
    
    # Initialize output placeholders
    error_messages = ''
    elapses = []
    if num_instances == 0:
        return np.zeros((0, 4, 4)), np.ones((0, 3)), np.ones((0, 3)), np.ones((0, 3)), error_messages, elapses
    RTs = np.zeros((num_instances, 4, 4))
    bbox_scales = np.ones((num_instances, 3))
    
    result_rvecs = []
    result_tvecs = []

    # Loop over the predicted instances in the scene
    for i in range(num_instances):
        
        # Extract class information
        class_name = synset_names[class_ids[i]]
        class_id = class_ids[i]

        # Get Mask and NOCS predictions
        mask = masks[:, :, i]
        coord = coords[:, :, i, :]

        #Get the all 3D NOCS points and corresponding 2D image points
        NOCS_points = coord[mask == 1] - 0.5
        image_points = np.argwhere(mask == 1).astype(np.float32) # img_points must be np.float32
        image_points[:,[0, 1]] = image_points[:,[1, 0]] # Switch (height,width) to (width, height)

        # Remove outliers
        NOCS_points, image_points = remove_outliers_im(NOCS_points, image_points)

        # Solve correspondences between 3D NOCS points and 2D image points through the Efficient PnP algorithm
        start = time.time()
        retval, rvecs, tvecs, reprojectionError = cv2.solvePnPGeneric(
                objectPoints=NOCS_points, 
                imagePoints=image_points, 
                cameraMatrix=intrinsics, 
                distCoeffs=None,
                flags=cv2.SOLVEPNP_EPNP)
        elapsed = time.time() - start
        print('EPnP elapsed: ', elapsed)
        print('EPnP reprojectionError: ', reprojectionError)
        elapses.append(elapsed)
        
        # EPnP provides only one solution
        rvec = rvecs[0]
        tvec = tvecs[0]

        
        # Recover also the scale factor
        if with_scale:
            
            # Get the ODM (Object's Depth Mean point)
            dist = get_object_distance_from_depth(depth, mask, intrinsics, outlier_removal=True)
            
            # Recover the scale factor
            scale_factor = np.linalg.norm(dist) / np.linalg.norm(tvec)

            # Apply the scale to the translation
            tvec *= scale_factor
            print("TVEC",tvec)
            
            # Convert Rvec and Tvec to 4x4 RTs
            pred_RT = np.zeros((4, 4), dtype=np.float32) 

            pred_R = cv2.Rodrigues(rvec)[0]
            pred_RT[:3, :3] = np.diag([scale_factor, scale_factor, scale_factor]) / 1000 @ pred_R
            pred_RT[:3, 3] = tvec[:,0] / 1000
            pred_RT[3, 3] = 1
            
            # Add to result RTs
            RTs[i, :, :] = pred_RT

        result_rvecs.append(rvec)
        result_tvecs.append(tvec)

        # Recover the object dimensions in the NOCSpace
        abs_coord_pts = np.abs(NOCS_points)
        bbox_scales[i, :] = 2*np.amax(abs_coord_pts, axis=0)
        bbox_coordinates_3D = get_3d_bbox(bbox_scales[i, :], 0) # (3,N)
        bbox_coordinates_3D = bbox_coordinates_3D.transpose()+0.5 # (N,3)

        if False:
            """Drawing the NOCS and BBOX in 3D via Open3D"""

            pcl = o3d.geometry.PointCloud()
            pcl.points = o3d.utility.Vector3dVector(NOCS_points)
            pcl.colors = o3d.utility.Vector3dVector(NOCS_points+0.5)

            print("\nLet's draw a box using o3d.geometry.LineSet.")
            lines = [
                [0, 1],
                [0, 2],
                [1, 3],
                [2, 3],
                [4, 5],
                [4, 6],
                [5, 7],
                [6, 7],
                [0, 4],
                [1, 5],
                [2, 6],
                [3, 7],
            ]
            colors = [[1, 0, 0] for a in range(len(lines))]
            line_set = o3d.geometry.LineSet(
                points=o3d.utility.Vector3dVector(bbox_coordinates_3D-0.5),
                lines=o3d.utility.Vector2iVector(lines)
            )
            line_set.colors = o3d.utility.Vector3dVector(colors)

            o3d.visualization.draw_geometries([pcl,line_set])

    return RTs, result_rvecs, result_tvecs, bbox_scales, error_messages, elapses

def get_object_distance_from_depth(depth, mask, intrinsics, outlier_removal=False):
    """Returns the distance from the camera (in millimeters) to the object of interest,
        by getting the closest point from the object's depth point-cloud.
    """
    
    #Get 3d points and 2d pixels
    pts, idxs = backproject(depth, intrinsics, mask)

     
    # Get the ODM (Object's Depth Mean point)
    if outlier_removal: # remove outliers from depth
        inliers = remove_outliers(pts)
        center = inliers.get_center()
        #print("Center point:", center)
        #print("L2-norm of center point:", np.linalg.norm(center))
        return np.linalg.norm(center)

    else:
        # Compute distances from the camera
        distances = np.linalg.norm(pts, axis=1)
        
        # Get distance of the point that is closest to the camera
        closest_point_idx = np.argmin(distances)

        return distances[closest_point_idx]

def display_inlier_outlier(cloud, ind):
    inlier_cloud = cloud.select_by_index(ind)
    outlier_cloud = cloud.select_by_index(ind, invert=True)

    print("Showing outliers (red) and inliers (gray): ")
    outlier_cloud.paint_uniform_color([1, 0, 0])
    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
    o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud])


def remove_outliers_im(points_3D, image_points):

    """Removes outliers from a set of 3D points and corresponding image points"""
    # Create Open3D pointcloud from NOCS points
    pcl = o3d.geometry.PointCloud()
    pcl.points = o3d.utility.Vector3dVector(points_3D)

    ## Remove outliers 
    cl, ind = pcl.remove_statistical_outlier(nb_neighbors=25,
                                                std_ratio=0.5)

    # Select the inliers
    inlier_points_3D = pcl.select_by_index(ind)
    inlier_image_points = image_points[ind]

    # convert back to numpy matrix
    inlier_points_3D = np.asarray(inlier_points_3D.points)

    return inlier_points_3D, inlier_image_points

def remove_outliers(points):
    """Removes outliers from a set of 3D points
    """

    pcl = o3d.geometry.PointCloud()
    pcl.points = o3d.utility.Vector3dVector(points)

    # Remove outliers
    cl, ind = pcl.remove_statistical_outlier(nb_neighbors=20,
                                            std_ratio=0.5)
    
    #display_inlier_outlier(pcl, ind)

    # Select the inliers
    inliers = pcl.select_by_index(ind)
    
    return inliers

def nms_old(bounding_boxes, confidence_scores, classIDs, maskz, coordz, threshold):
    
    # If no bounding boxes, return empty list
    if len(bounding_boxes) == 0:
        return [], []

    # Bounding boxes
    boxes = np.array(bounding_boxes)

    # coordinates of bounding boxes
    start_y = boxes[:, 0]
    start_x = boxes[:, 1]
    end_y = boxes[:, 2]
    end_x = boxes[:, 3]

    # Confidence scores of bounding boxes
    score = np.array(confidence_scores)

    # Picked bounding boxes
    picked_indices = []

    # Compute areas of bounding boxes
    areas = (end_x - start_x + 1) * (end_y - start_y + 1)

    # Sort by confidence score of bounding boxes
    order = np.argsort(score)

    # Iterate bounding boxes
    while order.size > 0:
        # The index of largest confidence score
        index = order[-1]

        # Pick the bounding box with largest confidence score
        picked_indices.append(index)

        # Compute ordinates of intersection-over-union(IOU)
        x1 = np.maximum(start_x[index], start_x[order[:-1]])
        x2 = np.minimum(end_x[index], end_x[order[:-1]])
        y1 = np.maximum(start_y[index], start_y[order[:-1]])
        y2 = np.minimum(end_y[index], end_y[order[:-1]])

        # Compute areas of intersection-over-union
        w = np.maximum(0.0, x2 - x1 + 1)
        h = np.maximum(0.0, y2 - y1 + 1)
        intersection = w * h

        # Compute the ratio between intersection and union
        ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)

        left = np.where(ratio < threshold)
        order = order[left]

    allindices = np.arange(0,len(classIDs))

    indices2delete = np.setdiff1d(allindices,picked_indices)
    
    print('allindices: {} | indices2keep: {} | indices2delete: {}'.format(allindices, picked_indices, indices2delete))

    return indices2delete

def fix_background_nocs(nocs_img):
    """
    The background in the nocs image from SOM, is not black completely. Let's make it so.

    Inputs NOCS [h,w,3], [0-255], uint8. Weird background values
    Output NOCS [h,w,3], [0-255], uint8. Black background
    """
    # Convert to grayscale
    nocs_img_gray = cv2.cvtColor(nocs_img, cv2.COLOR_RGB2GRAY)
    # Use OTSU to threshold the NOCS image
    rect, nocs_binarized = cv2.threshold(nocs_img_gray,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
    
    # From [0-255] to [0-1]
    nocs_binary = nocs_binarized/255.0
    # From single channel to 3 channel
    nocs_bin_stack = np.stack([nocs_binary, nocs_binary, nocs_binary], axis=2)
    # Apply the mask to the original nocs image
    nocs_processed = nocs_img * nocs_bin_stack

    return nocs_processed

def clean_and_intersect_nocs(nocs, mask, show=True):
    """
    The mask and nocs are not the same. The mask seems bigger; sometimes the NOCS is.
    So if we mask the nocs, we get some background values in the NOCS points.
    Let's avoid this.

    nocs [h,w,3] in range [0-255]
    mask [h,w,3] in range [0-1]
    """
    
    # Sum the RGB values of each pixel
    s = np.sum(nocs, axis=2)
    
    # Get the indices where both MASK and NOCS are zero.
    x,y = np.where((mask!=0) & (s!=0))
    # Make it binary image
    xy_image = np.where((mask!=0) & (s!=0), 1, 0)
    # Convert to 3 channels
    xy_image3 = np.stack([xy_image, xy_image, xy_image], axis=2)
    
    # apply new mask to nocs
    nocs_cleaned = nocs * xy_image3
    nocs_cleaned = nocs_cleaned.astype('float32')
    # apply new mask to mask
    mask_cleaned = mask  * xy_image
    mask_cleaned = mask_cleaned.astype('uint8')
    
    points = np.where(xy_image != 0)

    if show:
        mask_to_show = np.stack([mask_cleaned, mask_cleaned, mask_cleaned], axis=2)
        mask_to_show *= 255
        cv2.imshow("nocs cleaned", nocs_cleaned[:,:,::-1])
        cv2.imshow("mask cleaned", mask_to_show)
        cv2.waitKey(0) 
        cv2.destroyAllWindows()

    return nocs_cleaned, mask_cleaned, points