<filename>models/model_mrcnn3d.py
import tensorflow as tf
import numpy as np
import math
import sys
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import tf_util
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/nn_distance'))
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/sampling'))
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/grouping'))
import tf_nndistance
from tf_sampling import farthest_point_sample, gather_point
from tf_grouping import query_ball_point, group_point
from pointnet_util import pointnet_fp_module, pointnet_sa_module
import config
from sklearn.neighbors import NearestNeighbors
from tensorflow.python.ops import array_ops

def placeholder_inputs(config):
    pc_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_POINT, 3))
    pc_ins_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_GROUP, config.NUM_POINT_INS, 3))
    group_label_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_POINT))
    group_indicator_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_GROUP))
    seg_label_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_POINT))
    seg_label_per_group_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_GROUP))
    bbox_ins_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_GROUP, 6))
    return pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, seg_label_per_group_pl, bbox_ins_pl

def multi_encoding_net(xyz, points, npoint, radius_list, nsample_list, mlp_list, mlp_list2, is_training, bn_decay, scope, bn=True, use_xyz=False, output_shift=False, shift_pred=None, fps_idx=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint: int32 -- #points sampled in farthest point sampling
            radius: list of float32 -- search radius in local region
            nsample: list of int32 -- how many points in each local region
            mlp: list of list of int32 -- output size for MLP on each point
            use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
        Return:
            new_xyz: (batch_size, npoint, 3) TF tensor
            new_points: (batch_size, npoint, mlp2[-1]) TF tensor
            shift_pred: (batch_size, npoint, 3) TF tensor
            fps_idx: (batch_size, npoint) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        if fps_idx is None:
            fps_idx = farthest_point_sample(npoint, xyz) # (batch_size, npoint)
        new_xyz = gather_point(xyz, fps_idx) # (batch_size, npoint, 3)
        new_points_list = []
        group_xyz_list = []
        for i in range(len(radius_list)):
            radius = radius_list[i]
            nsample = nsample_list[i]
            idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
            grouped_xyz = group_point(xyz, idx)
            grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # [B, nseed, nsmp, 3]
            if shift_pred is not None:
                grouped_xyz -= tf.tile(tf.expand_dims(shift_pred, 2), [1,1,nsample,1])
            if points is not None:
                grouped_points = group_point(points, idx)
                if use_xyz:
                    grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1)
            else:
                grouped_points = grouped_xyz
            for j,num_out_channel in enumerate(mlp_list[i]):
                grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                                padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                                scope='conv_prev_%d_%d'%(i,j), bn_decay=bn_decay)
            new_points_list.append(tf.reduce_max(grouped_points, axis=[2]))
        new_points = tf.concat(new_points_list, axis=-1) # (batch_size, npoint, \sum_k{mlp[k][-1]})
        for i,num_out_channel in enumerate(mlp_list2):
            new_points = tf_util.conv1d(new_points, num_out_channel, 1,
                                               padding='VALID', stride=1, bn=bn, is_training=is_training,
                                               scope='conv_post_%d'%i, bn_decay=bn_decay)
        if output_shift:
            shift_pred = tf_util.conv1d(new_points, 4, 1,
                                        padding='VALID', stride=1, scope='conv_shift_pred', activation_fn=None)
        return new_xyz, new_points, shift_pred, fps_idx

def shift_pred_net(xyz, points, npoint_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
        Return:
            pc_seed: (batch_size, npoint_seed, 3) TF tensor
            shift_pred_seed_4d: (batch_size, npoint_seed, 4) TF tensor
            ind_seed: (batch_size, npoint_seed) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        ind_seed = farthest_point_sample(npoint_seed, xyz) # (batch_size, npoint_seed)
        pc_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None # do not use color for shift prediction

        if return_fullfea:
            new_xyz = tf.concat((pc_seed, xyz), 1)
        else:
            new_xyz = pc_seed

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=0.8, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')

        # Feature Propagation layers
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 4, 1,
                             padding='VALID', stride=1, scope='conv_shift_pred', activation_fn=None)
        if return_fullfea:
            shift_pred_seed_4d, shift_pred_full_4d = tf.split(net, [npoint_seed, num_point], axis=1)
            end_points['shift_pred_full_4d'] = shift_pred_full_4d
        else:
            shift_pred_seed_4d = net

        end_points['pc_seed'] = pc_seed
        end_points['shift_pred_seed_4d'] = shift_pred_seed_4d
        end_points['ind_seed'] = ind_seed

        return end_points

def sem_net(xyz, points, npoint_sem, num_category, ind_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        npoint_seed = ind_seed.get_shape()[1].value

        ind_sem = farthest_point_sample(npoint_sem, xyz) # (batch_size, npoint_sem)
        new_xyz_sem = gather_point(xyz, ind_sem) # (batch_size, npoint_sem, 3)
        end_points['ind_sem'] = ind_sem
        new_xyz_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)

        if return_fullfea:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem, xyz), 1)
        else:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem), 1)
        
        l0_xyz = xyz
        l0_points = None

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=128, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=64, radius=0.8, nsample=32, mlp=[256,256,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        l5_xyz, l5_points, l5_indices = pointnet_sa_module(l4_xyz, l4_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer5')
        # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')

        # with FPN
        if return_fullfea:
            end_points['sem_fea_full_l5'] = pointnet_fp_module(xyz, l5_xyz, points, l5_points, [], is_training, bn_decay, scope='fa_layer0_fpn')
            end_points['sem_fea_full_l4'] = pointnet_fp_module(xyz, l4_xyz, points, l4_points, [], is_training, bn_decay, scope='fa_layer1_fpn')
            end_points['sem_fea_full_l3'] = pointnet_fp_module(xyz, l3_xyz, points, l3_points, [], is_training, bn_decay, scope='fa_layer2_fpn')
            end_points['sem_fea_full_l2'] = pointnet_fp_module(xyz, l2_xyz, points, l2_points, [], is_training, bn_decay, scope='fa_layer3_fpn')
            end_points['sem_fea_full_l1'] = pointnet_fp_module(xyz, l1_xyz, points, l1_points, [], is_training, bn_decay, scope='fa_layer4_fpn')

        # Feature Propagation layers
        l4_points = pointnet_fp_module(l4_xyz, l5_xyz, l4_points, l5_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        if return_fullfea:
            sem_fea_seed, sem_fea, sem_fea_full = tf.split(net, [npoint_seed, npoint_sem, num_point], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            end_points['sem_fea_full'] = sem_fea_full
            net = end_points['sem_fea_full']
        else:
            sem_fea_seed, sem_fea = tf.split(net, [npoint_seed, npoint_sem], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            net = end_points['sem_fea']

        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def sem_net_simp(xyz, points, num_category, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None
        end_points = {}

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=128, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=64, radius=0.8, nsample=32, mlp=[256,256,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        l5_xyz, l5_points, l5_indices = pointnet_sa_module(l4_xyz, l4_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer5')

        # Feature Propagation layers
        l4_points = pointnet_fp_module(l4_xyz, l5_xyz, l4_points, l5_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points


def sem_net2(xyz, points, npoint_sem, num_category, ind_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        npoint_seed = ind_seed.get_shape()[1].value

        ind_sem = farthest_point_sample(npoint_sem, xyz) # (batch_size, npoint_sem)
        new_xyz_sem = gather_point(xyz, ind_sem) # (batch_size, npoint_sem, 3)
        end_points['ind_sem'] = ind_sem
        new_xyz_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)

        if return_fullfea:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem, xyz), 1)
        else:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem), 1)
        
        l0_xyz = xyz
        l0_points = None

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=512, radius=0.2, nsample=64, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=128, radius=0.4, nsample=64, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3')

        # with FPN
        if return_fullfea:
            end_points['sem_fea_full_l3'] = pointnet_fp_module(xyz, l3_xyz, points, l3_points, [], is_training, bn_decay, scope='fa_layer2_fpn')
            end_points['sem_fea_full_l2'] = pointnet_fp_module(xyz, l2_xyz, points, l2_points, [], is_training, bn_decay, scope='fa_layer3_fpn')
            end_points['sem_fea_full_l1'] = pointnet_fp_module(xyz, l1_xyz, points, l1_points, [], is_training, bn_decay, scope='fa_layer4_fpn')

        # Feature Propagation layers
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer2')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer3')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        if return_fullfea:
            sem_fea_seed, sem_fea, sem_fea_full = tf.split(net, [npoint_seed, npoint_sem, num_point], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            end_points['sem_fea_full'] = sem_fea_full
        else:
            sem_fea_seed, sem_fea = tf.split(net, [npoint_seed, npoint_sem], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea

        # net = end_points['sem_fea_full']
        net = end_points['sem_fea']
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def sem_net2_simp(xyz, points, num_category, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None
        end_points = {}

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=512, radius=0.2, nsample=64, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=128, radius=0.4, nsample=64, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3')

        # Feature Propagation layers
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer2')
        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer3')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def pn2_fea_extractor(xyz, points, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
        Return:
            new_points: (batch_size, ndataset, channel_out) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = points

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=48, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=48, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=48, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer4')


        # Feature Propagation layers
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,128], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [128,64], is_training, bn_decay, scope='fa_layer2')
        new_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, [64,64,64], is_training, bn_decay, scope='fa_layer3')

        return new_points


def single_encoding_net(pc, mlp_list, mlp_list2, scope, is_training, bn_decay):
    ''' The encoding network for instance
    Input:
        pc: [B, N, 3]
    Return:
        net: [B, nfea]
    '''
    with tf.variable_scope(scope) as myscope:
        net = tf.expand_dims(pc, 2)
        for i,num_out_channel in enumerate(mlp_list):
            net = tf_util.conv2d(net, num_out_channel, [1,1],
                                 padding='VALID', stride=[1,1],
                                 bn=True, is_training=is_training,
                                 scope='conv%d'%i, bn_decay=bn_decay)
        net = tf.reduce_max(net, axis=[1])
        net = tf.squeeze(net, 1)
        for i,num_out_channel in enumerate(mlp_list2):
            net = tf_util.fully_connected(net, num_out_channel, bn=True, is_training=is_training,
                                          scope='fc%d'%i, bn_decay=bn_decay)
        return net

def fea_trans_net(input_fea, mlp_list, scope, is_training, bn_decay):
    with tf.variable_scope(scope) as myscope:
        net = input_fea
        nlayer = len(mlp_list)
        for i,num_out_channel in enumerate(mlp_list):
            if i<nlayer-1:
                net = tf_util.conv1d(net, num_out_channel, 1, padding='VALID', bn=True, is_training=is_training,
                                     scope='conv%d'%i, bn_decay=bn_decay)
            else:
                net = tf_util.conv1d(net, num_out_channel, 1, padding='VALID', activation_fn=None, scope='conv%d'%i)
        return net

def sample(mean, log_var):
    # Sample z
    z = mean + tf.exp(log_var/2.0) * tf.random_normal(tf.shape(mean), 0, 1, dtype=tf.float32)
    return z

def decoding_net(feat, num_point, scope, is_training, bn_decay):
    ''' The decoding network for shape generation
    Input:
        feat: [B, nsmp, nfea]
    Return:
        pc: [B, nsmp, num_point, 3]
    '''
    with tf.variable_scope(scope) as myscope:
        nsmp = feat.get_shape()[1].value
        nfea = feat.get_shape()[2].value
        feat = tf.reshape(feat, [-1, nfea])
        # UPCONV Decoder
        if num_point<=3072 and num_point>1536:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[5,5], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 1024
        elif num_point<=1536 and num_point>896:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 484
        elif num_point<=896 and num_point>384:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv4', activation_fn=None)
            num_point_conv = 256
        else:
            raise('Exception')
        pc_upconv = tf.reshape(net, [-1, num_point_conv, 3])

        num_point_fc = num_point - num_point_conv
        # FC Decoder
        net = tf_util.fully_connected(feat, 512, bn=True, is_training=is_training, scope='de_fc2', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, scope='de_fc3', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, num_point_fc*3, activation_fn=None, scope='de_fc4')
        pc_fc = tf.reshape(net, [-1, num_point_fc, 3])
        # Merge
        pc = tf.concat([pc_upconv, pc_fc], axis=1)
        pc = tf.reshape(pc, [-1, nsmp, num_point, 3])
        return pc

def decoding_conf_net(feat, num_point, scope, is_training, bn_decay):
    ''' The decoding network for shape generation & confidence generation for each generated point
    Input:
        feat: [B, nsmp, nfea]
    Return:
        pc: [B, nsmp, num_point, 3]
    '''
    with tf.variable_scope(scope) as myscope:
        nsmp = feat.get_shape()[1].value
        nfea = feat.get_shape()[2].value
        feat = tf.reshape(feat, [-1, nfea])
        # UPCONV Decoder
        if num_point<=3072 and num_point>1536:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[5,5], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 1024
        elif num_point<=1536 and num_point>896:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 484
        elif num_point<=896 and num_point>384:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv4', activation_fn=None)
            num_point_conv = 256
        else:
            raise('Exception')
        pc_upconv = tf.reshape(net, [-1, num_point_conv, 4])

        num_point_fc = num_point - num_point_conv
        # FC Decoder
        net = tf_util.fully_connected(feat, 512, bn=True, is_training=is_training, scope='de_fc2', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, scope='de_fc3', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, num_point_fc*4, activation_fn=None, scope='de_fc4')
        pc_fc = tf.reshape(net, [-1, num_point_fc, 4])
        # Merge
        pc = tf.concat([pc_upconv, pc_fc], axis=1)
        pc = tf.reshape(pc, [-1, nsmp, num_point, 4])
        pc_conf = pc[:,:,:,3]
        pc_xyz = pc[:,:,:,:3]
        return pc_xyz, pc_conf

def shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins_gt, num_category, scope, is_training, bn_decay=None, nsmp=128, return_fullfea=False):
    ''' Shape proposal generation
    Inputs:
        pc: [B, NUM_POINT, 3]
        pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys
        group_label: [B, NUM_POINT]
        group_indicator: [B, NUM_GROUP]
        bbox_ins_gt: [B, NUM_GROUP, 6], only used during training
    Returns:
        fb_logits: [B, NUM_SAMPLE, 2] confidence logits (before softmax)
        fb_prob: [B, NUM_SAMPLE, 2] confidence probabilities
        bbox_ins: [B, NUM_SAMPLE, (x, y, z, l, w, h)]
        entity_fea: [B, NUM_POINT, nfea] entity feature for each point
        center_pos: [B, NUM_POINT, 3] center coordinate for each point, in world coord sys
    '''
    with tf.variable_scope(scope) as myscope:
        # Parameter extraction
        batch_size = pc.get_shape()[0].value
        ngroup = pc_ins.get_shape()[1].value
        nsmp_ins = pc_ins.get_shape()[2].value
        end_points = {}

        # Shift prediction, ind_seed [B, nsmp], shift_pred_seed [B, nsmp, 3]
        end_points = shift_pred_net(pc, None, nsmp, end_points, 'shift_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        pc_seed = end_points['pc_seed']
        shift_pred_seed_4d = end_points['shift_pred_seed_4d']
        ind_seed = end_points['ind_seed']
        shift_pred_seed = tf.multiply(shift_pred_seed_4d[:,:,:3], shift_pred_seed_4d[:,:,3:])

        # Semantic prediction, sem_fea_seed [B, nsmp, nfea]
        end_points = sem_net(pc, None, 10000 if return_fullfea else 1024, num_category, ind_seed, end_points, 'sem_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        # end_points = sem_net2(pc, None, 4096, num_category, ind_seed, end_points, 'sem_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        sem_fea_seed = end_points['sem_fea_seed']

        # Encode instance, pcfea_ins_centered [B, ngroup, nfea_ins], pc_ins_center [B, ngroup, 1, 3]
        pc_ins_center = (tf.reduce_max(pc_ins, 2, keep_dims=True)+tf.reduce_min(pc_ins, 2, keep_dims=True))/2 # [B, ngroup, 1, 3] -> requires random padding for pc_ins generation
        pc_ins_centered = pc_ins-pc_ins_center
        idx = tf.where(tf.greater(group_indicator, 0))
        pc_ins_centered_list = tf.gather_nd(pc_ins_centered, idx)
        pcfea_ins_centered_list = single_encoding_net(pc_ins_centered_list, [64, 256, 512], [256], 'instance_encoder', is_training, bn_decay)
        nfea_ins = pcfea_ins_centered_list.get_shape()[1].value
        pcfea_ins_centered = tf.scatter_nd(tf.cast(idx,tf.int32), pcfea_ins_centered_list, tf.constant([batch_size, ngroup, nfea_ins])) # [B, ngroup, nfea_ins]
        
        # Collect instance feature for seed points [B, nsmp, nfea_seed]
        idx = tf.where(tf.greater_equal(ind_seed,0))
        ind_seed_aug = tf.concat((tf.expand_dims(tf.cast(idx[:,0],tf.int32),-1),tf.reshape(ind_seed,[-1,1])),1)
        group_label_seed = tf.reshape(tf.gather_nd(group_label, ind_seed_aug), [-1, nsmp]) # [B, nsmp]
        idx = tf.where(tf.greater_equal(group_label_seed,0))
        group_label_seed_aug = tf.concat((tf.expand_dims(tf.cast(idx[:,0],tf.int32),-1),tf.reshape(group_label_seed,[-1,1])),1)
        pcfea_ins_seed = tf.reshape(tf.gather_nd(pcfea_ins_centered, group_label_seed_aug), [-1, nsmp, nfea_ins])
        pc_ins_centered_seed = tf.reshape(tf.gather_nd(pc_ins_centered, group_label_seed_aug), [-1, nsmp, nsmp_ins, 3])
        pc_ins_center_seed = tf.reshape(tf.gather_nd(pc_ins_center, group_label_seed_aug), [-1, nsmp, 1, 3])
       
        # Encode context, pcfea_seed [B, nsmp, nfea_context]
        # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.25,0.5,1.0,1.5], [128,256,256,512], [[64,64,128], [64,128,256], [64,128,256], [64,128,256]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)
        _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.25,0.5,1.0,1.5], [96,192,192,384], [[64,64,128], [64,128,128], [64,128,128], [64,128,128]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)
        # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.4,0.8,1.5], [192,256,384], [[64,128,192], [64,128,192], [64,128,192]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)

        # Compute foreground/background score [B, nsmp, 2]
        fb_logits = fea_trans_net(pcfea_seed, [256, 64, 2], 'fb_logits', is_training, bn_decay)
        # fb_logits = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 64, 2], 'fb_logits', is_training, bn_decay)
        fb_prob = tf.nn.softmax(fb_logits, -1)

        #### CVAE
        # Compute mu and sigma [B, nsmp, 512]
        # mu_sigma_c = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 512, 512], 'mu_sigma_c', is_training, bn_decay)
        # mu_sigma_x = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed, pcfea_ins_seed), axis=-1), [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)
        mu_sigma_c = fea_trans_net(pcfea_seed, [256, 512, 512], 'mu_sigma_c', is_training, bn_decay)
        mu_sigma_x = fea_trans_net(tf.concat((pcfea_seed, pcfea_ins_seed), axis=-1), [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)
        # mu_sigma_x = fea_trans_net(pcfea_ins_seed, [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)

        # Sample z [B, nsmp, 256]
        mean = mu_sigma_x[:,:,:256]
        log_var = mu_sigma_x[:,:,256:]
        log_var = tf.clip_by_value(log_var, -10.0, 1.0)
        cmean = mu_sigma_c[:,:,:256]
        clog_var = mu_sigma_c[:,:,256:]
        clog_var = tf.clip_by_value(clog_var, -10.0, 1.0)
        zi = sample(mean, log_var)
        zc = cmean
        z = tf.cond(is_training, lambda: zi, lambda: zc)
        
        # Decode shapes pc [B, nsmp, nsmp_ins, 3]
        gcfeat = tf_util.conv1d(pcfea_seed, 256, 1, padding='VALID', bn=True, is_training=is_training,
                                scope='dec_fc', bn_decay=bn_decay)
        feat = tf.concat((z, gcfeat), axis=-1)
        pc_ins_pred, pc_ins_conf_pred = decoding_conf_net(feat, nsmp_ins, 'decoder', is_training=is_training, bn_decay=bn_decay)
        # pc_ins_pred = decoding_net(feat, nsmp_ins, 'decoder', is_training=is_training, bn_decay=bn_decay)
        pc_ins_pred = pc_ins_pred + tf.stop_gradient(tf.expand_dims(shift_pred_seed, 2))
        # pc_ins_pred = pc_ins_pred + tf.expand_dims(shift_pred_seed, 2)

        # Collect bbox for reconstructions, pc_ins_conf_pred [B, nsmp, nsmp_ins]
        pc_ins_pred_world_coord = pc_ins_pred + tf.expand_dims(pc_seed, 2)
        bbox_ins_pred = tf.concat(((tf.reduce_max(pc_ins_pred_world_coord, 2)+tf.reduce_min(pc_ins_pred_world_coord, 2))/2, 
                                  tf.reduce_max(pc_ins_pred_world_coord, 2)-tf.reduce_min(pc_ins_pred_world_coord, 2)), 2) # [B, nsmp, 6] -> center + l,w,h

        # pc_ins_pred_world_coord = pc_ins_pred + tf.expand_dims(pc_seed, 2) # [B, nsmp, nsmp_ins, 3]
        # inconf_pt_mask = 1000.0*tf.cast(tf.greater(pc_ins_conf_pred,0.5), tf.float32) 
        # inconf_pt_mask = tf.multiply(inconf_pt_mask, 1-tf.cast(tf.less_equal(pc_ins_conf_pred, 1e-3+tf.reduce_min(pc_ins_conf_pred, 2, keep_dims=True)), tf.float32))
        # inconf_pt_mask = tf.expand_dims(inconf_pt_mask, -1) # [B, nsmp, nsmp_ins, 1]

        # max_edge = tf.reduce_max(pc_ins_pred_world_coord-inconf_pt_mask, 2)
        # min_edge = tf.reduce_min(pc_ins_pred_world_coord+inconf_pt_mask, 2)
        # bbox_ins_pred = tf.concat(((max_edge+min_edge)/2, 
        #                           max_edge-min_edge), 2) # [B, nsmp, 6] -> center + l,w,h

        # # Random shift and scale box during training
        # bbox_ins_pred = tf.concat( (bbox_ins_pred[:,:,:3]+bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=0.0, stddev=0.1),
        #     bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=1.0, stddev=0.1)), axis=-1)

        # Propagate seed feature and center position
        if return_fullfea:
            entity_fea = pointnet_fp_module(pc, pc_seed, None, pcfea_seed, [], is_training=False, bn_decay=None, scope='entity_fea_prop', bn=False)
            # shift_pred = pointnet_fp_module(pc, pc_seed, None, shift_pred_seed, [], is_training=False, bn_decay=None, scope='shift_prop', bn=False)
            shift_pred = tf.multiply(end_points['shift_pred_full_4d'][:,:,:3], end_points['shift_pred_full_4d'][:,:,3:])
            center_pos = pc+shift_pred
            end_points['entity_fea'] = entity_fea # [B, N, 256] entity feature of each point
            end_points['center_pos'] = center_pos # [B, N, 3] center location in the world coord sys

        # Store end_points
        end_points['shift_pred_seed'] = shift_pred_seed # [B, nsmp, 3], offset from seed point to ins center
        end_points['shift_pred_seed_4d'] = shift_pred_seed_4d # [B, nsmp, 4], offset from seed point to ins center
        end_points['pc_seed'] = pc_seed # [B, nsmp, 3], seed point coordinate in world coord sys
        end_points['ind_seed'] = ind_seed # [B, nsmp], seed index
        end_points['pc_ins_centered_seed'] = pc_ins_centered_seed # [B, nsmp, nsmp_ins, 3], centered gt instance point cloud for each seed
        end_points['pc_ins_center_seed'] = pc_ins_center_seed # [B, nsmp, 1, 3], gt instance center for each seed in world coord sys
        end_points['mean'] = mean # [B, nsmp, 256]
        end_points['log_var'] = log_var
        end_points['cmean'] = cmean
        end_points['clog_var'] = clog_var
        end_points['fb_logits'] = fb_logits # [B, nsmp, 2] foreground/backgroud logits
        end_points['fb_prob'] = fb_prob # [B, nsmp, 2] foreground/background probability
        end_points['pc_ins_pred'] = pc_ins_pred # [B, nsmp, nsmp_ins, 3], in local sys, needs to add pc_seed
        end_points['pc_ins_conf_pred'] = pc_ins_conf_pred # [B, nsmp, nsmp_ins] per-point confidence
        end_points['bbox_ins_pred'] = bbox_ins_pred # [B, nsmp, 6]

        return end_points #, alexnetmodel

def nms_3d(boxes, scores, pre_nms_limit, max_output_size, iou_threshold=0.5, score_threshold=float('-inf')):
    ''' Non maximum suppression in 3D
    Inputs:
        boxes: [B, N, 6] center + l,w,h
        scores: [B, N] prob between 0 and 1
    Outputs:
        selected_indices: [B, M]
    '''

    batch_size = scores.shape[0]    
    num_box_input = scores.shape[1]    
    sidx = np.argsort(-scores, 1) # [B, N] from large to small
    selected_indices = -np.ones((batch_size, max_output_size), dtype=np.int32) # [B, M]
    for i in range(batch_size):
        cursidx = sidx[i,:]
        curscores = scores[i,:]
        curvolume = boxes[i,:,3]*boxes[i,:,4]*boxes[i,:,5]
        if pre_nms_limit>0:
            cursidx = cursidx[:pre_nms_limit]
        cursidx = cursidx[curscores[cursidx]>score_threshold]
        count = 0
        while len(cursidx)>0 and count<max_output_size:
            selected_indices[i,count] = cursidx[0]
            count += 1
            vA = np.maximum(boxes[i,[cursidx[0]],:3]-boxes[i,[cursidx[0]],3:]/2, boxes[i,cursidx,:3]-boxes[i,cursidx,3:]/2)
            vB = np.minimum(boxes[i,[cursidx[0]],:3]+boxes[i,[cursidx[0]],3:]/2, boxes[i,cursidx,:3]+boxes[i,cursidx,3:]/2)
            intersection_cube = np.maximum(vB-vA,0)
            intersection_volume = intersection_cube[:,0]*intersection_cube[:,1]*intersection_cube[:,2]
            iou = np.divide(intersection_volume,curvolume[cursidx]+curvolume[cursidx[0]]-intersection_volume+1e-8)
            cursidx = np.delete(cursidx, np.where(iou>iou_threshold)[0])
    return selected_indices

def gather_selection(source, selected_idx, max_selected_size):
    '''
    Inputs:
        source: [B, N, C]
        selected_idx: [B, M], -1 means not selecting anything
    Returns:
        target: [B, M, C], 0 padded
    '''
    batch_size = source.get_shape()[0].value
    fea_size = source.get_shape()[2].value
    pos_idx = tf.cast(tf.where(tf.greater_equal(selected_idx,0)), tf.int32)
    selected_idx_vec = tf.gather_nd(selected_idx, pos_idx)
    target_vec = tf.gather_nd(source, tf.concat((tf.expand_dims(pos_idx[:,0],-1), tf.reshape(selected_idx_vec,[-1,1])),1))
    target = tf.scatter_nd(pos_idx, target_vec, tf.constant([batch_size, max_selected_size, fea_size]))
    return target

def trim_zeros_graph(boxes, name=None):
    """Often boxes are represented with matrices of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 6] matrix of boxes.
    non_zeros: [N] a 1D boolean mask identifying the rows to keep
    """
    non_zeros = tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1), tf.bool)
    boxes = tf.boolean_mask(boxes, non_zeros, name=name)
    return boxes, non_zeros

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

def box_shrink(box, pc):
    ''' Shrink bounding box so that it is tight with respect to pc
    Inputs:
        box: [B, NUM_SAMPLE, 6]
        pc: [B, NUM_POINT, 3]
    Returns:
        box: [B, NUM_SAMPLE, 6]
    '''
    pc_aug = tf.expand_dims(pc, 1) # [B, 1, NUM_POINT, 3]
    box_aug = tf.expand_dims(box, 2) # [B, NUM_SAMPLE, 1, 6]
    box_masks = tf.logical_and(pc_aug>=(box_aug[:,:,:,:3]-box_aug[:,:,:,3:]/2),
                               pc_aug<=(box_aug[:,:,:,:3]+box_aug[:,:,:,3:]/2)) # [B, NUM_SAMPLE, NUM_POINT, 3]
    box_masks = tf.logical_and(tf.logical_and(box_masks[:,:,:,0], box_masks[:,:,:,1]), box_masks[:,:,:,2]) # [B, NUM_SAMPLE, NUM_POINT]
    box_out_masks = 1-tf.cast(tf.expand_dims(box_masks, -1), tf.float32) # [B, NUM_SAMPLE, NUM_POINT, 1]
    gamma = 1e4 # a large number for the box estimation trick
    box_max = tf.reduce_max(pc_aug-gamma*box_out_masks,2) # [B, NUM_SAMPLE, 3]
    box_min = tf.reduce_min(pc_aug+gamma*box_out_masks,2)
    box = tf.concat( ((box_max+box_min)/2, box_max-box_min+1e-3), 2) # [B, NUM_SAMPLE, 6]
    keep = tf.greater(box_max-box_min, 0)
    keep = tf.logical_and(tf.logical_and(keep[:,:,0], keep[:,:,1]), keep[:,:,2])
    keep = tf.expand_dims(tf.cast(keep, tf.float32), -1)
    box = tf.multiply(box, keep)
    return box

def box_refinement(box, gt_box):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (center_x, center_y, center_z, l, w, h)]
    """
    box = tf.cast(box, tf.float32)
    gt_box = tf.cast(gt_box, tf.float32)

    dz = (gt_box[:,2] - box[:,2]) / (box[:,5]+1e-8)
    dy = (gt_box[:,1] - box[:,1]) / (box[:,4]+1e-8)
    dx = (gt_box[:,0] - box[:,0]) / (box[:,3]+1e-8)
    dh = tf.log(gt_box[:,5] / (box[:,5]+1e-8))
    dw = tf.log(gt_box[:,4] / (box[:,4]+1e-8))
    dl = tf.log(gt_box[:,3] / (box[:,3]+1e-8))

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

def apply_box_delta(box, delta):
    ''' Apply bounding box delta to refine box
    Inputs:
        box: [NUM_SAMPLE, 6]
        delta: [NUM_SAMPLE, 6]
    Returns:
        box_refined: [NUM_SAMPLE, 6]
    '''
    delta = tf.stack([delta[:,2], delta[:,1], delta[:,0], delta[:,5], delta[:,4], delta[:,3]], axis=1)
    box_refined_part2 = tf.multiply(tf.exp(delta[:,3:]), box[:,3:])
    box_refined_part1 = tf.multiply(delta[:,:3], box[:,3:])+box[:,:3]
    box_refined = tf.concat((box_refined_part1, 1.3*box_refined_part2), axis=1)
    return box_refined

def sample_points_within_box(masks, nsmp):
    '''
    Inputs:
        masks: [nmask, npoint]
        nsmp: scalar
    Returns:
        masks_selection_idx: [nmask, nsmp]
    '''
    if masks.shape[0]==0:
        return np.zeros((0, nsmp), dtype=np.int32)
    else:
        masks_selection_idx = [np.random.choice(np.where(masks[i,:])[0], nsmp, replace=True) if len(np.where(masks[i,:])[0])>0 else np.zeros(nsmp, dtype=np.int32) for i in np.arange(masks.shape[0])]
        masks_selection_idx = np.stack(masks_selection_idx, 0).astype(np.int32)
        return masks_selection_idx

def spn_target_gen(proposals, proposal_seed_class_ids, gt_class_ids, gt_boxes):
    ''' SPN target generation
    Inputs:
        proposals: [NUM_SAMPLE, 6]
        proposal_seed_class_ids: [NUM_SAMPLE] - 0 is background class and 1 is foreground class
        gt_class_ids: [NUM_GROUP] - 0 is background class and 1 is foreground class
        gt_boxes: [NUM_GROUP, 6]
    Returns:
        spn_match: [NUM_SAMPLE], 1 = positive, -1 = negative, 0 = neutral
    '''
    # Remove zero padding
    non_zeros = tf.where(tf.greater(gt_class_ids, 0))[:,0]
    gt_boxes = tf.gather(gt_boxes, non_zeros, axis=0, name="spn_trim_gt_boxes")
    gt_class_ids = tf.gather(gt_class_ids, non_zeros, axis=0, name="spn_trim_gt_class_ids")

    # Compute IoU [NUM_SAMPLE, NUM_GROUP_TRIMMED]
    proposals_aug = tf.expand_dims(proposals, 1)
    gt_boxes_aug = tf.expand_dims(gt_boxes, 0)
    proposal_volume = tf.multiply(tf.multiply(proposals_aug[:,:,3], proposals_aug[:,:,4]), proposals_aug[:,:,5])
    gt_boxes_volume = tf.multiply(tf.multiply(gt_boxes_aug[:,:,3], gt_boxes_aug[:,:,4]), gt_boxes_aug[:,:,5])
    vA = tf.maximum(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]-gt_boxes_aug[:,:,3:]/2)
    vB = tf.minimum(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]+gt_boxes_aug[:,:,3:]/2)
    intersection_cube = tf.maximum(vB-vA,0)
    intersection_volume = tf.multiply(tf.multiply(intersection_cube[:,:,0], intersection_cube[:,:,1]), intersection_cube[:,:,2])
    ious = tf.divide(intersection_volume,proposal_volume+gt_boxes_volume-intersection_volume+1e-8)

    # Determine positive and negative ROIs
    roi_iou_max = tf.reduce_max(ious, axis=1)
    # 1. Positive ROIs are those with >=0.5 IoU with a GT box and gt class id is 1
    # We also guarantee that the boxes with the largest IoU with a GT box is also treated as positive
    spn_match_positive = tf.logical_and(tf.greater_equal(roi_iou_max, 0.5), tf.equal(proposal_seed_class_ids,1))
    masked_ious = tf.multiply(ious, tf.cast(tf.equal(tf.expand_dims(proposal_seed_class_ids,-1),1), tf.float32))
    positive_aug_idx = tf.argmax(masked_ious,0)
    positive_aug_idx = tf.boolean_mask(positive_aug_idx, tf.reduce_max(masked_ious, 0)>0)
    updates = tf.cast(tf.ones_like(positive_aug_idx), tf.float32)
    spn_match_positive_aug = tf.cond(
        tf.size(positive_aug_idx)>0,
        true_fn = lambda: tf.scatter_nd(positive_aug_idx, updates, spn_match_positive.shape),
        false_fn = lambda: tf.constant(0.0))
    spn_match_positive = tf.greater(tf.cast(spn_match_positive, tf.float32)+spn_match_positive_aug, 0)
    # 2. Negative ROIs are those with < 0.5 with every GT box and not a postive ROI
    spn_match_negative = tf.logical_and(tf.less(roi_iou_max, 0.5), tf.logical_not(spn_match_positive))

    spn_match = tf.cast(spn_match_positive, tf.float32)-tf.cast(spn_match_negative, tf.float32)

    return spn_match


def detection_target_gen(proposals, gt_class_ids, gt_boxes, gt_masks, pc, config):
    ''' Generate detection targets for training
    Inputs:
        proposals: [SPN_NMS_MAX_SIZE, 6], zero padded
        gt_class_ids: [NUM_GROUP] - 0 is background class and objects start from 1
        gt_boxes: [NUM_GROUP, 6]
        gt_masks: [NUM_POINT, NUM_GROUP]
        pc: [NUM_POINT, 3]
    Returns:
        rois: [TRAIN_ROIS_PER_IMAGE, 6], zero padded
        roi_gt_class_ids: [TRAIN_ROIS_PER_IMAGE] - 0 is invalid class and used for padding, zero padded
        deltas: [TRAIN_ROIS_PER_IMAGE, 6], zero padded
        masks_selection_idx: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded
        masks: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - binary mask, zero padded
    '''
    # Remove zero padding
    proposals, _ = trim_zeros_graph(proposals, name='trim_proposals')
    gt_boxes, non_zeros = trim_zeros_graph(gt_boxes, name="trim_gt_boxes")
    gt_class_ids = tf.boolean_mask(gt_class_ids, non_zeros,
                                   name="trim_gt_class_ids")
    gt_masks = tf.gather(gt_masks, tf.where(non_zeros)[:, 0], axis=1,
                         name="trim_gt_masks")

    # Remove empty proposals
    pc_aug = tf.expand_dims(pc,1) # [N, 1, 3]
    proposals_aug = tf.expand_dims(proposals, 0) # [1, NP, 6]
    roi_masks = tf.logical_and(pc_aug>=(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2),
                               pc_aug<=(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,0], roi_masks[:,:,1]), roi_masks[:,:,2]) # [N, NP]
    non_empty_idx = tf.where(tf.greater(tf.reduce_sum(tf.cast(roi_masks, tf.float32), 0), 0))[:,0]
    roi_masks = tf.gather(roi_masks, non_empty_idx, axis=1) # [N, NP']
    proposals = tf.gather(proposals, non_empty_idx, axis=0) # [NP', 6]

    # Compute IoU [n_proposal, n_gt_boxes]
    proposals_aug = tf.expand_dims(proposals, 1)
    gt_boxes_aug = tf.expand_dims(gt_boxes, 0)
    proposal_volume = tf.multiply(tf.multiply(proposals_aug[:,:,3], proposals_aug[:,:,4]), proposals_aug[:,:,5])
    gt_boxes_volume = tf.multiply(tf.multiply(gt_boxes_aug[:,:,3], gt_boxes_aug[:,:,4]), gt_boxes_aug[:,:,5])
    vA = tf.maximum(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]-gt_boxes_aug[:,:,3:]/2)
    vB = tf.minimum(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]+gt_boxes_aug[:,:,3:]/2)
    intersection_cube = tf.maximum(vB-vA,0)
    intersection_volume = tf.multiply(tf.multiply(intersection_cube[:,:,0], intersection_cube[:,:,1]), intersection_cube[:,:,2])
    ious = tf.divide(intersection_volume,proposal_volume+gt_boxes_volume-intersection_volume+1e-8)
    
    # Determine positive and negative ROIs
    roi_iou_max = tf.reduce_max(ious, axis=1)
    # 1. Positive ROIs are those with >= 0.5 IoU with a GT box
    positive_indices = tf.where(roi_iou_max >= 0.5)[:, 0]
    # 2. Negative ROIs are those with < 0.5 with every GT box
    negative_indices = tf.where(roi_iou_max < 0.5)[:, 0]

    # Subsample ROIs. Aim for 33% positive
    # Positive ROIs
    positive_count = int(config.TRAIN_ROIS_PER_IMAGE *
                         config.ROI_POSITIVE_RATIO)
    positive_indices = tf.random_shuffle(positive_indices)[:positive_count]
    positive_count = tf.shape(positive_indices)[0]
    # Negative ROIs. Add enough to maintain positive:negative ratio.
    r = 1.0 / config.ROI_POSITIVE_RATIO
    negative_count = tf.cast(r * tf.cast(positive_count, tf.float32), tf.int32) - positive_count
    negative_indices = tf.random_shuffle(negative_indices)[:negative_count]
    # Gather selected ROIs [POSITIVE_COUNT/NEGATIVE_COUNT, 6]
    positive_rois = tf.gather(proposals, positive_indices)
    negative_rois = tf.gather(proposals, negative_indices)

    # Assign positive ROIs to GT boxes. roi_gt_boxes: [POSITIVE_COUNT, 6], roi_gt_class_ids: [POSITIVE_COUNT]
    positive_ious = tf.gather(ious, positive_indices)
    roi_gt_box_assignment = tf.cond(
        tf.greater(tf.shape(positive_ious)[1], 0),
        true_fn = lambda: tf.argmax(positive_ious, axis=1),
        false_fn = lambda: tf.cast(tf.constant([]),tf.int64)
    )
    roi_gt_boxes = tf.gather(gt_boxes, roi_gt_box_assignment)
    roi_gt_class_ids = tf.gather(gt_class_ids, roi_gt_box_assignment)

    # Compute bbox refinement for positive ROIs. delta: [POSITIVE_COUNT, 6]
    deltas = box_refinement(positive_rois, roi_gt_boxes)
    deltas /= config.BBOX_STD_DEV

    # Compute mask targets
    roi_gt_masks = tf.cast(tf.gather(gt_masks, roi_gt_box_assignment, axis=1), tf.bool) # [NUM_POINT, POSITIVE_COUNT]
    positive_roi_masks = tf.gather(roi_masks, positive_indices, axis=1)
    masks_full = tf.transpose(tf.logical_and(positive_roi_masks, roi_gt_masks)) # [POSITIVE_COUNT, NUM_POINT]
    masks_selection_idx = tf.stop_gradient(tf.py_func(sample_points_within_box, 
        [tf.transpose(positive_roi_masks), config.NUM_POINT_INS_MASK], tf.int32)) # [POSITIVE_COUNT, NUM_POINT_INS_MASK]
    smp_idx = tf.reshape(tf.tile(tf.reshape(tf.range(positive_count),[-1,1]),[1, config.NUM_POINT_INS_MASK]),[-1,1])
    smp_idx = tf.concat((smp_idx, tf.reshape(masks_selection_idx,[-1,1])),1)
    masks = tf.reshape(tf.gather_nd(masks_full, smp_idx),[-1, config.NUM_POINT_INS_MASK]) # [POSITIVE_COUNT, NUM_POINT_INS_MASK]

    # Append negative ROIs and pad bbox deltas and masks that
    # are not used for negative ROIs with zeros.
    rois = tf.concat([positive_rois, negative_rois], axis=0)
    N = tf.shape(negative_rois)[0]
    P = tf.maximum(config.TRAIN_ROIS_PER_IMAGE - tf.shape(rois)[0], 0)
    rois = tf.pad(rois, [(0, P), (0, 0)])
    roi_gt_class_ids = tf.pad(roi_gt_class_ids, [(0, N + P)])
    deltas = tf.pad(deltas, [(0, N + P), (0, 0)])
    masks_selection_idx = tf.pad(masks_selection_idx, [[0, N + P], (0, 0)])
    masks = tf.pad(masks, [[0, N + P], (0, 0)])

    return rois, roi_gt_class_ids, deltas, masks_selection_idx, masks

def mask_selection_gen(proposals, pc, num_rois, config, empty_removal=True):
    ''' Generate detection targets for training
    Inputs:
        proposals: [SPN_NMS_MAX_SIZE, 6], zero padded
        pc: [NUM_POINT, 3]
    Returns:
        proposals: [NUM_ROIS, 6]
        masks_selection_idx: [NUM_ROIS, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded
    '''
    # Remove zero padding
    proposals, _ = trim_zeros_graph(proposals, name='trim_proposals')

    # Remove empty proposals
    pc_aug = tf.expand_dims(pc,1) # [N, 1, 3]
    proposals_aug = tf.expand_dims(proposals, 0) # [1, NP, 6]
    roi_masks = tf.logical_and(pc_aug>=(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2-1e-3),
                               pc_aug<=(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2+1e-3))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,0], roi_masks[:,:,1]), roi_masks[:,:,2]) # [N, NP]
    if empty_removal:
        non_empty_idx = tf.where(tf.greater(tf.reduce_sum(tf.cast(roi_masks, tf.float32), 0), 0))[:,0]
        roi_masks = tf.gather(roi_masks, non_empty_idx, axis=1) # [N, NP']
        proposals = tf.gather(proposals, non_empty_idx, axis=0) # [NP', 6]
    proposals_count = tf.shape(proposals)[0]

    # Generate mask selection index
    masks_selection_idx = tf.stop_gradient(tf.py_func(sample_points_within_box, 
        [tf.transpose(roi_masks), config.NUM_POINT_INS_MASK], tf.int32)) # [NP', NUM_POINT_INS_MASK]

    # Append negative ROIs and pad bbox deltas and masks that
    # are not used for negative ROIs with zeros.
    P = tf.maximum(num_rois - proposals_count, 0)
    masks_selection_idx = tf.pad(masks_selection_idx, [[0, P], (0, 0)])
    proposals = tf.pad(proposals, [(0, P), (0, 0)])

    return proposals, masks_selection_idx

def points_cropping(pc, pc_fea, pc_center, rois, masks_selection_idx, num_rois, num_point_per_roi, normalize_crop_region=True):
    ''' Crop points for network heads, in analogy to ROIAlign
    Inputs:
        pc: [B, NUM_POINT, 3]
        pc_fea: [B, NUM_POINT, NFEA]
        pc_center: [B, NUM_POINT, 3]
        rois: [B, NUM_ROIS, 6], zero padded
        masks_selection_idx: [B, NUM_ROIS, NUM_POINT_PER_ROI]
    Returns:
        pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
    '''
    batch_size = pc.get_shape()[0].value
    smp_idx = tf.reshape(tf.tile(tf.reshape(tf.range(batch_size),[-1,1]),[1, num_rois*num_point_per_roi]),[-1,1])
    smp_idx = tf.concat((smp_idx, tf.reshape(masks_selection_idx,[-1,1])),1)
    pc_fea_cropped = tf.reshape(tf.gather_nd(pc_fea, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_center_cropped = tf.reshape(tf.gather_nd(pc_center, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_coord_cropped_unnormalized = tf.reshape(tf.gather_nd(pc, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_coord_cropped = pc_coord_cropped_unnormalized

    # convert world coord to local
    rois_center = tf.expand_dims(rois[:,:,:3], 2)
    pc_coord_cropped = pc_coord_cropped-rois_center
    pc_center_cropped = pc_center_cropped-rois_center
    if normalize_crop_region:
        # scale box to [1,1,1]
        rois = rois+tf.cast(tf.equal(tf.reduce_sum(rois, 2, keep_dims=True),0),tf.float32)
        rois_size = tf.expand_dims(rois[:,:,3:], 2)
        pc_coord_cropped = tf.divide(pc_coord_cropped, rois_size)
        pc_center_cropped = tf.divide(pc_center_cropped, rois_size)
    return pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized

def refine_detections(rois, probs, deltas, pc, fb_prob, sem_prob, config):
    '''Refine classified proposals and filter overlaps and return final
    detections.
    Inputs:
        rois: [NUM_ROIS, 6], zero padded, in world coord sys
        probs: [NUM_ROIS, NUM_CATEGORY] - 0 is background class and objects start from 1
        deltas: [NUM_ROIS, NUM_CATEGORY, 6]
        pc: [NUM_POINT, 3]
        fb_prob: [NUM_ROIS]
        sem_prob: [NUM_ROIS]
    Returns:
        detections: [NUM_DETECTIONS, (center_x, center_y, center_z, l, w, h, class_id, score)]
    '''
    # Class IDs per ROI
    class_ids = tf.argmax(probs, axis=1, output_type=tf.int32)
    # Class probability of the top class of each ROI
    indices = tf.stack([tf.range(probs.shape[0]), class_ids], axis=1)
    class_scores = tf.gather_nd(probs, indices) # [NUM_ROIS]
    # Class-specific bounding box deltas
    deltas_specific = tf.gather_nd(deltas, indices) # [NUM_ROIS, 6]
    # Apply bounding box deltas

    # Shape: [NUM_ROIS, (center_x, center_y, center_z, l, w, h)]
    refined_rois = apply_box_delta(rois, deltas_specific * config.BBOX_STD_DEV)
    # Shrink boxes
    if config.SHRINK_BOX:
        refined_rois = tf.squeeze(box_shrink(tf.expand_dims(refined_rois, 0), tf.expand_dims(pc,0)),0)

    # Filter out background boxes
    keep = tf.where(class_ids > 0)[:, 0]
    # keep = tf.where(tf.logical_and(tf.logical_and(class_ids > 0, fb_prob>0.01), sem_prob>0.01))[:, 0]
    # Filter out low confidence boxes
    if config.DETECTION_MIN_CONFIDENCE:
        conf_keep = tf.where(class_scores >= config.DETECTION_MIN_CONFIDENCE)[:, 0]
        keep = tf.sets.set_intersection(tf.expand_dims(keep, 0),
                                        tf.expand_dims(conf_keep, 0))
        keep = tf.sparse_tensor_to_dense(keep)[0]

    # Apply per-class NMS
    # 1. Prepare variables
    pre_nms_class_ids = tf.gather(class_ids, keep)
    pre_nms_scores = tf.gather(class_scores, keep) # [NUM_KEEP]
    # pre_nms_scores = tf.gather(class_scores*fb_prob*sem_prob, keep) # [NUM_KEEP]
    pre_nms_rois = tf.gather(refined_rois,   keep)
    unique_pre_nms_class_ids = tf.unique(pre_nms_class_ids)[0]

    def nms_keep_map(class_id):
        """Apply Non-Maximum Suppression on ROIs of the given class."""
        # Indices of ROIs of the given class
        ixs = tf.where(tf.equal(pre_nms_class_ids, class_id))[:, 0]
        # Apply NMS

        class_keep = tf.py_func(nms_3d, [tf.expand_dims(tf.gather(pre_nms_rois, ixs),0),
            tf.expand_dims(tf.gather(pre_nms_scores, ixs),0),
            -1, config.DETECTION_MAX_INSTANCES,
            config.DETECTION_NMS_THRESHOLD, float('-inf')], tf.int32)
        class_keep = tf.squeeze(class_keep, 0) # [DETECTION_MAX_INSTANCES], -1 padded
        class_keep = tf.gather(class_keep, tf.where(class_keep > -1)[:,0]) # [<=DETECTION_MAX_INSTANCES], no padding

        # Map indices
        class_keep = tf.gather(keep, tf.gather(ixs, class_keep))
        # Pad with -1 so returned tensors have the same shape
        gap = config.DETECTION_MAX_INSTANCES - tf.shape(class_keep)[0]
        class_keep = tf.pad(class_keep, [(0, gap)],
                            mode='CONSTANT', constant_values=-1)
        # Set shape so map_fn() can infer result shape
        class_keep.set_shape([config.DETECTION_MAX_INSTANCES])
        return class_keep

    # 2. Map over class IDs
    nms_keep = tf.map_fn(nms_keep_map, unique_pre_nms_class_ids,
                         dtype=tf.int64)
    # 3. Merge results into one list, and remove -1 padding
    nms_keep = tf.reshape(nms_keep, [-1])
    nms_keep = tf.gather(nms_keep, tf.where(nms_keep > -1)[:, 0])
    # 4. Compute intersection between keep and nms_keep
    keep = tf.sets.set_intersection(tf.expand_dims(keep, 0),
                                    tf.expand_dims(nms_keep, 0))
    keep = tf.sparse_tensor_to_dense(keep)[0]
    # Keep top detections
    roi_count = config.DETECTION_MAX_INSTANCES
    class_scores_keep = tf.gather(class_scores, keep)
    # class_scores_keep = tf.gather(class_scores*fb_prob*sem_prob, keep)
    num_keep = tf.minimum(tf.shape(class_scores_keep)[0], roi_count)
    top_ids = tf.nn.top_k(class_scores_keep, k=num_keep, sorted=True)[1]
    keep = tf.gather(keep, top_ids)

    # Arrange output as [N, (center_x, center_y, center_z, l, w, h, class_id, score)]
    detections = tf.concat([
        tf.gather(refined_rois, keep),
        tf.to_float(tf.gather(class_ids, keep))[..., tf.newaxis],
        tf.gather(class_scores, keep)[..., tf.newaxis]
        ], axis=1)

    # Pad with zeros if detections < DETECTION_MAX_INSTANCES
    gap = config.DETECTION_MAX_INSTANCES - tf.shape(detections)[0]
    detections = tf.pad(detections, [(0, gap), (0, 0)], "CONSTANT")
    return detections

def classification_head(pc, pc_fea, num_category, mlp_list, mlp_list2, is_training, bn_decay, scope, bn=True):
    ''' Classification head for both class id prediction and bbox delta regression
    Inputs:
        pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        num_category: scalar
    Returns:
        logits: [B, NUM_ROIS, NUM_CATEGORY]
        probs: [B, NUM_ROIS, NUM_CATEGORY]
        bbox_deltas: [B, NUM_ROIS, NUM_CATEGORY, (dz, dy, dx, log(dh), log(dw), log(dl))]
    '''
    with tf.variable_scope(scope) as myscope:
        num_rois = pc.get_shape()[1].value
        grouped_points = tf.concat((pc_fea, pc), -1)
        for i,num_out_channel in enumerate(mlp_list):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_prev_%d'%i, bn_decay=bn_decay)
        new_points = tf.reduce_max(grouped_points, axis=2)
        for i,num_out_channel in enumerate(mlp_list2):
            new_points = tf_util.conv1d(new_points, num_out_channel, 1,
                                        padding='VALID', stride=1, bn=bn, is_training=is_training,
                                        scope='conv_post_%d'%i, bn_decay=bn_decay)
        logits = tf_util.conv1d(new_points, num_category, 1, padding='VALID',
                                stride=1, scope='conv_classify', activation_fn=None)
        probs = tf.nn.softmax(logits, 2)
        bbox_deltas = tf_util.conv1d(new_points, num_category*6, 1, padding='VALID',
                                     stride=1, scope='conv_bbox_regress', activation_fn=None)
        bbox_deltas = tf.reshape(bbox_deltas, [-1, num_rois, num_category, 6])
        return logits, probs, bbox_deltas

def segmentation_head(pc, pc_fea, num_category, mlp_list, mlp_list2, mlp_list3, is_training, bn_decay, scope, bn=True):
    ''' Segmentation head
    Inputs:
        pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        num_category: scalar
    Returns:
        masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
    '''
    with tf.variable_scope(scope) as myscope:
        num_rois = pc.get_shape()[1].value
        num_point_per_roi = pc.get_shape()[2].value
        grouped_points = tf.concat((pc_fea, pc), -1)
        for i,num_out_channel in enumerate(mlp_list):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_prev_%d'%i, bn_decay=bn_decay)
        local_feat = grouped_points
        for i,num_out_channel in enumerate(mlp_list2):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_%d'%i, bn_decay=bn_decay)
        global_feat = tf.reduce_max(grouped_points, axis=2, keep_dims=True)
        global_feat_expanded = tf.tile(global_feat, [1, 1, num_point_per_roi, 1])
        new_points = tf.concat((global_feat_expanded, local_feat), -1)
        for i,num_out_channel in enumerate(mlp_list3):
            new_points = tf_util.conv2d(new_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_post_%d'%i, bn_decay=bn_decay)

        masks = tf_util.conv2d(new_points, num_category, [1, 1], padding='VALID',
                               stride=[1,1], scope='conv_seg', activation_fn=None)
        return masks

def dict_stop_gradient(dict_in):
    keys = dict_in.keys()
    for key in keys:
        dict_in[key] = tf.stop_gradient(dict_in[key])
    return dict_in

def select_segmentation(mrcnn3d_masks, class_ids):
    ''' Convert segmentation into point cloud label
    Inputs:
        mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
        class_ids: [B, NUM_ROIS]
    Returns:
        mrcnn3d_mask_selected: [B, NUM_ROIS, NUM_POINT_PER_ROI]
    '''
    batch_size = mrcnn3d_masks.get_shape()[0].value
    num_rois = mrcnn3d_masks.get_shape()[1].value
    num_point_per_roi = mrcnn3d_masks.get_shape()[2].value
    num_category = mrcnn3d_masks.get_shape()[3].value

    mrcnn3d_masks = tf.reshape(mrcnn3d_masks, [-1, num_point_per_roi, num_category])
    mrcnn3d_masks = tf.transpose(mrcnn3d_masks, perm=[0, 2, 1]) # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI]
    class_ids = tf.cast(tf.reshape(class_ids, [-1]), tf.int32)
    class_ids_aug = tf.stack([tf.range(batch_size*num_rois, dtype=tf.int32), class_ids], 1)
    mrcnn3d_mask_selected = tf.gather_nd(mrcnn3d_masks, class_ids_aug) #[-1, NUM_POINT_PER_ROI]
    mrcnn3d_mask_selected = tf.reshape(mrcnn3d_mask_selected, [batch_size, num_rois, num_point_per_roi])

    return mrcnn3d_mask_selected

def unmold_segmentation(mrcnn3d_masks, rois, class_ids, pc_coord_cropped, pc):
    ''' Convert segmentation into point cloud label
    Inputs:
        mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
        rois: [B, NUM_ROIS, 6]
        class_ids: [B, NUM_ROIS]
        pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc: [B, NUM_POINT, 3]
    Returns:
        mrcnn3d_mask_unmolded: [B, NUM_ROIS, NUM_POINT]
    '''
    batch_size = mrcnn3d_masks.get_shape()[0].value
    num_rois = mrcnn3d_masks.get_shape()[1].value
    num_point_per_roi = mrcnn3d_masks.get_shape()[2].value
    num_category = mrcnn3d_masks.get_shape()[3].value
    num_point = pc.get_shape()[1].value

    mrcnn3d_masks = tf.reshape(mrcnn3d_masks, [-1, num_point_per_roi, num_category])
    mrcnn3d_masks = tf.transpose(mrcnn3d_masks, perm=[0, 2, 1]) # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI]
    class_ids = tf.cast(tf.reshape(class_ids, [-1]), tf.int32)
    class_ids_aug = tf.stack([tf.range(batch_size*num_rois, dtype=tf.int32), class_ids], 1)
    mrcnn3d_masks = tf.gather_nd(mrcnn3d_masks, class_ids_aug) #[-1, NUM_POINT_PER_ROI]

    # [B, NUM_ROIS, NUM_POINT, NUM_POINT_PER_ROI]
    dist = tf.reduce_sum(tf.square(tf.expand_dims(tf.expand_dims(pc, 1),3)-tf.expand_dims(pc_coord_cropped,2)), -1)
    min_idx = tf.argmin(dist, 3, output_type=tf.int32) # [B, NUM_ROIS, NUM_POINT]
    min_idx = tf.reshape(min_idx, [-1, num_point]) # [-1, NUM_POINT]
    min_idx_aug = tf.tile(tf.expand_dims(tf.range(batch_size*num_rois, dtype=tf.int32),-1), [1, num_point])
    min_idx_aug = tf.stack([tf.reshape(min_idx_aug, [-1]), tf.reshape(min_idx, [-1])], 1)
    mrcnn3d_mask_unmolded = tf.reshape(tf.gather_nd(mrcnn3d_masks, min_idx_aug), [batch_size, num_rois, num_point])

    # Mask out regions outside rois
    pc_aug = tf.expand_dims(pc, 1) # [B, 1, NUM_POINT, 3]
    rois_aug = tf.expand_dims(rois, 2) # [B, NUM_ROIS, 1, 6]
    roi_masks = tf.logical_and(pc_aug>=(rois_aug[:,:,:,:3]-rois_aug[:,:,:,3:]/2),
                               pc_aug<=(rois_aug[:,:,:,:3]+rois_aug[:,:,:,3:]/2))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,:,0], roi_masks[:,:,:,1]), roi_masks[:,:,:,2]) # [B, NUM_ROIS, NUM_POINT]
    roi_masks = tf.cast(roi_masks, tf.float32)

    mrcnn3d_mask_unmolded = tf.multiply(mrcnn3d_mask_unmolded, roi_masks)
    return mrcnn3d_mask_unmolded


def mrcnn_3d(pc, pc_ins, group_label, group_indicator, seg_label, seg_label_per_group, bbox_ins, config, is_training, mode='training', bn_decay=None):
    ''' Shape proposal generation
    Inputs:
        pc: [B, NUM_POINT, 3]
        color: [B, NUM_POINT, 3]
        pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys
        group_label: [B, NUM_POINT]
        group_indicator: [B, NUM_GROUP]
        seg_label: [B, NUM_POINT]
        bbox_ins: [B, NUM_GROUP, 6]
    Returns:
        
    '''
    assert mode in ['training', 'inference']
    if not config.USE_COLOR:
        color = None
    if 'SPN' in config.TRAIN_MODULE and mode=='training':
        end_points = shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins, config.NUM_CATEGORY, scope='shape_proposal_net', is_training=is_training, bn_decay=bn_decay, nsmp=config.NUM_SAMPLE, return_fullfea=False)
    else:
        end_points = shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins, config.NUM_CATEGORY, scope='shape_proposal_net', is_training=tf.constant(False), bn_decay=None, nsmp=config.NUM_SAMPLE, return_fullfea=True)
        end_points = dict_stop_gradient(end_points)
    if config.SHRINK_BOX:
        end_points['bbox_ins_pred'] = box_shrink(end_points['bbox_ins_pred'], pc)
    group_label_onehot = tf.one_hot(group_label, depth=config.NUM_GROUP, axis=-1) #[B, NUM_POINT, NUM_GROUP]
    # seg_label_per_group = tf.multiply(tf.cast(tf.expand_dims(seg_label,-1), tf.float32), group_label_onehot)
    # seg_label_per_group = tf.cast(tf.round(tf.divide(tf.reduce_sum(seg_label_per_group, 1),tf.reduce_sum(group_label_onehot, 1)+1e-8)), tf.int32) #[B, NUM_GROUP]

    if 'MRCNN' in config.TRAIN_MODULE or mode=='inference':
        SPN_NMS_MAX_SIZE = config.SPN_NMS_MAX_SIZE_TRAINING if mode == "training"\
            else config.SPN_NMS_MAX_SIZE_INFERENCE
        # 3D non maximum suppression - selected_indices: [B, M], spn_rois: [B, M, 6]
        selected_indices = tf.stop_gradient(tf.py_func(nms_3d, [end_points['bbox_ins_pred'], end_points['fb_prob'][:,:,1], config.SPN_PRE_NMS_LIMIT, SPN_NMS_MAX_SIZE, config.SPN_IOU_THRESHOLD, config.SPN_SCORE_THRESHOLD], tf.int32))
        spn_rois = gather_selection(end_points['bbox_ins_pred'], selected_indices, SPN_NMS_MAX_SIZE)

        if mode=='training':
            # Detection target generation - rois: [B, TRAIN_ROIS_PER_IMAGE, 6], target_class_ids: [B, TRAIN_ROIS_PER_IMAGE]
            # target_bbox: [B, TRAIN_ROIS_PER_IMAGE, 6], target_mask_selection_idx: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK]
            # target_mask: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK], all zero padded
            names = ["rois", "target_class_ids", "target_bbox", "target_mask_selection_idx", "target_mask"]
            outputs = batch_slice(
                [spn_rois, seg_label_per_group, bbox_ins, group_label_onehot, pc],
                lambda v, w, x, y, z: detection_target_gen(v, w, x, y, z, config),
                config.BATCH_SIZE, names=names)
            rois, target_class_ids, target_bbox, target_mask_selection_idx, target_mask = outputs

            # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### color
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], color), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_pn2_fea = pn2_fea_extractor(pc, None, 'pn2_fea_extractor', is_training, bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            sem_fea_full_l1 = tf_util.conv1d(end_points['sem_fea_full_l1'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn1', bn_decay=bn_decay)
            sem_fea_full_l2 = tf_util.conv1d(end_points['sem_fea_full_l2'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn2', bn_decay=bn_decay)
            sem_fea_full_l3 = tf_util.conv1d(end_points['sem_fea_full_l3'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn3', bn_decay=bn_decay)
            sem_fea_full_l4 = tf_util.conv1d(end_points['sem_fea_full_l4'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn4', bn_decay=bn_decay)
            sem_fea_full_l5 = tf_util.conv1d(end_points['sem_fea_full_l5'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn5', bn_decay=bn_decay)
            pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5), -1), 
                end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            # Classification and bbox refinement head
            mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY, 
                [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')
            # mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
            #     pc_fea_cropped, config.NUM_CATEGORY, 
            #     [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')

            # Mask prediction head
            mrcnn3d_mask = segmentation_head(pc_coord_cropped,
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY,
                [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')
            # mrcnn3d_mask = segmentation_head(pc_coord_cropped,
            #     pc_fea_cropped, config.NUM_CATEGORY,
            #     [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')
        elif mode=='inference':
            # rois: [B, NUM_ROIS, 6]
            names = ["rois", "mask_selection_idx"]
            outputs = batch_slice(
                [spn_rois, pc],
                lambda x, y: mask_selection_gen(x, y, SPN_NMS_MAX_SIZE, config, empty_removal=True),
                config.BATCH_SIZE, names=names)
            rois, mask_selection_idx = outputs

            # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_pn2_fea = pn2_fea_extractor(pc, color, 'pn2_fea_extractor', is_training, bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            sem_fea_full_l1 = tf_util.conv1d(end_points['sem_fea_full_l1'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn1', bn_decay=bn_decay)
            sem_fea_full_l2 = tf_util.conv1d(end_points['sem_fea_full_l2'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn2', bn_decay=bn_decay)
            sem_fea_full_l3 = tf_util.conv1d(end_points['sem_fea_full_l3'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn3', bn_decay=bn_decay)
            sem_fea_full_l4 = tf_util.conv1d(end_points['sem_fea_full_l4'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn4', bn_decay=bn_decay)
            sem_fea_full_l5 = tf_util.conv1d(end_points['sem_fea_full_l5'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn5', bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            #### generate fb_conf and sem_conf
            # pc: [B, N, 3], fb_prob: [B, nsmp, 2], pc_seed: [B, nsmp, 3] -> fb_prob: [B, N]
            midx = tf.argmin(tf.reduce_sum(tf.square(tf.expand_dims(pc, 2)-tf.expand_dims(end_points['pc_seed'],1)),-1), 2)
            midx_aug = tf.tile(tf.reshape(tf.range(config.BATCH_SIZE, dtype=tf.int64),[-1,1]), [1,config.NUM_POINT])
            midx_aug = tf.stack((tf.reshape(midx_aug, [-1]), tf.reshape(midx, [-1])), 1)
            fb_prob = tf.reshape(tf.gather_nd(end_points['fb_prob'], midx_aug), [config.BATCH_SIZE, config.NUM_POINT, 2])
            fb_prob = fb_prob[:,:,1]
            sem_prob = tf.nn.softmax(end_points['sem_class_logits'], -1) #[B, NUM_POINT, NUM_CATEGORY]

            ##### pn2
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea, tf.expand_dims(fb_prob, -1), sem_prob), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            # pc_fea_cropped, fb_prob_cropped, sem_prob_cropped = tf.split(pc_fea_cropped, [832, 1, config.NUM_CATEGORY], -1)
            ##### sem fpn fea
            pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5, tf.expand_dims(fb_prob, -1), sem_prob), -1), 
                end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            pc_fea_cropped, fb_prob_cropped, sem_prob_cropped = tf.split(pc_fea_cropped, [4*128+5*64, 1, config.NUM_CATEGORY], -1)
            
            fb_prob_cropped = tf.reduce_mean(tf.squeeze(fb_prob_cropped, -1),-1) # [B, NUM_ROIS]
            end_points['fb_prob_cropped'] = fb_prob_cropped
            sem_prob_cropped = tf.reduce_mean(sem_prob_cropped, 2) # [B, NUM_ROIS, NUM_CATEGORY]

            # Classification and bbox refinement head - mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
            # mrcnn3d_class: [B, NUM_ROIS, NUM_CATEGORY],
            # mrcnn3d_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6]
            mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY, 
                [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')

            midx = tf.argmax(mrcnn3d_class_logits, -1) # [B, NUM_ROIS]
            midx = tf.stack((tf.range(config.BATCH_SIZE*SPN_NMS_MAX_SIZE, dtype=tf.int64), tf.reshape(midx,[-1])), 1)
            sem_prob_cropped = tf.gather_nd(tf.reshape(sem_prob_cropped, [-1, config.NUM_CATEGORY]), midx)
            sem_prob_cropped = tf.reshape(sem_prob_cropped, [config.BATCH_SIZE, SPN_NMS_MAX_SIZE])
            end_points['sem_prob_cropped'] = sem_prob_cropped

            # Generate detections: [B, DETECTION_MAX_INSTANCES, (center_x, center_y, center_z, l, w, h, class_id, score)]
            # detections = batch_slice(
            #     [rois, mrcnn3d_class, mrcnn3d_bbox, pc],
            #     lambda x, y, w, z: refine_detections(x, y, w, z, config),
            #     config.BATCH_SIZE)
            detections = batch_slice(
                [rois, mrcnn3d_class, mrcnn3d_bbox, pc, fb_prob_cropped, sem_prob_cropped],
                lambda u, v, x, y, w, z: refine_detections(u, v, x, y, w, z, config),
                config.BATCH_SIZE)

            # Re-crop point cloud for mask prediction
            names = ["rois_final", "mask_selection_idx_final"]
            outputs = batch_slice(
                [detections[:,:,:6], pc],
                lambda x, y: mask_selection_gen(x, y, config.DETECTION_MAX_INSTANCES, config, empty_removal=False),
                config.BATCH_SIZE, names=names)
            rois_final, mask_selection_idx_final = outputs
            # Points cropping - pc_fea_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5), -1), 
                end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            # Mask prediction head
            mrcnn3d_mask = segmentation_head(pc_coord_cropped_final,
                tf.concat((pc_fea_cropped_final, pc_center_cropped_final), -1), config.NUM_CATEGORY,
                [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')

            # Unmold segmentation
            # mrcnn3d_mask_unmolded = unmold_segmentation(tf.nn.sigmoid(mrcnn3d_mask), rois_final, detections[:,:,6], pc_coord_cropped_final_unnormalized, pc)
            mrcnn3d_mask_selected = select_segmentation(tf.nn.sigmoid(mrcnn3d_mask), detections[:,:,6])

    # Update end_points
    end_points['group_label'] = group_label
    end_points['seg_label'] = seg_label
    end_points['seg_label_per_group'] = seg_label_per_group #[B, NUM_GROUP]
    end_points['bbox_ins'] = bbox_ins #[B, NUM_GROUP, 6]
    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        end_points['selected_indices'] = selected_indices #[B, SPN_NMS_MAX_SIZE]
        end_points['spn_rois'] = spn_rois #[B, SPN_NMS_MAX_SIZE, 6]
        end_points['rois'] = rois #[B, NUM_ROIS, 6]
        end_points['target_class_ids'] = target_class_ids #[B, NUM_ROIS]
        end_points['target_bbox'] = target_bbox #[B, NUM_ROIS, 6]
        end_points['target_mask_selection_idx'] = target_mask_selection_idx #[B, NUM_ROIS, NUM_POINT_PER_ROI]
        end_points['target_mask'] = target_mask #[B, NUM_ROIS, NUM_POINT_PER_ROI]
        end_points['mrcnn3d_class_logits'] = mrcnn3d_class_logits #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_class'] = mrcnn3d_class #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_bbox'] = mrcnn3d_bbox #[B, NUM_ROIS, NUM_CATEGORY, 6]
        end_points['mrcnn3d_mask'] = mrcnn3d_mask #[B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
    elif mode=='inference':
        end_points['selected_indices'] = selected_indices #[B, SPN_NMS_MAX_SIZE]
        end_points['spn_rois'] = spn_rois #[B, SPN_NMS_MAX_SIZE, 6]
        end_points['rois'] = rois #[B, NUM_ROIS, 6]
        end_points['mrcnn3d_class_logits'] = mrcnn3d_class_logits #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_class'] = mrcnn3d_class #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_bbox'] = mrcnn3d_bbox #[B, NUM_ROIS, NUM_CATEGORY, 6]
        end_points['detections'] = detections #[B, DETECTION_MAX_INSTANCES, 6+2]
        end_points['mrcnn3d_mask'] = mrcnn3d_mask #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NUM_CATEGORY]
        # end_points['mrcnn3d_mask_unmolded'] = mrcnn3d_mask_unmolded #[B, DETECTION_MAX_INSTANCES, NUM_POINT]
        end_points['mrcnn3d_mask_selected'] = mrcnn3d_mask_selected #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI]
        end_points['pc_coord_cropped_final_unnormalized'] = pc_coord_cropped_final_unnormalized #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
    return end_points#, alexnetmodel

def get_sem_class_loss(sem_class_logits, labels, smpw):
    '''
    Inputs:
       sem_class_logits: [B, npoint, num_category]
       label: [B, npoint]
    '''
    labels = tf.cast(labels, tf.int32)
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=sem_class_logits, weights=smpw)
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=sem_class_logits)
    loss = tf.divide(tf.reduce_sum(tf.multiply(loss, smpw)), tf.reduce_sum(smpw)+1e-8)
    return loss

def smooth_l1_loss(y_true, y_pred):
    """Implements Smooth-L1 loss.
    y_true and y_pred are typically: [N, 6], but could be any shape.
    """
    diff = tf.abs(y_true - y_pred)
    less_than_one = tf.cast(tf.less(diff, 1.0), tf.float32)
    loss = (less_than_one * 0.5 * diff**2) + (1 - less_than_one) * (diff - 0.5)
    return loss

def get_spn_class_loss(fb_logits, spn_match):
    '''
    Inputs:
       fb_logits: [B, nsmp, 2]
       spn_match: [B, nsmp]
    '''
    # Only postive and negative contribute to loss but not neutral
    fb_logits = tf.reshape(fb_logits, [-1, 2])
    spn_match = tf.reshape(spn_match, [-1])
    valid_mask = tf.not_equal(spn_match, 0)
    fb_logits = tf.boolean_mask(fb_logits, valid_mask)
    spn_match = tf.cast(tf.equal(spn_match, 1), tf.int32)
    spn_match = tf.boolean_mask(spn_match, valid_mask)
    loss = tf.cond(tf.size(spn_match)>0,
        lambda: tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=spn_match, logits=fb_logits)),
        lambda: tf.constant(0.0))

    return loss

def focal_loss(prediction_tensor, target_tensor, alpha=0.25, gamma=2):
    """Compute focal loss for predictions.
        Multi-labels Focal loss formula:
            FL = -alpha * (z-p)^gamma * log(p) -(1-alpha) * p^gamma * log(1-p)
                 ,which alpha = 0.25, gamma = 2, p = sigmoid(x), z = target_tensor.
    Args:
     prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
     target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
     alpha: A scalar tensor for focal loss alpha hyper-parameter
     gamma: A scalar tensor for focal loss gamma hyper-parameter
    Returns:
        loss: A (scalar) tensor representing the value of the loss function
    """
    sigmoid_p = tf.nn.sigmoid(prediction_tensor)
    zeros = array_ops.zeros_like(sigmoid_p, dtype=sigmoid_p.dtype)
    
    # For poitive prediction, only need consider front part loss, back part is 0;
    # target_tensor > zeros <=> z=1, so poitive coefficient = z - p.
    pos_p_sub = array_ops.where(target_tensor > zeros, target_tensor - sigmoid_p, zeros)
    
    # For negative prediction, only need consider back part loss, front part is 0;
    # target_tensor > zeros <=> z=1, so negative coefficient = 0.
    neg_p_sub = array_ops.where(target_tensor > zeros, zeros, sigmoid_p)
    per_entry_cross_ent = - alpha * (pos_p_sub ** gamma) * tf.log(tf.clip_by_value(sigmoid_p, 1e-8, 1.0)) \
                          - (1 - alpha) * (neg_p_sub ** gamma) * tf.log(tf.clip_by_value(1.0 - sigmoid_p, 1e-8, 1.0))
    return tf.reduce_sum(per_entry_cross_ent, -1)

def get_mrcnn3d_class_focal_loss(mrcnn3d_class_logits, gt_class_ids, roi_valid_mask, config):
    '''
    Inputs:
       mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
       gt_class_ids: [B, NUM_ROIS], zero padded
       roi_valid_mask: [B, NUM_ROIS]
    '''
    gt_class_ids_onehot = tf.one_hot(gt_class_ids, depth=config.NUM_CATEGORY, axis=-1)
    loss = focal_loss(mrcnn3d_class_logits, gt_class_ids_onehot)
    loss = tf.multiply(loss, roi_valid_mask)
    # loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(roi_valid_mask)+1e-8)
    loss = tf.reduce_sum(loss)

    return loss

def get_mrcnn3d_class_loss(mrcnn3d_class_logits, gt_class_ids, roi_valid_mask):
    '''
    Inputs:
       mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
       gt_class_ids: [B, NUM_ROIS], zero padded
       roi_valid_mask: [B, NUM_ROIS]
    '''
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=gt_class_ids, logits=mrcnn3d_class_logits)
    loss = tf.multiply(loss, roi_valid_mask)
    loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(roi_valid_mask)+1e-8)

    return loss

def get_mrcnn3d_bbox_loss(gt_bbox, gt_class_ids, pred_bbox, roi_valid_mask, num_category):
    '''
    Inputs:
       gt_bbox: [B, NUM_ROIS, 6]
       gt_class_ids: [B, NUM_ROIS], zero padded
       pred_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6]
       roi_valid_mask: [B, NUM_ROIS]
       num_category: scalar
    '''
    # Only foreground box contribute to the loss
    gt_bbox = tf.reshape(gt_bbox, [-1, 6])
    gt_class_ids = tf.reshape(gt_class_ids, [-1])
    pred_bbox = tf.reshape(pred_bbox, [-1, num_category, 6])
    roi_valid_mask = tf.reshape(roi_valid_mask, [-1])

    gt_selected_indices = tf.where(tf.logical_and(tf.greater(roi_valid_mask, 0), tf.greater(gt_class_ids, 0)))[:,0]
    gt_selected_indices = tf.cast(gt_selected_indices, tf.int32)
    pred_selected_indices = tf.concat((tf.reshape(gt_selected_indices, [-1,1]),
        tf.reshape(tf.gather(gt_class_ids, gt_selected_indices), [-1,1])), axis=1)

    gt_bbox = tf.gather(gt_bbox, gt_selected_indices, axis=0)
    pred_bbox = tf.gather_nd(pred_bbox, pred_selected_indices)

    loss = tf.cond(tf.size(gt_bbox)>0,
        lambda: tf.reduce_mean(tf.reduce_sum(smooth_l1_loss(y_true=gt_bbox, y_pred=pred_bbox),1),0),
        lambda: tf.constant(0.0))

    return loss

def get_mrcnn3d_mask_loss(gt_masks, gt_class_ids, pred_masks, roi_valid_mask, num_category, num_point_per_roi):
    '''
    Inputs:
       gt_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI]
       gt_class_ids: [B, NUM_ROIS], zero padded
       pred_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
       roi_valid_mask: [B, NUM_ROIS]
       num_category: scalar
       num_point_per_roi: scalar
    '''
    # Only foreground box contribute to the loss
    gt_masks = tf.reshape(gt_masks, [-1, num_point_per_roi])
    gt_class_ids = tf.reshape(gt_class_ids, [-1])
    pred_masks = tf.reshape(pred_masks, [-1, num_point_per_roi, num_category])
    pred_masks = tf.transpose(pred_masks, perm=[0,2,1])
    roi_valid_mask = tf.reshape(roi_valid_mask, [-1])

    gt_selected_indices = tf.where(tf.logical_and(tf.greater(roi_valid_mask, 0), tf.greater(gt_class_ids, 0)))[:,0]
    gt_selected_indices = tf.cast(gt_selected_indices, tf.int32)
    pred_selected_indices = tf.concat((tf.reshape(gt_selected_indices, [-1,1]),
        tf.reshape(tf.gather(gt_class_ids, gt_selected_indices), [-1,1])), axis=1)

    gt_masks = tf.gather(gt_masks, gt_selected_indices, axis=0) # [N, NUM_POINT_PER_ROI]
    gt_masks = tf.cast(gt_masks, tf.float32)
    pred_masks = tf.gather_nd(pred_masks, pred_selected_indices) # [N, NUM_POINT_PER_ROI]

    loss = tf.cond(tf.size(gt_masks)>0,
        lambda: tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=gt_masks, logits=pred_masks)),
        lambda: tf.constant(0.0))

    return loss

def get_loss(end_points, config, alpha, smpw, mode='training'):
    batch_size = config.BATCH_SIZE
    nsmp = config.NUM_SAMPLE
    nsmp_ins = config.NUM_POINT_INS
    bbox_size = tf.reduce_max(end_points['pc_ins_centered_seed'], axis=2, keep_dims=True)-tf.reduce_min(end_points['pc_ins_centered_seed'], axis=2, keep_dims=True)
    radius = 1e-8 + tf.sqrt(tf.reduce_sum(tf.square(bbox_size/2), axis=-1, keep_dims=True)) # [B, nsmp, 1, 1]
    # radius = 1e-8 + tf.reduce_max(tf.sqrt(tf.reduce_sum(tf.square(end_points['pc_ins_centered_seed']), axis=-1, keep_dims=True)), axis=2, keep_dims=True) # [B, nsmp, 1, 1]
    shift_gt_seed = end_points['pc_ins_center_seed']-tf.expand_dims(end_points['pc_seed'],2) # [B, nsmp, 1, 3]
    shift_dist = tf.sqrt(tf.reduce_sum(tf.square(shift_gt_seed), 3, keep_dims=True)+1e-8)
    shift_gt_seed_normalized_4d = tf.concat((tf.divide(shift_gt_seed, shift_dist), tf.divide(shift_dist, radius)), -1)
    shift_pred_seed_normalized_4d = tf.concat((tf.expand_dims(end_points['shift_pred_seed_4d'],2)[:,:,:,:3], tf.divide(tf.expand_dims(end_points['shift_pred_seed_4d'],2)[:,:,:,3:], radius)), -1)

    # Fg/Bg loss, spn_match: [B, nsmp]
    fb_score_gt = tf.squeeze(gather_selection(tf.expand_dims(end_points['group_label'],-1), end_points['ind_seed'], end_points['ind_seed'].get_shape()[1].value),-1)
    fb_score_gt = tf.cast(tf.greater(fb_score_gt,0), tf.float32) # [B, nsmp]
    spn_match = batch_slice(
        [end_points['bbox_ins_pred'], fb_score_gt, tf.cast(tf.greater(end_points['seg_label_per_group'],0),tf.float32), end_points['bbox_ins']],
        lambda w, x, y, z: spn_target_gen(w, x, y, z),
        batch_size, names=["spn_match"])
    spn_match = tf.stop_gradient(tf.cast(spn_match, tf.int32))
    end_points['spn_match'] = spn_match
    spn_class_loss = get_spn_class_loss(end_points['fb_logits'], spn_match)
    # spn_class_loss = tf.constant(0.0)

    # # Reconstruction loss
    # pc_ins_pred = end_points['pc_ins_pred']
    # pc_ins_pred_normalized = tf.reshape(tf.div(pc_ins_pred, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    # pc_ins_gt_normalized = tf.reshape(tf.div(end_points['pc_ins_centered_seed']+shift_gt_seed, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    # recon_valid_mask = tf.reshape(fb_score_gt, [-1])
    # recon_valid_mask = tf.stop_gradient(recon_valid_mask)
    # dists_forward,_,dists_backward,_ = tf_nndistance.nn_distance(pc_ins_pred_normalized, pc_ins_gt_normalized)
    # recons_loss = tf.reduce_mean(dists_forward+dists_backward, axis=-1) # B*nsmp
    # recons_loss = tf.divide(tf.reduce_sum(tf.multiply(recons_loss, recon_valid_mask)),
    #     tf.reduce_sum(recon_valid_mask)+1e-8)
    # # recons_loss = tf.constant(0.0)

    # Reconstruction loss & confidence loss
    pc_ins_pred = end_points['pc_ins_pred']
    pc_ins_pred_normalized = tf.reshape(tf.div(pc_ins_pred, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    pc_ins_gt_normalized = tf.reshape(tf.div(end_points['pc_ins_centered_seed']+shift_gt_seed, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    recon_valid_mask = tf.reshape(fb_score_gt, [-1])
    recon_valid_mask = tf.stop_gradient(recon_valid_mask)
    end_points['recon_valid_mask'] = recon_valid_mask
    dists_forward,_,dists_backward,_ = tf_nndistance.nn_distance(pc_ins_pred_normalized, pc_ins_gt_normalized)
    # dists_forward [B*nsmp, nsmp_ins]
    dist_thres = 0.04
    width = 0.03
    dists = tf.reshape(dists_forward, [batch_size, nsmp, nsmp_ins])

    conf_gt = tf.reshape(1.0 / (1.0 + tf.exp(-5*(dists-dist_thres)/width)), [-1, nsmp_ins]) # [B*nsmp, nsmp_ins]
    end_points['conf_gt'] = conf_gt
    conf_pred = tf.reshape(end_points['pc_ins_conf_pred'], [-1, nsmp_ins]) # [B*nsmp, nsmp_ins]
    conf_loss = tf.reduce_mean(smooth_l1_loss(conf_gt, conf_pred), axis=-1) # [B*nsmp]
    conf_loss = tf.divide(tf.reduce_sum(tf.multiply(conf_loss, recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)

    recons_loss = tf.reduce_mean(dists_forward+dists_backward, axis=-1) # B*nsmp
    recons_loss = tf.divide(tf.reduce_sum(tf.multiply(recons_loss, recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)
    # recons_loss = tf.constant(0.0)

    # Shift loss
    shift_loss = tf.reduce_sum(smooth_l1_loss(shift_gt_seed_normalized_4d, shift_pred_seed_normalized_4d), axis=-1)
    shift_loss = tf.reduce_sum(tf.multiply(tf.reshape(shift_loss, [-1]), recon_valid_mask))
    shift_loss = tf.divide(shift_loss, tf.reduce_sum(recon_valid_mask)+1e-8)
    # shift_loss = tf.constant(0.0)

    # Sem loss
    ind_sem = end_points['ind_sem']
    nsmp_sem = ind_sem.get_shape()[1].value
    sem_labels = end_points['seg_label']
    ind_sem_aug = tf.tile(tf.reshape(tf.range(batch_size),[-1, 1]), [1, nsmp_sem])
    ind_sem_aug = tf.concat( (tf.reshape(ind_sem_aug, [-1, 1]), tf.reshape(ind_sem, [-1, 1])), 1 )
    sem_labels = tf.reshape(tf.gather_nd(sem_labels, ind_sem_aug), [batch_size, nsmp_sem])
    smpw = tf.reshape(tf.gather_nd(smpw, ind_sem_aug), [batch_size, nsmp_sem])
    sem_labels = tf.cast(sem_labels, tf.int32)
    sem_loss = tf.losses.sparse_softmax_cross_entropy(labels=sem_labels, logits=end_points['sem_class_logits'], weights=smpw)
    end_points['sem_labels'] = sem_labels
    # sem_loss = tf.constant(0.0)

    # KL loss
    mean = end_points['mean'] # [B, nsmp, 256]
    log_var = end_points['log_var']
    cmean = end_points['cmean']
    clog_var = end_points['clog_var']
    kl_loss = 0.5 * tf.reduce_mean( log_var - clog_var + (tf.exp(clog_var) + (mean-cmean)**2)/tf.exp(log_var) - 1.0, 2) 
    kl_loss = tf.divide(tf.reduce_sum(tf.multiply(tf.reshape(kl_loss, [-1]), recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)
    # kl_loss = tf.constant(0.0)

    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        # MRCNN3D classification loss
        roi_valid_mask = tf.cast(tf.not_equal(tf.reduce_sum(tf.abs(end_points['rois']), axis=-1),0), tf.float32)
        mrcnn3d_class_loss = get_mrcnn3d_class_loss(end_points['mrcnn3d_class_logits'], end_points['target_class_ids'], roi_valid_mask)
        # mrcnn3d_class_loss = get_mrcnn3d_class_focal_loss(end_points['mrcnn3d_class_logits'], end_points['target_class_ids'], roi_valid_mask, config)

        # MRCNN3D bbox loss
        mrcnn3d_bbox_loss = get_mrcnn3d_bbox_loss(end_points['target_bbox'], end_points['target_class_ids'], end_points['mrcnn3d_bbox'], roi_valid_mask, config.NUM_CATEGORY)

        # MRCNN3D mask loss
        mrcnn3d_mask_loss = get_mrcnn3d_mask_loss(end_points['target_mask'], end_points['target_class_ids'], end_points['mrcnn3d_mask'], roi_valid_mask, config.NUM_CATEGORY, config.NUM_POINT_INS_MASK)
    
    if 'SPN' in config.TRAIN_MODULE:
        loss = kl_loss * alpha + recons_loss + shift_loss + spn_class_loss + sem_loss + conf_loss
        if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
            loss += mrcnn3d_class_loss + mrcnn3d_bbox_loss + mrcnn3d_mask_loss
    elif 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        loss = mrcnn3d_class_loss + mrcnn3d_bbox_loss + mrcnn3d_mask_loss
    else:
        loss = tf.constant(0.0)

    # Store end_points
    end_points['spn_class_loss'] = spn_class_loss
    end_points['recons_loss'] = recons_loss
    end_points['shift_loss'] = shift_loss
    end_points['conf_loss'] = conf_loss
    end_points['sem_loss'] = sem_loss
    end_points['kl_loss'] = kl_loss
    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        end_points['mrcnn3d_class_loss'] = mrcnn3d_class_loss
        end_points['mrcnn3d_bbox_loss'] = mrcnn3d_bbox_loss
        end_points['mrcnn3d_mask_loss'] = mrcnn3d_mask_loss
    end_points['loss'] = loss
    
    return loss, end_points


if __name__=='__main__':
    # boxes = np.array([[0,0,0,1.,1.,1.],[0.1,0.1,0.1,1.,1.,1.],[1.,1.,1.,1.,1.,1.],[0.8,0.8,0.8,1.,1.,1.]])
    # scores = np.array([1.0,1.0,0.5,1.0])
    # boxes = np.reshape(boxes,[1,4,6])
    # scores = np.reshape(scores,[1,4])
    # select_idx = nms_3d(boxes, scores, 4, iou_threshold=0.3, score_threshold=0.0)
    # print(select_idx)
    myconfig = config.Config()
    with tf.Graph().as_default():
        pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, bbox_ins_pl = placeholder_inputs(myconfig)
        end_points = mrcnn_3d(pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, bbox_ins_pl, myconfig, tf.constant(True), bn_decay=None)
        loss, end_points = get_loss(end_points, myconfig, 1.0)
        print(end_points)

<filename>models/model_mrcnn3d.py
import tensorflow as tf
import numpy as np
import math
import sys
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import tf_util
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/nn_distance'))
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/sampling'))
sys.path.append(os.path.join(ROOT_DIR, 'tf_ops/grouping'))
import tf_nndistance
from tf_sampling import farthest_point_sample, gather_point
from tf_grouping import query_ball_point, group_point
from pointnet_util import pointnet_fp_module, pointnet_sa_module
import config
from sklearn.neighbors import NearestNeighbors
from tensorflow.python.ops import array_ops

def placeholder_inputs(config):
    pc_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_POINT, 3))
    pc_ins_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_GROUP, config.NUM_POINT_INS, 3))
    group_label_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_POINT))
    group_indicator_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_GROUP))
    seg_label_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_POINT))
    seg_label_per_group_pl = tf.placeholder(tf.int32, shape=(config.BATCH_SIZE, config.NUM_GROUP))
    bbox_ins_pl = tf.placeholder(tf.float32, shape=(config.BATCH_SIZE, config.NUM_GROUP, 6))
    return pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, seg_label_per_group_pl, bbox_ins_pl

def multi_encoding_net(xyz, points, npoint, radius_list, nsample_list, mlp_list, mlp_list2, is_training, bn_decay, scope, bn=True, use_xyz=False, output_shift=False, shift_pred=None, fps_idx=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint: int32 -- #points sampled in farthest point sampling
            radius: list of float32 -- search radius in local region
            nsample: list of int32 -- how many points in each local region
            mlp: list of list of int32 -- output size for MLP on each point
            use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features
        Return:
            new_xyz: (batch_size, npoint, 3) TF tensor
            new_points: (batch_size, npoint, mlp2[-1]) TF tensor
            shift_pred: (batch_size, npoint, 3) TF tensor
            fps_idx: (batch_size, npoint) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        if fps_idx is None:
            fps_idx = farthest_point_sample(npoint, xyz) # (batch_size, npoint)
        new_xyz = gather_point(xyz, fps_idx) # (batch_size, npoint, 3)
        new_points_list = []
        group_xyz_list = []
        for i in range(len(radius_list)):
            radius = radius_list[i]
            nsample = nsample_list[i]
            idx, pts_cnt = query_ball_point(radius, nsample, xyz, new_xyz)
            grouped_xyz = group_point(xyz, idx)
            grouped_xyz -= tf.tile(tf.expand_dims(new_xyz, 2), [1,1,nsample,1]) # [B, nseed, nsmp, 3]
            if shift_pred is not None:
                grouped_xyz -= tf.tile(tf.expand_dims(shift_pred, 2), [1,1,nsample,1])
            if points is not None:
                grouped_points = group_point(points, idx)
                if use_xyz:
                    grouped_points = tf.concat([grouped_points, grouped_xyz], axis=-1)
            else:
                grouped_points = grouped_xyz
            for j,num_out_channel in enumerate(mlp_list[i]):
                grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                                padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                                scope='conv_prev_%d_%d'%(i,j), bn_decay=bn_decay)
            new_points_list.append(tf.reduce_max(grouped_points, axis=[2]))
        new_points = tf.concat(new_points_list, axis=-1) # (batch_size, npoint, \sum_k{mlp[k][-1]})
        for i,num_out_channel in enumerate(mlp_list2):
            new_points = tf_util.conv1d(new_points, num_out_channel, 1,
                                               padding='VALID', stride=1, bn=bn, is_training=is_training,
                                               scope='conv_post_%d'%i, bn_decay=bn_decay)
        if output_shift:
            shift_pred = tf_util.conv1d(new_points, 4, 1,
                                        padding='VALID', stride=1, scope='conv_shift_pred', activation_fn=None)
        return new_xyz, new_points, shift_pred, fps_idx

def shift_pred_net(xyz, points, npoint_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
        Return:
            pc_seed: (batch_size, npoint_seed, 3) TF tensor
            shift_pred_seed_4d: (batch_size, npoint_seed, 4) TF tensor
            ind_seed: (batch_size, npoint_seed) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        ind_seed = farthest_point_sample(npoint_seed, xyz) # (batch_size, npoint_seed)
        pc_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None # do not use color for shift prediction

        if return_fullfea:
            new_xyz = tf.concat((pc_seed, xyz), 1)
        else:
            new_xyz = pc_seed

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=0.8, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')

        # Feature Propagation layers
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 4, 1,
                             padding='VALID', stride=1, scope='conv_shift_pred', activation_fn=None)
        if return_fullfea:
            shift_pred_seed_4d, shift_pred_full_4d = tf.split(net, [npoint_seed, num_point], axis=1)
            end_points['shift_pred_full_4d'] = shift_pred_full_4d
        else:
            shift_pred_seed_4d = net

        end_points['pc_seed'] = pc_seed
        end_points['shift_pred_seed_4d'] = shift_pred_seed_4d
        end_points['ind_seed'] = ind_seed

        return end_points

def sem_net(xyz, points, npoint_sem, num_category, ind_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        npoint_seed = ind_seed.get_shape()[1].value

        ind_sem = farthest_point_sample(npoint_sem, xyz) # (batch_size, npoint_sem)
        new_xyz_sem = gather_point(xyz, ind_sem) # (batch_size, npoint_sem, 3)
        end_points['ind_sem'] = ind_sem
        new_xyz_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)

        if return_fullfea:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem, xyz), 1)
        else:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem), 1)
        
        l0_xyz = xyz
        l0_points = None

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=128, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=64, radius=0.8, nsample=32, mlp=[256,256,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        l5_xyz, l5_points, l5_indices = pointnet_sa_module(l4_xyz, l4_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer5')
        # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')

        # with FPN
        if return_fullfea:
            end_points['sem_fea_full_l5'] = pointnet_fp_module(xyz, l5_xyz, points, l5_points, [], is_training, bn_decay, scope='fa_layer0_fpn')
            end_points['sem_fea_full_l4'] = pointnet_fp_module(xyz, l4_xyz, points, l4_points, [], is_training, bn_decay, scope='fa_layer1_fpn')
            end_points['sem_fea_full_l3'] = pointnet_fp_module(xyz, l3_xyz, points, l3_points, [], is_training, bn_decay, scope='fa_layer2_fpn')
            end_points['sem_fea_full_l2'] = pointnet_fp_module(xyz, l2_xyz, points, l2_points, [], is_training, bn_decay, scope='fa_layer3_fpn')
            end_points['sem_fea_full_l1'] = pointnet_fp_module(xyz, l1_xyz, points, l1_points, [], is_training, bn_decay, scope='fa_layer4_fpn')

        # Feature Propagation layers
        l4_points = pointnet_fp_module(l4_xyz, l5_xyz, l4_points, l5_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        if return_fullfea:
            sem_fea_seed, sem_fea, sem_fea_full = tf.split(net, [npoint_seed, npoint_sem, num_point], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            end_points['sem_fea_full'] = sem_fea_full
            net = end_points['sem_fea_full']
        else:
            sem_fea_seed, sem_fea = tf.split(net, [npoint_seed, npoint_sem], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            net = end_points['sem_fea']

        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def sem_net_simp(xyz, points, num_category, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None
        end_points = {}

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.1, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.2, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=128, radius=0.4, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=64, radius=0.8, nsample=32, mlp=[256,256,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4')
        l5_xyz, l5_points, l5_indices = pointnet_sa_module(l4_xyz, l4_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer5')

        # Feature Propagation layers
        l4_points = pointnet_fp_module(l4_xyz, l5_xyz, l4_points, l5_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer2')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer3')
        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer4')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points


def sem_net2(xyz, points, npoint_sem, num_category, ind_seed, end_points, scope, is_training, bn_decay=None, return_fullfea=False):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        npoint_seed = ind_seed.get_shape()[1].value

        ind_sem = farthest_point_sample(npoint_sem, xyz) # (batch_size, npoint_sem)
        new_xyz_sem = gather_point(xyz, ind_sem) # (batch_size, npoint_sem, 3)
        end_points['ind_sem'] = ind_sem
        new_xyz_seed = gather_point(xyz, ind_seed) # (batch_size, npoint_seed, 3)

        if return_fullfea:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem, xyz), 1)
        else:
            new_xyz = tf.concat((new_xyz_seed, new_xyz_sem), 1)
        
        l0_xyz = xyz
        l0_points = None

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=512, radius=0.2, nsample=64, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=128, radius=0.4, nsample=64, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3')

        # with FPN
        if return_fullfea:
            end_points['sem_fea_full_l3'] = pointnet_fp_module(xyz, l3_xyz, points, l3_points, [], is_training, bn_decay, scope='fa_layer2_fpn')
            end_points['sem_fea_full_l2'] = pointnet_fp_module(xyz, l2_xyz, points, l2_points, [], is_training, bn_decay, scope='fa_layer3_fpn')
            end_points['sem_fea_full_l1'] = pointnet_fp_module(xyz, l1_xyz, points, l1_points, [], is_training, bn_decay, scope='fa_layer4_fpn')

        # Feature Propagation layers
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer2')
        l0_points = pointnet_fp_module(new_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer3')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        if return_fullfea:
            sem_fea_seed, sem_fea, sem_fea_full = tf.split(net, [npoint_seed, npoint_sem, num_point], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea
            end_points['sem_fea_full'] = sem_fea_full
        else:
            sem_fea_seed, sem_fea = tf.split(net, [npoint_seed, npoint_sem], axis=1)
            end_points['sem_fea_seed'] = sem_fea_seed
            end_points['sem_fea'] = sem_fea

        # net = end_points['sem_fea_full']
        net = end_points['sem_fea']
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def sem_net2_simp(xyz, points, num_category, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
            npoint_sem: int32 -- #points to sample for fast training
            num_category: int32 -- #output category
            ind_seed: (batch_size, npoint_seed) sampling index of seed points
        Return:
            sem_fea_seed: (batch_size, npoint_seed, nfea)
            sem_fea: (batch_size, npoint_sem, nfea)
            sem_fea_full: (batch_size, ndataset, nfea)
            sem_class_logits: (batch_size, npoint_sem, num_category)
            ind_sem: (batch_size, npoint_sem) sampling index of sem points
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = None
        end_points = {}

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=512, radius=0.2, nsample=64, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=128, radius=0.4, nsample=64, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer3')

        # Feature Propagation layers
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,256], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,128], is_training, bn_decay, scope='fa_layer2')
        l0_points = pointnet_fp_module(l0_xyz, l1_xyz, None, l1_points, [128,128,128], is_training, bn_decay, scope='fa_layer3')

        # FC layers
        net = tf_util.conv1d(l0_points, 128, 1, padding='VALID', bn=True, is_training=is_training, scope='fc1', bn_decay=bn_decay)
        net = tf_util.dropout(net, keep_prob=0.5, is_training=is_training, scope='dp1')
        sem_class_logits = tf_util.conv1d(net, num_category, 1, padding='VALID', activation_fn=None, scope='fc2')
        end_points['sem_class_logits'] = sem_class_logits

        return end_points

def pn2_fea_extractor(xyz, points, scope, is_training, bn_decay=None):
    ''' Encode multiple context.
        Input:
            xyz: (batch_size, ndataset, 3) TF tensor
            points: (batch_size, ndataset, channel) TF tensor
        Return:
            new_points: (batch_size, ndataset, channel_out) TF tensor
    '''
    with tf.variable_scope(scope) as sc:
        batch_size = xyz.get_shape()[0].value
        num_point = xyz.get_shape()[1].value
        l0_xyz = xyz
        l0_points = points

        # Layer 1
        l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=48, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
        l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=48, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
        l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=48, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3')
        l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=None, radius=None, nsample=None, mlp=[256,512,1024], mlp2=None, group_all=True, is_training=is_training, bn_decay=bn_decay, scope='layer4')


        # Feature Propagation layers
        l3_points = pointnet_fp_module(l3_xyz, l4_xyz, l3_points, l4_points, [256,256], is_training, bn_decay, scope='fa_layer0')
        l2_points = pointnet_fp_module(l2_xyz, l3_xyz, l2_points, l3_points, [256,128], is_training, bn_decay, scope='fa_layer1')
        l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [128,64], is_training, bn_decay, scope='fa_layer2')
        new_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, [64,64,64], is_training, bn_decay, scope='fa_layer3')

        return new_points


def single_encoding_net(pc, mlp_list, mlp_list2, scope, is_training, bn_decay):
    ''' The encoding network for instance
    Input:
        pc: [B, N, 3]
    Return:
        net: [B, nfea]
    '''
    with tf.variable_scope(scope) as myscope:
        net = tf.expand_dims(pc, 2)
        for i,num_out_channel in enumerate(mlp_list):
            net = tf_util.conv2d(net, num_out_channel, [1,1],
                                 padding='VALID', stride=[1,1],
                                 bn=True, is_training=is_training,
                                 scope='conv%d'%i, bn_decay=bn_decay)
        net = tf.reduce_max(net, axis=[1])
        net = tf.squeeze(net, 1)
        for i,num_out_channel in enumerate(mlp_list2):
            net = tf_util.fully_connected(net, num_out_channel, bn=True, is_training=is_training,
                                          scope='fc%d'%i, bn_decay=bn_decay)
        return net

def fea_trans_net(input_fea, mlp_list, scope, is_training, bn_decay):
    with tf.variable_scope(scope) as myscope:
        net = input_fea
        nlayer = len(mlp_list)
        for i,num_out_channel in enumerate(mlp_list):
            if i<nlayer-1:
                net = tf_util.conv1d(net, num_out_channel, 1, padding='VALID', bn=True, is_training=is_training,
                                     scope='conv%d'%i, bn_decay=bn_decay)
            else:
                net = tf_util.conv1d(net, num_out_channel, 1, padding='VALID', activation_fn=None, scope='conv%d'%i)
        return net

def sample(mean, log_var):
    # Sample z
    z = mean + tf.exp(log_var/2.0) * tf.random_normal(tf.shape(mean), 0, 1, dtype=tf.float32)
    return z

def decoding_net(feat, num_point, scope, is_training, bn_decay):
    ''' The decoding network for shape generation
    Input:
        feat: [B, nsmp, nfea]
    Return:
        pc: [B, nsmp, num_point, 3]
    '''
    with tf.variable_scope(scope) as myscope:
        nsmp = feat.get_shape()[1].value
        nfea = feat.get_shape()[2].value
        feat = tf.reshape(feat, [-1, nfea])
        # UPCONV Decoder
        if num_point<=3072 and num_point>1536:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[5,5], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 1024
        elif num_point<=1536 and num_point>896:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 484
        elif num_point<=896 and num_point>384:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 3, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv4', activation_fn=None)
            num_point_conv = 256
        else:
            raise('Exception')
        pc_upconv = tf.reshape(net, [-1, num_point_conv, 3])

        num_point_fc = num_point - num_point_conv
        # FC Decoder
        net = tf_util.fully_connected(feat, 512, bn=True, is_training=is_training, scope='de_fc2', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, scope='de_fc3', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, num_point_fc*3, activation_fn=None, scope='de_fc4')
        pc_fc = tf.reshape(net, [-1, num_point_fc, 3])
        # Merge
        pc = tf.concat([pc_upconv, pc_fc], axis=1)
        pc = tf.reshape(pc, [-1, nsmp, num_point, 3])
        return pc

def decoding_conf_net(feat, num_point, scope, is_training, bn_decay):
    ''' The decoding network for shape generation & confidence generation for each generated point
    Input:
        feat: [B, nsmp, nfea]
    Return:
        pc: [B, nsmp, num_point, 3]
    '''
    with tf.variable_scope(scope) as myscope:
        nsmp = feat.get_shape()[1].value
        nfea = feat.get_shape()[2].value
        feat = tf.reshape(feat, [-1, nfea])
        # UPCONV Decoder
        if num_point<=3072 and num_point>1536:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[5,5], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 1024
        elif num_point<=1536 and num_point>896:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[2,2], stride=[1,1], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[3,3], padding='VALID', scope='upconv4', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv5', activation_fn=None)
            num_point_conv = 484
        elif num_point<=896 and num_point>384:
            conv_feat = tf.expand_dims(tf.expand_dims(feat, 1),1)
            net = tf_util.conv2d_transpose(conv_feat, 512, kernel_size=[3,3], stride=[1,1], padding='VALID', scope='upconv1', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 256, kernel_size=[3,3], stride=[2,2], padding='VALID', scope='upconv2', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 128, kernel_size=[4,4], stride=[2,2], padding='VALID', scope='upconv3', bn=True, bn_decay=bn_decay, is_training=is_training)
            net = tf_util.conv2d_transpose(net, 4, kernel_size=[1,1], stride=[1,1], padding='VALID', scope='upconv4', activation_fn=None)
            num_point_conv = 256
        else:
            raise('Exception')
        pc_upconv = tf.reshape(net, [-1, num_point_conv, 4])

        num_point_fc = num_point - num_point_conv
        # FC Decoder
        net = tf_util.fully_connected(feat, 512, bn=True, is_training=is_training, scope='de_fc2', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, scope='de_fc3', bn_decay=bn_decay)
        net = tf_util.fully_connected(net, num_point_fc*4, activation_fn=None, scope='de_fc4')
        pc_fc = tf.reshape(net, [-1, num_point_fc, 4])
        # Merge
        pc = tf.concat([pc_upconv, pc_fc], axis=1)
        pc = tf.reshape(pc, [-1, nsmp, num_point, 4])
        pc_conf = pc[:,:,:,3]
        pc_xyz = pc[:,:,:,:3]
        return pc_xyz, pc_conf

def shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins_gt, num_category, scope, is_training, bn_decay=None, nsmp=128, return_fullfea=False):
    ''' Shape proposal generation
    Inputs:
        pc: [B, NUM_POINT, 3]
        pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys
        group_label: [B, NUM_POINT]
        group_indicator: [B, NUM_GROUP]
        bbox_ins_gt: [B, NUM_GROUP, 6], only used during training
    Returns:
        fb_logits: [B, NUM_SAMPLE, 2] confidence logits (before softmax)
        fb_prob: [B, NUM_SAMPLE, 2] confidence probabilities
        bbox_ins: [B, NUM_SAMPLE, (x, y, z, l, w, h)]
        entity_fea: [B, NUM_POINT, nfea] entity feature for each point
        center_pos: [B, NUM_POINT, 3] center coordinate for each point, in world coord sys
    '''
    with tf.variable_scope(scope) as myscope:
        # Parameter extraction
        batch_size = pc.get_shape()[0].value
        ngroup = pc_ins.get_shape()[1].value
        nsmp_ins = pc_ins.get_shape()[2].value
        end_points = {}

        # Shift prediction, ind_seed [B, nsmp], shift_pred_seed [B, nsmp, 3]
        end_points = shift_pred_net(pc, None, nsmp, end_points, 'shift_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        pc_seed = end_points['pc_seed']
        shift_pred_seed_4d = end_points['shift_pred_seed_4d']
        ind_seed = end_points['ind_seed']
        shift_pred_seed = tf.multiply(shift_pred_seed_4d[:,:,:3], shift_pred_seed_4d[:,:,3:])

        # Semantic prediction, sem_fea_seed [B, nsmp, nfea]
        end_points = sem_net(pc, None, 10000 if return_fullfea else 1024, num_category, ind_seed, end_points, 'sem_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        # end_points = sem_net2(pc, None, 4096, num_category, ind_seed, end_points, 'sem_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea)
        sem_fea_seed = end_points['sem_fea_seed']

        # Encode instance, pcfea_ins_centered [B, ngroup, nfea_ins], pc_ins_center [B, ngroup, 1, 3]
        pc_ins_center = (tf.reduce_max(pc_ins, 2, keep_dims=True)+tf.reduce_min(pc_ins, 2, keep_dims=True))/2 # [B, ngroup, 1, 3] -> requires random padding for pc_ins generation
        pc_ins_centered = pc_ins-pc_ins_center
        idx = tf.where(tf.greater(group_indicator, 0))
        pc_ins_centered_list = tf.gather_nd(pc_ins_centered, idx)
        pcfea_ins_centered_list = single_encoding_net(pc_ins_centered_list, [64, 256, 512], [256], 'instance_encoder', is_training, bn_decay)
        nfea_ins = pcfea_ins_centered_list.get_shape()[1].value
        pcfea_ins_centered = tf.scatter_nd(tf.cast(idx,tf.int32), pcfea_ins_centered_list, tf.constant([batch_size, ngroup, nfea_ins])) # [B, ngroup, nfea_ins]
        
        # Collect instance feature for seed points [B, nsmp, nfea_seed]
        idx = tf.where(tf.greater_equal(ind_seed,0))
        ind_seed_aug = tf.concat((tf.expand_dims(tf.cast(idx[:,0],tf.int32),-1),tf.reshape(ind_seed,[-1,1])),1)
        group_label_seed = tf.reshape(tf.gather_nd(group_label, ind_seed_aug), [-1, nsmp]) # [B, nsmp]
        idx = tf.where(tf.greater_equal(group_label_seed,0))
        group_label_seed_aug = tf.concat((tf.expand_dims(tf.cast(idx[:,0],tf.int32),-1),tf.reshape(group_label_seed,[-1,1])),1)
        pcfea_ins_seed = tf.reshape(tf.gather_nd(pcfea_ins_centered, group_label_seed_aug), [-1, nsmp, nfea_ins])
        pc_ins_centered_seed = tf.reshape(tf.gather_nd(pc_ins_centered, group_label_seed_aug), [-1, nsmp, nsmp_ins, 3])
        pc_ins_center_seed = tf.reshape(tf.gather_nd(pc_ins_center, group_label_seed_aug), [-1, nsmp, 1, 3])
       
        # Encode context, pcfea_seed [B, nsmp, nfea_context]
        # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.25,0.5,1.0,1.5], [128,256,256,512], [[64,64,128], [64,128,256], [64,128,256], [64,128,256]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)
        _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.25,0.5,1.0,1.5], [96,192,192,384], [[64,64,128], [64,128,128], [64,128,128], [64,128,128]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)
        # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.4,0.8,1.5], [192,256,384], [[64,128,192], [64,128,192], [64,128,192]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed)

        # Compute foreground/background score [B, nsmp, 2]
        fb_logits = fea_trans_net(pcfea_seed, [256, 64, 2], 'fb_logits', is_training, bn_decay)
        # fb_logits = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 64, 2], 'fb_logits', is_training, bn_decay)
        fb_prob = tf.nn.softmax(fb_logits, -1)

        #### CVAE
        # Compute mu and sigma [B, nsmp, 512]
        # mu_sigma_c = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 512, 512], 'mu_sigma_c', is_training, bn_decay)
        # mu_sigma_x = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed, pcfea_ins_seed), axis=-1), [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)
        mu_sigma_c = fea_trans_net(pcfea_seed, [256, 512, 512], 'mu_sigma_c', is_training, bn_decay)
        mu_sigma_x = fea_trans_net(tf.concat((pcfea_seed, pcfea_ins_seed), axis=-1), [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)
        # mu_sigma_x = fea_trans_net(pcfea_ins_seed, [256, 512, 512], 'mu_sigma_x', is_training, bn_decay)

        # Sample z [B, nsmp, 256]
        mean = mu_sigma_x[:,:,:256]
        log_var = mu_sigma_x[:,:,256:]
        log_var = tf.clip_by_value(log_var, -10.0, 1.0)
        cmean = mu_sigma_c[:,:,:256]
        clog_var = mu_sigma_c[:,:,256:]
        clog_var = tf.clip_by_value(clog_var, -10.0, 1.0)
        zi = sample(mean, log_var)
        zc = cmean
        z = tf.cond(is_training, lambda: zi, lambda: zc)
        
        # Decode shapes pc [B, nsmp, nsmp_ins, 3]
        gcfeat = tf_util.conv1d(pcfea_seed, 256, 1, padding='VALID', bn=True, is_training=is_training,
                                scope='dec_fc', bn_decay=bn_decay)
        feat = tf.concat((z, gcfeat), axis=-1)
        pc_ins_pred, pc_ins_conf_pred = decoding_conf_net(feat, nsmp_ins, 'decoder', is_training=is_training, bn_decay=bn_decay)
        # pc_ins_pred = decoding_net(feat, nsmp_ins, 'decoder', is_training=is_training, bn_decay=bn_decay)
        pc_ins_pred = pc_ins_pred + tf.stop_gradient(tf.expand_dims(shift_pred_seed, 2))
        # pc_ins_pred = pc_ins_pred + tf.expand_dims(shift_pred_seed, 2)

        # Collect bbox for reconstructions, pc_ins_conf_pred [B, nsmp, nsmp_ins]
        pc_ins_pred_world_coord = pc_ins_pred + tf.expand_dims(pc_seed, 2)
        bbox_ins_pred = tf.concat(((tf.reduce_max(pc_ins_pred_world_coord, 2)+tf.reduce_min(pc_ins_pred_world_coord, 2))/2, 
                                  tf.reduce_max(pc_ins_pred_world_coord, 2)-tf.reduce_min(pc_ins_pred_world_coord, 2)), 2) # [B, nsmp, 6] -> center + l,w,h

        # pc_ins_pred_world_coord = pc_ins_pred + tf.expand_dims(pc_seed, 2) # [B, nsmp, nsmp_ins, 3]
        # inconf_pt_mask = 1000.0*tf.cast(tf.greater(pc_ins_conf_pred,0.5), tf.float32) 
        # inconf_pt_mask = tf.multiply(inconf_pt_mask, 1-tf.cast(tf.less_equal(pc_ins_conf_pred, 1e-3+tf.reduce_min(pc_ins_conf_pred, 2, keep_dims=True)), tf.float32))
        # inconf_pt_mask = tf.expand_dims(inconf_pt_mask, -1) # [B, nsmp, nsmp_ins, 1]

        # max_edge = tf.reduce_max(pc_ins_pred_world_coord-inconf_pt_mask, 2)
        # min_edge = tf.reduce_min(pc_ins_pred_world_coord+inconf_pt_mask, 2)
        # bbox_ins_pred = tf.concat(((max_edge+min_edge)/2, 
        #                           max_edge-min_edge), 2) # [B, nsmp, 6] -> center + l,w,h

        # # Random shift and scale box during training
        # bbox_ins_pred = tf.concat( (bbox_ins_pred[:,:,:3]+bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=0.0, stddev=0.1),
        #     bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=1.0, stddev=0.1)), axis=-1)

        # Propagate seed feature and center position
        if return_fullfea:
            entity_fea = pointnet_fp_module(pc, pc_seed, None, pcfea_seed, [], is_training=False, bn_decay=None, scope='entity_fea_prop', bn=False)
            # shift_pred = pointnet_fp_module(pc, pc_seed, None, shift_pred_seed, [], is_training=False, bn_decay=None, scope='shift_prop', bn=False)
            shift_pred = tf.multiply(end_points['shift_pred_full_4d'][:,:,:3], end_points['shift_pred_full_4d'][:,:,3:])
            center_pos = pc+shift_pred
            end_points['entity_fea'] = entity_fea # [B, N, 256] entity feature of each point
            end_points['center_pos'] = center_pos # [B, N, 3] center location in the world coord sys

        # Store end_points
        end_points['shift_pred_seed'] = shift_pred_seed # [B, nsmp, 3], offset from seed point to ins center
        end_points['shift_pred_seed_4d'] = shift_pred_seed_4d # [B, nsmp, 4], offset from seed point to ins center
        end_points['pc_seed'] = pc_seed # [B, nsmp, 3], seed point coordinate in world coord sys
        end_points['ind_seed'] = ind_seed # [B, nsmp], seed index
        end_points['pc_ins_centered_seed'] = pc_ins_centered_seed # [B, nsmp, nsmp_ins, 3], centered gt instance point cloud for each seed
        end_points['pc_ins_center_seed'] = pc_ins_center_seed # [B, nsmp, 1, 3], gt instance center for each seed in world coord sys
        end_points['mean'] = mean # [B, nsmp, 256]
        end_points['log_var'] = log_var
        end_points['cmean'] = cmean
        end_points['clog_var'] = clog_var
        end_points['fb_logits'] = fb_logits # [B, nsmp, 2] foreground/backgroud logits
        end_points['fb_prob'] = fb_prob # [B, nsmp, 2] foreground/background probability
        end_points['pc_ins_pred'] = pc_ins_pred # [B, nsmp, nsmp_ins, 3], in local sys, needs to add pc_seed
        end_points['pc_ins_conf_pred'] = pc_ins_conf_pred # [B, nsmp, nsmp_ins] per-point confidence
        end_points['bbox_ins_pred'] = bbox_ins_pred # [B, nsmp, 6]

        return end_points #, alexnetmodel

def nms_3d(boxes, scores, pre_nms_limit, max_output_size, iou_threshold=0.5, score_threshold=float('-inf')):
    ''' Non maximum suppression in 3D
    Inputs:
        boxes: [B, N, 6] center + l,w,h
        scores: [B, N] prob between 0 and 1
    Outputs:
        selected_indices: [B, M]
    '''

    batch_size = scores.shape[0]    
    num_box_input = scores.shape[1]    
    sidx = np.argsort(-scores, 1) # [B, N] from large to small
    selected_indices = -np.ones((batch_size, max_output_size), dtype=np.int32) # [B, M]
    for i in range(batch_size):
        cursidx = sidx[i,:]
        curscores = scores[i,:]
        curvolume = boxes[i,:,3]*boxes[i,:,4]*boxes[i,:,5]
        if pre_nms_limit>0:
            cursidx = cursidx[:pre_nms_limit]
        cursidx = cursidx[curscores[cursidx]>score_threshold]
        count = 0
        while len(cursidx)>0 and count<max_output_size:
            selected_indices[i,count] = cursidx[0]
            count += 1
            vA = np.maximum(boxes[i,[cursidx[0]],:3]-boxes[i,[cursidx[0]],3:]/2, boxes[i,cursidx,:3]-boxes[i,cursidx,3:]/2)
            vB = np.minimum(boxes[i,[cursidx[0]],:3]+boxes[i,[cursidx[0]],3:]/2, boxes[i,cursidx,:3]+boxes[i,cursidx,3:]/2)
            intersection_cube = np.maximum(vB-vA,0)
            intersection_volume = intersection_cube[:,0]*intersection_cube[:,1]*intersection_cube[:,2]
            iou = np.divide(intersection_volume,curvolume[cursidx]+curvolume[cursidx[0]]-intersection_volume+1e-8)
            cursidx = np.delete(cursidx, np.where(iou>iou_threshold)[0])
    return selected_indices

def gather_selection(source, selected_idx, max_selected_size):
    '''
    Inputs:
        source: [B, N, C]
        selected_idx: [B, M], -1 means not selecting anything
    Returns:
        target: [B, M, C], 0 padded
    '''
    batch_size = source.get_shape()[0].value
    fea_size = source.get_shape()[2].value
    pos_idx = tf.cast(tf.where(tf.greater_equal(selected_idx,0)), tf.int32)
    selected_idx_vec = tf.gather_nd(selected_idx, pos_idx)
    target_vec = tf.gather_nd(source, tf.concat((tf.expand_dims(pos_idx[:,0],-1), tf.reshape(selected_idx_vec,[-1,1])),1))
    target = tf.scatter_nd(pos_idx, target_vec, tf.constant([batch_size, max_selected_size, fea_size]))
    return target

def trim_zeros_graph(boxes, name=None):
    """Often boxes are represented with matrices of shape [N, 4] and
    are padded with zeros. This removes zero boxes.

    boxes: [N, 6] matrix of boxes.
    non_zeros: [N] a 1D boolean mask identifying the rows to keep
    """
    non_zeros = tf.cast(tf.reduce_sum(tf.abs(boxes), axis=1), tf.bool)
    boxes = tf.boolean_mask(boxes, non_zeros, name=name)
    return boxes, non_zeros

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

def box_shrink(box, pc):
    ''' Shrink bounding box so that it is tight with respect to pc
    Inputs:
        box: [B, NUM_SAMPLE, 6]
        pc: [B, NUM_POINT, 3]
    Returns:
        box: [B, NUM_SAMPLE, 6]
    '''
    pc_aug = tf.expand_dims(pc, 1) # [B, 1, NUM_POINT, 3]
    box_aug = tf.expand_dims(box, 2) # [B, NUM_SAMPLE, 1, 6]
    box_masks = tf.logical_and(pc_aug>=(box_aug[:,:,:,:3]-box_aug[:,:,:,3:]/2),
                               pc_aug<=(box_aug[:,:,:,:3]+box_aug[:,:,:,3:]/2)) # [B, NUM_SAMPLE, NUM_POINT, 3]
    box_masks = tf.logical_and(tf.logical_and(box_masks[:,:,:,0], box_masks[:,:,:,1]), box_masks[:,:,:,2]) # [B, NUM_SAMPLE, NUM_POINT]
    box_out_masks = 1-tf.cast(tf.expand_dims(box_masks, -1), tf.float32) # [B, NUM_SAMPLE, NUM_POINT, 1]
    gamma = 1e4 # a large number for the box estimation trick
    box_max = tf.reduce_max(pc_aug-gamma*box_out_masks,2) # [B, NUM_SAMPLE, 3]
    box_min = tf.reduce_min(pc_aug+gamma*box_out_masks,2)
    box = tf.concat( ((box_max+box_min)/2, box_max-box_min+1e-3), 2) # [B, NUM_SAMPLE, 6]
    keep = tf.greater(box_max-box_min, 0)
    keep = tf.logical_and(tf.logical_and(keep[:,:,0], keep[:,:,1]), keep[:,:,2])
    keep = tf.expand_dims(tf.cast(keep, tf.float32), -1)
    box = tf.multiply(box, keep)
    return box

def box_refinement(box, gt_box):
    """Compute refinement needed to transform box to gt_box.
    box and gt_box are [N, (center_x, center_y, center_z, l, w, h)]
    """
    box = tf.cast(box, tf.float32)
    gt_box = tf.cast(gt_box, tf.float32)

    dz = (gt_box[:,2] - box[:,2]) / (box[:,5]+1e-8)
    dy = (gt_box[:,1] - box[:,1]) / (box[:,4]+1e-8)
    dx = (gt_box[:,0] - box[:,0]) / (box[:,3]+1e-8)
    dh = tf.log(gt_box[:,5] / (box[:,5]+1e-8))
    dw = tf.log(gt_box[:,4] / (box[:,4]+1e-8))
    dl = tf.log(gt_box[:,3] / (box[:,3]+1e-8))

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

def apply_box_delta(box, delta):
    ''' Apply bounding box delta to refine box
    Inputs:
        box: [NUM_SAMPLE, 6]
        delta: [NUM_SAMPLE, 6]
    Returns:
        box_refined: [NUM_SAMPLE, 6]
    '''
    delta = tf.stack([delta[:,2], delta[:,1], delta[:,0], delta[:,5], delta[:,4], delta[:,3]], axis=1)
    box_refined_part2 = tf.multiply(tf.exp(delta[:,3:]), box[:,3:])
    box_refined_part1 = tf.multiply(delta[:,:3], box[:,3:])+box[:,:3]
    box_refined = tf.concat((box_refined_part1, 1.3*box_refined_part2), axis=1)
    return box_refined

def sample_points_within_box(masks, nsmp):
    '''
    Inputs:
        masks: [nmask, npoint]
        nsmp: scalar
    Returns:
        masks_selection_idx: [nmask, nsmp]
    '''
    if masks.shape[0]==0:
        return np.zeros((0, nsmp), dtype=np.int32)
    else:
        masks_selection_idx = [np.random.choice(np.where(masks[i,:])[0], nsmp, replace=True) if len(np.where(masks[i,:])[0])>0 else np.zeros(nsmp, dtype=np.int32) for i in np.arange(masks.shape[0])]
        masks_selection_idx = np.stack(masks_selection_idx, 0).astype(np.int32)
        return masks_selection_idx

def spn_target_gen(proposals, proposal_seed_class_ids, gt_class_ids, gt_boxes):
    ''' SPN target generation
    Inputs:
        proposals: [NUM_SAMPLE, 6]
        proposal_seed_class_ids: [NUM_SAMPLE] - 0 is background class and 1 is foreground class
        gt_class_ids: [NUM_GROUP] - 0 is background class and 1 is foreground class
        gt_boxes: [NUM_GROUP, 6]
    Returns:
        spn_match: [NUM_SAMPLE], 1 = positive, -1 = negative, 0 = neutral
    '''
    # Remove zero padding
    non_zeros = tf.where(tf.greater(gt_class_ids, 0))[:,0]
    gt_boxes = tf.gather(gt_boxes, non_zeros, axis=0, name="spn_trim_gt_boxes")
    gt_class_ids = tf.gather(gt_class_ids, non_zeros, axis=0, name="spn_trim_gt_class_ids")

    # Compute IoU [NUM_SAMPLE, NUM_GROUP_TRIMMED]
    proposals_aug = tf.expand_dims(proposals, 1)
    gt_boxes_aug = tf.expand_dims(gt_boxes, 0)
    proposal_volume = tf.multiply(tf.multiply(proposals_aug[:,:,3], proposals_aug[:,:,4]), proposals_aug[:,:,5])
    gt_boxes_volume = tf.multiply(tf.multiply(gt_boxes_aug[:,:,3], gt_boxes_aug[:,:,4]), gt_boxes_aug[:,:,5])
    vA = tf.maximum(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]-gt_boxes_aug[:,:,3:]/2)
    vB = tf.minimum(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]+gt_boxes_aug[:,:,3:]/2)
    intersection_cube = tf.maximum(vB-vA,0)
    intersection_volume = tf.multiply(tf.multiply(intersection_cube[:,:,0], intersection_cube[:,:,1]), intersection_cube[:,:,2])
    ious = tf.divide(intersection_volume,proposal_volume+gt_boxes_volume-intersection_volume+1e-8)

    # Determine positive and negative ROIs
    roi_iou_max = tf.reduce_max(ious, axis=1)
    # 1. Positive ROIs are those with >=0.5 IoU with a GT box and gt class id is 1
    # We also guarantee that the boxes with the largest IoU with a GT box is also treated as positive
    spn_match_positive = tf.logical_and(tf.greater_equal(roi_iou_max, 0.5), tf.equal(proposal_seed_class_ids,1))
    masked_ious = tf.multiply(ious, tf.cast(tf.equal(tf.expand_dims(proposal_seed_class_ids,-1),1), tf.float32))
    positive_aug_idx = tf.argmax(masked_ious,0)
    positive_aug_idx = tf.boolean_mask(positive_aug_idx, tf.reduce_max(masked_ious, 0)>0)
    updates = tf.cast(tf.ones_like(positive_aug_idx), tf.float32)
    spn_match_positive_aug = tf.cond(
        tf.size(positive_aug_idx)>0,
        true_fn = lambda: tf.scatter_nd(positive_aug_idx, updates, spn_match_positive.shape),
        false_fn = lambda: tf.constant(0.0))
    spn_match_positive = tf.greater(tf.cast(spn_match_positive, tf.float32)+spn_match_positive_aug, 0)
    # 2. Negative ROIs are those with < 0.5 with every GT box and not a postive ROI
    spn_match_negative = tf.logical_and(tf.less(roi_iou_max, 0.5), tf.logical_not(spn_match_positive))

    spn_match = tf.cast(spn_match_positive, tf.float32)-tf.cast(spn_match_negative, tf.float32)

    return spn_match


def detection_target_gen(proposals, gt_class_ids, gt_boxes, gt_masks, pc, config):
    ''' Generate detection targets for training
    Inputs:
        proposals: [SPN_NMS_MAX_SIZE, 6], zero padded
        gt_class_ids: [NUM_GROUP] - 0 is background class and objects start from 1
        gt_boxes: [NUM_GROUP, 6]
        gt_masks: [NUM_POINT, NUM_GROUP]
        pc: [NUM_POINT, 3]
    Returns:
        rois: [TRAIN_ROIS_PER_IMAGE, 6], zero padded
        roi_gt_class_ids: [TRAIN_ROIS_PER_IMAGE] - 0 is invalid class and used for padding, zero padded
        deltas: [TRAIN_ROIS_PER_IMAGE, 6], zero padded
        masks_selection_idx: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded
        masks: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - binary mask, zero padded
    '''
    # Remove zero padding
    proposals, _ = trim_zeros_graph(proposals, name='trim_proposals')
    gt_boxes, non_zeros = trim_zeros_graph(gt_boxes, name="trim_gt_boxes")
    gt_class_ids = tf.boolean_mask(gt_class_ids, non_zeros,
                                   name="trim_gt_class_ids")
    gt_masks = tf.gather(gt_masks, tf.where(non_zeros)[:, 0], axis=1,
                         name="trim_gt_masks")

    # Remove empty proposals
    pc_aug = tf.expand_dims(pc,1) # [N, 1, 3]
    proposals_aug = tf.expand_dims(proposals, 0) # [1, NP, 6]
    roi_masks = tf.logical_and(pc_aug>=(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2),
                               pc_aug<=(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,0], roi_masks[:,:,1]), roi_masks[:,:,2]) # [N, NP]
    non_empty_idx = tf.where(tf.greater(tf.reduce_sum(tf.cast(roi_masks, tf.float32), 0), 0))[:,0]
    roi_masks = tf.gather(roi_masks, non_empty_idx, axis=1) # [N, NP']
    proposals = tf.gather(proposals, non_empty_idx, axis=0) # [NP', 6]

    # Compute IoU [n_proposal, n_gt_boxes]
    proposals_aug = tf.expand_dims(proposals, 1)
    gt_boxes_aug = tf.expand_dims(gt_boxes, 0)
    proposal_volume = tf.multiply(tf.multiply(proposals_aug[:,:,3], proposals_aug[:,:,4]), proposals_aug[:,:,5])
    gt_boxes_volume = tf.multiply(tf.multiply(gt_boxes_aug[:,:,3], gt_boxes_aug[:,:,4]), gt_boxes_aug[:,:,5])
    vA = tf.maximum(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]-gt_boxes_aug[:,:,3:]/2)
    vB = tf.minimum(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2, gt_boxes_aug[:,:,:3]+gt_boxes_aug[:,:,3:]/2)
    intersection_cube = tf.maximum(vB-vA,0)
    intersection_volume = tf.multiply(tf.multiply(intersection_cube[:,:,0], intersection_cube[:,:,1]), intersection_cube[:,:,2])
    ious = tf.divide(intersection_volume,proposal_volume+gt_boxes_volume-intersection_volume+1e-8)
    
    # Determine positive and negative ROIs
    roi_iou_max = tf.reduce_max(ious, axis=1)
    # 1. Positive ROIs are those with >= 0.5 IoU with a GT box
    positive_indices = tf.where(roi_iou_max >= 0.5)[:, 0]
    # 2. Negative ROIs are those with < 0.5 with every GT box
    negative_indices = tf.where(roi_iou_max < 0.5)[:, 0]

    # Subsample ROIs. Aim for 33% positive
    # Positive ROIs
    positive_count = int(config.TRAIN_ROIS_PER_IMAGE *
                         config.ROI_POSITIVE_RATIO)
    positive_indices = tf.random_shuffle(positive_indices)[:positive_count]
    positive_count = tf.shape(positive_indices)[0]
    # Negative ROIs. Add enough to maintain positive:negative ratio.
    r = 1.0 / config.ROI_POSITIVE_RATIO
    negative_count = tf.cast(r * tf.cast(positive_count, tf.float32), tf.int32) - positive_count
    negative_indices = tf.random_shuffle(negative_indices)[:negative_count]
    # Gather selected ROIs [POSITIVE_COUNT/NEGATIVE_COUNT, 6]
    positive_rois = tf.gather(proposals, positive_indices)
    negative_rois = tf.gather(proposals, negative_indices)

    # Assign positive ROIs to GT boxes. roi_gt_boxes: [POSITIVE_COUNT, 6], roi_gt_class_ids: [POSITIVE_COUNT]
    positive_ious = tf.gather(ious, positive_indices)
    roi_gt_box_assignment = tf.cond(
        tf.greater(tf.shape(positive_ious)[1], 0),
        true_fn = lambda: tf.argmax(positive_ious, axis=1),
        false_fn = lambda: tf.cast(tf.constant([]),tf.int64)
    )
    roi_gt_boxes = tf.gather(gt_boxes, roi_gt_box_assignment)
    roi_gt_class_ids = tf.gather(gt_class_ids, roi_gt_box_assignment)

    # Compute bbox refinement for positive ROIs. delta: [POSITIVE_COUNT, 6]
    deltas = box_refinement(positive_rois, roi_gt_boxes)
    deltas /= config.BBOX_STD_DEV

    # Compute mask targets
    roi_gt_masks = tf.cast(tf.gather(gt_masks, roi_gt_box_assignment, axis=1), tf.bool) # [NUM_POINT, POSITIVE_COUNT]
    positive_roi_masks = tf.gather(roi_masks, positive_indices, axis=1)
    masks_full = tf.transpose(tf.logical_and(positive_roi_masks, roi_gt_masks)) # [POSITIVE_COUNT, NUM_POINT]
    masks_selection_idx = tf.stop_gradient(tf.py_func(sample_points_within_box, 
        [tf.transpose(positive_roi_masks), config.NUM_POINT_INS_MASK], tf.int32)) # [POSITIVE_COUNT, NUM_POINT_INS_MASK]
    smp_idx = tf.reshape(tf.tile(tf.reshape(tf.range(positive_count),[-1,1]),[1, config.NUM_POINT_INS_MASK]),[-1,1])
    smp_idx = tf.concat((smp_idx, tf.reshape(masks_selection_idx,[-1,1])),1)
    masks = tf.reshape(tf.gather_nd(masks_full, smp_idx),[-1, config.NUM_POINT_INS_MASK]) # [POSITIVE_COUNT, NUM_POINT_INS_MASK]

    # Append negative ROIs and pad bbox deltas and masks that
    # are not used for negative ROIs with zeros.
    rois = tf.concat([positive_rois, negative_rois], axis=0)
    N = tf.shape(negative_rois)[0]
    P = tf.maximum(config.TRAIN_ROIS_PER_IMAGE - tf.shape(rois)[0], 0)
    rois = tf.pad(rois, [(0, P), (0, 0)])
    roi_gt_class_ids = tf.pad(roi_gt_class_ids, [(0, N + P)])
    deltas = tf.pad(deltas, [(0, N + P), (0, 0)])
    masks_selection_idx = tf.pad(masks_selection_idx, [[0, N + P], (0, 0)])
    masks = tf.pad(masks, [[0, N + P], (0, 0)])

    return rois, roi_gt_class_ids, deltas, masks_selection_idx, masks

def mask_selection_gen(proposals, pc, num_rois, config, empty_removal=True):
    ''' Generate detection targets for training
    Inputs:
        proposals: [SPN_NMS_MAX_SIZE, 6], zero padded
        pc: [NUM_POINT, 3]
    Returns:
        proposals: [NUM_ROIS, 6]
        masks_selection_idx: [NUM_ROIS, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded
    '''
    # Remove zero padding
    proposals, _ = trim_zeros_graph(proposals, name='trim_proposals')

    # Remove empty proposals
    pc_aug = tf.expand_dims(pc,1) # [N, 1, 3]
    proposals_aug = tf.expand_dims(proposals, 0) # [1, NP, 6]
    roi_masks = tf.logical_and(pc_aug>=(proposals_aug[:,:,:3]-proposals_aug[:,:,3:]/2-1e-3),
                               pc_aug<=(proposals_aug[:,:,:3]+proposals_aug[:,:,3:]/2+1e-3))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,0], roi_masks[:,:,1]), roi_masks[:,:,2]) # [N, NP]
    if empty_removal:
        non_empty_idx = tf.where(tf.greater(tf.reduce_sum(tf.cast(roi_masks, tf.float32), 0), 0))[:,0]
        roi_masks = tf.gather(roi_masks, non_empty_idx, axis=1) # [N, NP']
        proposals = tf.gather(proposals, non_empty_idx, axis=0) # [NP', 6]
    proposals_count = tf.shape(proposals)[0]

    # Generate mask selection index
    masks_selection_idx = tf.stop_gradient(tf.py_func(sample_points_within_box, 
        [tf.transpose(roi_masks), config.NUM_POINT_INS_MASK], tf.int32)) # [NP', NUM_POINT_INS_MASK]

    # Append negative ROIs and pad bbox deltas and masks that
    # are not used for negative ROIs with zeros.
    P = tf.maximum(num_rois - proposals_count, 0)
    masks_selection_idx = tf.pad(masks_selection_idx, [[0, P], (0, 0)])
    proposals = tf.pad(proposals, [(0, P), (0, 0)])

    return proposals, masks_selection_idx

def points_cropping(pc, pc_fea, pc_center, rois, masks_selection_idx, num_rois, num_point_per_roi, normalize_crop_region=True):
    ''' Crop points for network heads, in analogy to ROIAlign
    Inputs:
        pc: [B, NUM_POINT, 3]
        pc_fea: [B, NUM_POINT, NFEA]
        pc_center: [B, NUM_POINT, 3]
        rois: [B, NUM_ROIS, 6], zero padded
        masks_selection_idx: [B, NUM_ROIS, NUM_POINT_PER_ROI]
    Returns:
        pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
    '''
    batch_size = pc.get_shape()[0].value
    smp_idx = tf.reshape(tf.tile(tf.reshape(tf.range(batch_size),[-1,1]),[1, num_rois*num_point_per_roi]),[-1,1])
    smp_idx = tf.concat((smp_idx, tf.reshape(masks_selection_idx,[-1,1])),1)
    pc_fea_cropped = tf.reshape(tf.gather_nd(pc_fea, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_center_cropped = tf.reshape(tf.gather_nd(pc_center, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_coord_cropped_unnormalized = tf.reshape(tf.gather_nd(pc, smp_idx),[batch_size, num_rois, num_point_per_roi, -1])
    pc_coord_cropped = pc_coord_cropped_unnormalized

    # convert world coord to local
    rois_center = tf.expand_dims(rois[:,:,:3], 2)
    pc_coord_cropped = pc_coord_cropped-rois_center
    pc_center_cropped = pc_center_cropped-rois_center
    if normalize_crop_region:
        # scale box to [1,1,1]
        rois = rois+tf.cast(tf.equal(tf.reduce_sum(rois, 2, keep_dims=True),0),tf.float32)
        rois_size = tf.expand_dims(rois[:,:,3:], 2)
        pc_coord_cropped = tf.divide(pc_coord_cropped, rois_size)
        pc_center_cropped = tf.divide(pc_center_cropped, rois_size)
    return pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized

def refine_detections(rois, probs, deltas, pc, fb_prob, sem_prob, config):
    '''Refine classified proposals and filter overlaps and return final
    detections.
    Inputs:
        rois: [NUM_ROIS, 6], zero padded, in world coord sys
        probs: [NUM_ROIS, NUM_CATEGORY] - 0 is background class and objects start from 1
        deltas: [NUM_ROIS, NUM_CATEGORY, 6]
        pc: [NUM_POINT, 3]
        fb_prob: [NUM_ROIS]
        sem_prob: [NUM_ROIS]
    Returns:
        detections: [NUM_DETECTIONS, (center_x, center_y, center_z, l, w, h, class_id, score)]
    '''
    # Class IDs per ROI
    class_ids = tf.argmax(probs, axis=1, output_type=tf.int32)
    # Class probability of the top class of each ROI
    indices = tf.stack([tf.range(probs.shape[0]), class_ids], axis=1)
    class_scores = tf.gather_nd(probs, indices) # [NUM_ROIS]
    # Class-specific bounding box deltas
    deltas_specific = tf.gather_nd(deltas, indices) # [NUM_ROIS, 6]
    # Apply bounding box deltas

    # Shape: [NUM_ROIS, (center_x, center_y, center_z, l, w, h)]
    refined_rois = apply_box_delta(rois, deltas_specific * config.BBOX_STD_DEV)
    # Shrink boxes
    if config.SHRINK_BOX:
        refined_rois = tf.squeeze(box_shrink(tf.expand_dims(refined_rois, 0), tf.expand_dims(pc,0)),0)

    # Filter out background boxes
    keep = tf.where(class_ids > 0)[:, 0]
    # keep = tf.where(tf.logical_and(tf.logical_and(class_ids > 0, fb_prob>0.01), sem_prob>0.01))[:, 0]
    # Filter out low confidence boxes
    if config.DETECTION_MIN_CONFIDENCE:
        conf_keep = tf.where(class_scores >= config.DETECTION_MIN_CONFIDENCE)[:, 0]
        keep = tf.sets.set_intersection(tf.expand_dims(keep, 0),
                                        tf.expand_dims(conf_keep, 0))
        keep = tf.sparse_tensor_to_dense(keep)[0]

    # Apply per-class NMS
    # 1. Prepare variables
    pre_nms_class_ids = tf.gather(class_ids, keep)
    pre_nms_scores = tf.gather(class_scores, keep) # [NUM_KEEP]
    # pre_nms_scores = tf.gather(class_scores*fb_prob*sem_prob, keep) # [NUM_KEEP]
    pre_nms_rois = tf.gather(refined_rois,   keep)
    unique_pre_nms_class_ids = tf.unique(pre_nms_class_ids)[0]

    def nms_keep_map(class_id):
        """Apply Non-Maximum Suppression on ROIs of the given class."""
        # Indices of ROIs of the given class
        ixs = tf.where(tf.equal(pre_nms_class_ids, class_id))[:, 0]
        # Apply NMS

        class_keep = tf.py_func(nms_3d, [tf.expand_dims(tf.gather(pre_nms_rois, ixs),0),
            tf.expand_dims(tf.gather(pre_nms_scores, ixs),0),
            -1, config.DETECTION_MAX_INSTANCES,
            config.DETECTION_NMS_THRESHOLD, float('-inf')], tf.int32)
        class_keep = tf.squeeze(class_keep, 0) # [DETECTION_MAX_INSTANCES], -1 padded
        class_keep = tf.gather(class_keep, tf.where(class_keep > -1)[:,0]) # [<=DETECTION_MAX_INSTANCES], no padding

        # Map indices
        class_keep = tf.gather(keep, tf.gather(ixs, class_keep))
        # Pad with -1 so returned tensors have the same shape
        gap = config.DETECTION_MAX_INSTANCES - tf.shape(class_keep)[0]
        class_keep = tf.pad(class_keep, [(0, gap)],
                            mode='CONSTANT', constant_values=-1)
        # Set shape so map_fn() can infer result shape
        class_keep.set_shape([config.DETECTION_MAX_INSTANCES])
        return class_keep

    # 2. Map over class IDs
    nms_keep = tf.map_fn(nms_keep_map, unique_pre_nms_class_ids,
                         dtype=tf.int64)
    # 3. Merge results into one list, and remove -1 padding
    nms_keep = tf.reshape(nms_keep, [-1])
    nms_keep = tf.gather(nms_keep, tf.where(nms_keep > -1)[:, 0])
    # 4. Compute intersection between keep and nms_keep
    keep = tf.sets.set_intersection(tf.expand_dims(keep, 0),
                                    tf.expand_dims(nms_keep, 0))
    keep = tf.sparse_tensor_to_dense(keep)[0]
    # Keep top detections
    roi_count = config.DETECTION_MAX_INSTANCES
    class_scores_keep = tf.gather(class_scores, keep)
    # class_scores_keep = tf.gather(class_scores*fb_prob*sem_prob, keep)
    num_keep = tf.minimum(tf.shape(class_scores_keep)[0], roi_count)
    top_ids = tf.nn.top_k(class_scores_keep, k=num_keep, sorted=True)[1]
    keep = tf.gather(keep, top_ids)

    # Arrange output as [N, (center_x, center_y, center_z, l, w, h, class_id, score)]
    detections = tf.concat([
        tf.gather(refined_rois, keep),
        tf.to_float(tf.gather(class_ids, keep))[..., tf.newaxis],
        tf.gather(class_scores, keep)[..., tf.newaxis]
        ], axis=1)

    # Pad with zeros if detections < DETECTION_MAX_INSTANCES
    gap = config.DETECTION_MAX_INSTANCES - tf.shape(detections)[0]
    detections = tf.pad(detections, [(0, gap), (0, 0)], "CONSTANT")
    return detections

def classification_head(pc, pc_fea, num_category, mlp_list, mlp_list2, is_training, bn_decay, scope, bn=True):
    ''' Classification head for both class id prediction and bbox delta regression
    Inputs:
        pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        num_category: scalar
    Returns:
        logits: [B, NUM_ROIS, NUM_CATEGORY]
        probs: [B, NUM_ROIS, NUM_CATEGORY]
        bbox_deltas: [B, NUM_ROIS, NUM_CATEGORY, (dz, dy, dx, log(dh), log(dw), log(dl))]
    '''
    with tf.variable_scope(scope) as myscope:
        num_rois = pc.get_shape()[1].value
        grouped_points = tf.concat((pc_fea, pc), -1)
        for i,num_out_channel in enumerate(mlp_list):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_prev_%d'%i, bn_decay=bn_decay)
        new_points = tf.reduce_max(grouped_points, axis=2)
        for i,num_out_channel in enumerate(mlp_list2):
            new_points = tf_util.conv1d(new_points, num_out_channel, 1,
                                        padding='VALID', stride=1, bn=bn, is_training=is_training,
                                        scope='conv_post_%d'%i, bn_decay=bn_decay)
        logits = tf_util.conv1d(new_points, num_category, 1, padding='VALID',
                                stride=1, scope='conv_classify', activation_fn=None)
        probs = tf.nn.softmax(logits, 2)
        bbox_deltas = tf_util.conv1d(new_points, num_category*6, 1, padding='VALID',
                                     stride=1, scope='conv_bbox_regress', activation_fn=None)
        bbox_deltas = tf.reshape(bbox_deltas, [-1, num_rois, num_category, 6])
        return logits, probs, bbox_deltas

def segmentation_head(pc, pc_fea, num_category, mlp_list, mlp_list2, mlp_list3, is_training, bn_decay, scope, bn=True):
    ''' Segmentation head
    Inputs:
        pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
        num_category: scalar
    Returns:
        masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
    '''
    with tf.variable_scope(scope) as myscope:
        num_rois = pc.get_shape()[1].value
        num_point_per_roi = pc.get_shape()[2].value
        grouped_points = tf.concat((pc_fea, pc), -1)
        for i,num_out_channel in enumerate(mlp_list):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_prev_%d'%i, bn_decay=bn_decay)
        local_feat = grouped_points
        for i,num_out_channel in enumerate(mlp_list2):
            grouped_points = tf_util.conv2d(grouped_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_%d'%i, bn_decay=bn_decay)
        global_feat = tf.reduce_max(grouped_points, axis=2, keep_dims=True)
        global_feat_expanded = tf.tile(global_feat, [1, 1, num_point_per_roi, 1])
        new_points = tf.concat((global_feat_expanded, local_feat), -1)
        for i,num_out_channel in enumerate(mlp_list3):
            new_points = tf_util.conv2d(new_points, num_out_channel, [1, 1],
                                            padding='VALID', stride=[1,1], bn=bn, is_training=is_training,
                                            scope='conv_post_%d'%i, bn_decay=bn_decay)

        masks = tf_util.conv2d(new_points, num_category, [1, 1], padding='VALID',
                               stride=[1,1], scope='conv_seg', activation_fn=None)
        return masks

def dict_stop_gradient(dict_in):
    keys = dict_in.keys()
    for key in keys:
        dict_in[key] = tf.stop_gradient(dict_in[key])
    return dict_in

def select_segmentation(mrcnn3d_masks, class_ids):
    ''' Convert segmentation into point cloud label
    Inputs:
        mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
        class_ids: [B, NUM_ROIS]
    Returns:
        mrcnn3d_mask_selected: [B, NUM_ROIS, NUM_POINT_PER_ROI]
    '''
    batch_size = mrcnn3d_masks.get_shape()[0].value
    num_rois = mrcnn3d_masks.get_shape()[1].value
    num_point_per_roi = mrcnn3d_masks.get_shape()[2].value
    num_category = mrcnn3d_masks.get_shape()[3].value

    mrcnn3d_masks = tf.reshape(mrcnn3d_masks, [-1, num_point_per_roi, num_category])
    mrcnn3d_masks = tf.transpose(mrcnn3d_masks, perm=[0, 2, 1]) # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI]
    class_ids = tf.cast(tf.reshape(class_ids, [-1]), tf.int32)
    class_ids_aug = tf.stack([tf.range(batch_size*num_rois, dtype=tf.int32), class_ids], 1)
    mrcnn3d_mask_selected = tf.gather_nd(mrcnn3d_masks, class_ids_aug) #[-1, NUM_POINT_PER_ROI]
    mrcnn3d_mask_selected = tf.reshape(mrcnn3d_mask_selected, [batch_size, num_rois, num_point_per_roi])

    return mrcnn3d_mask_selected

def unmold_segmentation(mrcnn3d_masks, rois, class_ids, pc_coord_cropped, pc):
    ''' Convert segmentation into point cloud label
    Inputs:
        mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
        rois: [B, NUM_ROIS, 6]
        class_ids: [B, NUM_ROIS]
        pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
        pc: [B, NUM_POINT, 3]
    Returns:
        mrcnn3d_mask_unmolded: [B, NUM_ROIS, NUM_POINT]
    '''
    batch_size = mrcnn3d_masks.get_shape()[0].value
    num_rois = mrcnn3d_masks.get_shape()[1].value
    num_point_per_roi = mrcnn3d_masks.get_shape()[2].value
    num_category = mrcnn3d_masks.get_shape()[3].value
    num_point = pc.get_shape()[1].value

    mrcnn3d_masks = tf.reshape(mrcnn3d_masks, [-1, num_point_per_roi, num_category])
    mrcnn3d_masks = tf.transpose(mrcnn3d_masks, perm=[0, 2, 1]) # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI]
    class_ids = tf.cast(tf.reshape(class_ids, [-1]), tf.int32)
    class_ids_aug = tf.stack([tf.range(batch_size*num_rois, dtype=tf.int32), class_ids], 1)
    mrcnn3d_masks = tf.gather_nd(mrcnn3d_masks, class_ids_aug) #[-1, NUM_POINT_PER_ROI]

    # [B, NUM_ROIS, NUM_POINT, NUM_POINT_PER_ROI]
    dist = tf.reduce_sum(tf.square(tf.expand_dims(tf.expand_dims(pc, 1),3)-tf.expand_dims(pc_coord_cropped,2)), -1)
    min_idx = tf.argmin(dist, 3, output_type=tf.int32) # [B, NUM_ROIS, NUM_POINT]
    min_idx = tf.reshape(min_idx, [-1, num_point]) # [-1, NUM_POINT]
    min_idx_aug = tf.tile(tf.expand_dims(tf.range(batch_size*num_rois, dtype=tf.int32),-1), [1, num_point])
    min_idx_aug = tf.stack([tf.reshape(min_idx_aug, [-1]), tf.reshape(min_idx, [-1])], 1)
    mrcnn3d_mask_unmolded = tf.reshape(tf.gather_nd(mrcnn3d_masks, min_idx_aug), [batch_size, num_rois, num_point])

    # Mask out regions outside rois
    pc_aug = tf.expand_dims(pc, 1) # [B, 1, NUM_POINT, 3]
    rois_aug = tf.expand_dims(rois, 2) # [B, NUM_ROIS, 1, 6]
    roi_masks = tf.logical_and(pc_aug>=(rois_aug[:,:,:,:3]-rois_aug[:,:,:,3:]/2),
                               pc_aug<=(rois_aug[:,:,:,:3]+rois_aug[:,:,:,3:]/2))
    roi_masks = tf.logical_and(tf.logical_and(roi_masks[:,:,:,0], roi_masks[:,:,:,1]), roi_masks[:,:,:,2]) # [B, NUM_ROIS, NUM_POINT]
    roi_masks = tf.cast(roi_masks, tf.float32)

    mrcnn3d_mask_unmolded = tf.multiply(mrcnn3d_mask_unmolded, roi_masks)
    return mrcnn3d_mask_unmolded


def mrcnn_3d(pc, pc_ins, group_label, group_indicator, seg_label, seg_label_per_group, bbox_ins, config, is_training, mode='training', bn_decay=None):
    ''' Shape proposal generation
    Inputs:
        pc: [B, NUM_POINT, 3]
        color: [B, NUM_POINT, 3]
        pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys
        group_label: [B, NUM_POINT]
        group_indicator: [B, NUM_GROUP]
        seg_label: [B, NUM_POINT]
        bbox_ins: [B, NUM_GROUP, 6]
    Returns:
        
    '''
    assert mode in ['training', 'inference']
    if not config.USE_COLOR:
        color = None
    if 'SPN' in config.TRAIN_MODULE and mode=='training':
        end_points = shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins, config.NUM_CATEGORY, scope='shape_proposal_net', is_training=is_training, bn_decay=bn_decay, nsmp=config.NUM_SAMPLE, return_fullfea=False)
    else:
        end_points = shape_proposal_net(pc, pc_ins, group_label, group_indicator, bbox_ins, config.NUM_CATEGORY, scope='shape_proposal_net', is_training=tf.constant(False), bn_decay=None, nsmp=config.NUM_SAMPLE, return_fullfea=True)
        end_points = dict_stop_gradient(end_points)
    if config.SHRINK_BOX:
        end_points['bbox_ins_pred'] = box_shrink(end_points['bbox_ins_pred'], pc)
    group_label_onehot = tf.one_hot(group_label, depth=config.NUM_GROUP, axis=-1) #[B, NUM_POINT, NUM_GROUP]
    # seg_label_per_group = tf.multiply(tf.cast(tf.expand_dims(seg_label,-1), tf.float32), group_label_onehot)
    # seg_label_per_group = tf.cast(tf.round(tf.divide(tf.reduce_sum(seg_label_per_group, 1),tf.reduce_sum(group_label_onehot, 1)+1e-8)), tf.int32) #[B, NUM_GROUP]

    if 'MRCNN' in config.TRAIN_MODULE or mode=='inference':
        SPN_NMS_MAX_SIZE = config.SPN_NMS_MAX_SIZE_TRAINING if mode == "training"\
            else config.SPN_NMS_MAX_SIZE_INFERENCE
        # 3D non maximum suppression - selected_indices: [B, M], spn_rois: [B, M, 6]
        selected_indices = tf.stop_gradient(tf.py_func(nms_3d, [end_points['bbox_ins_pred'], end_points['fb_prob'][:,:,1], config.SPN_PRE_NMS_LIMIT, SPN_NMS_MAX_SIZE, config.SPN_IOU_THRESHOLD, config.SPN_SCORE_THRESHOLD], tf.int32))
        spn_rois = gather_selection(end_points['bbox_ins_pred'], selected_indices, SPN_NMS_MAX_SIZE)

        if mode=='training':
            # Detection target generation - rois: [B, TRAIN_ROIS_PER_IMAGE, 6], target_class_ids: [B, TRAIN_ROIS_PER_IMAGE]
            # target_bbox: [B, TRAIN_ROIS_PER_IMAGE, 6], target_mask_selection_idx: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK]
            # target_mask: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK], all zero padded
            names = ["rois", "target_class_ids", "target_bbox", "target_mask_selection_idx", "target_mask"]
            outputs = batch_slice(
                [spn_rois, seg_label_per_group, bbox_ins, group_label_onehot, pc],
                lambda v, w, x, y, z: detection_target_gen(v, w, x, y, z, config),
                config.BATCH_SIZE, names=names)
            rois, target_class_ids, target_bbox, target_mask_selection_idx, target_mask = outputs

            # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### color
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], color), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_pn2_fea = pn2_fea_extractor(pc, None, 'pn2_fea_extractor', is_training, bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            sem_fea_full_l1 = tf_util.conv1d(end_points['sem_fea_full_l1'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn1', bn_decay=bn_decay)
            sem_fea_full_l2 = tf_util.conv1d(end_points['sem_fea_full_l2'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn2', bn_decay=bn_decay)
            sem_fea_full_l3 = tf_util.conv1d(end_points['sem_fea_full_l3'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn3', bn_decay=bn_decay)
            sem_fea_full_l4 = tf_util.conv1d(end_points['sem_fea_full_l4'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn4', bn_decay=bn_decay)
            sem_fea_full_l5 = tf_util.conv1d(end_points['sem_fea_full_l5'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn5', bn_decay=bn_decay)
            pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5), -1), 
                end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), 
            #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            # Classification and bbox refinement head
            mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY, 
                [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')
            # mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
            #     pc_fea_cropped, config.NUM_CATEGORY, 
            #     [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')

            # Mask prediction head
            mrcnn3d_mask = segmentation_head(pc_coord_cropped,
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY,
                [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')
            # mrcnn3d_mask = segmentation_head(pc_coord_cropped,
            #     pc_fea_cropped, config.NUM_CATEGORY,
            #     [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')
        elif mode=='inference':
            # rois: [B, NUM_ROIS, 6]
            names = ["rois", "mask_selection_idx"]
            outputs = batch_slice(
                [spn_rois, pc],
                lambda x, y: mask_selection_gen(x, y, SPN_NMS_MAX_SIZE, config, empty_removal=True),
                config.BATCH_SIZE, names=names)
            rois, mask_selection_idx = outputs

            # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_pn2_fea = pn2_fea_extractor(pc, color, 'pn2_fea_extractor', is_training, bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            sem_fea_full_l1 = tf_util.conv1d(end_points['sem_fea_full_l1'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn1', bn_decay=bn_decay)
            sem_fea_full_l2 = tf_util.conv1d(end_points['sem_fea_full_l2'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn2', bn_decay=bn_decay)
            sem_fea_full_l3 = tf_util.conv1d(end_points['sem_fea_full_l3'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn3', bn_decay=bn_decay)
            sem_fea_full_l4 = tf_util.conv1d(end_points['sem_fea_full_l4'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn4', bn_decay=bn_decay)
            sem_fea_full_l5 = tf_util.conv1d(end_points['sem_fea_full_l5'], 64, 1, padding='VALID', bn=True, is_training=is_training, scope='fpn5', bn_decay=bn_decay)
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            #### generate fb_conf and sem_conf
            # pc: [B, N, 3], fb_prob: [B, nsmp, 2], pc_seed: [B, nsmp, 3] -> fb_prob: [B, N]
            midx = tf.argmin(tf.reduce_sum(tf.square(tf.expand_dims(pc, 2)-tf.expand_dims(end_points['pc_seed'],1)),-1), 2)
            midx_aug = tf.tile(tf.reshape(tf.range(config.BATCH_SIZE, dtype=tf.int64),[-1,1]), [1,config.NUM_POINT])
            midx_aug = tf.stack((tf.reshape(midx_aug, [-1]), tf.reshape(midx, [-1])), 1)
            fb_prob = tf.reshape(tf.gather_nd(end_points['fb_prob'], midx_aug), [config.BATCH_SIZE, config.NUM_POINT, 2])
            fb_prob = fb_prob[:,:,1]
            sem_prob = tf.nn.softmax(end_points['sem_class_logits'], -1) #[B, NUM_POINT, NUM_CATEGORY]

            ##### pn2
            # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea, tf.expand_dims(fb_prob, -1), sem_prob), -1), 
            #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            # pc_fea_cropped, fb_prob_cropped, sem_prob_cropped = tf.split(pc_fea_cropped, [832, 1, config.NUM_CATEGORY], -1)
            ##### sem fpn fea
            pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5, tf.expand_dims(fb_prob, -1), sem_prob), -1), 
                end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            pc_fea_cropped, fb_prob_cropped, sem_prob_cropped = tf.split(pc_fea_cropped, [4*128+5*64, 1, config.NUM_CATEGORY], -1)
            
            fb_prob_cropped = tf.reduce_mean(tf.squeeze(fb_prob_cropped, -1),-1) # [B, NUM_ROIS]
            end_points['fb_prob_cropped'] = fb_prob_cropped
            sem_prob_cropped = tf.reduce_mean(sem_prob_cropped, 2) # [B, NUM_ROIS, NUM_CATEGORY]

            # Classification and bbox refinement head - mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
            # mrcnn3d_class: [B, NUM_ROIS, NUM_CATEGORY],
            # mrcnn3d_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6]
            mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, 
                tf.concat((pc_fea_cropped, pc_center_cropped), -1), config.NUM_CATEGORY, 
                [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head')

            midx = tf.argmax(mrcnn3d_class_logits, -1) # [B, NUM_ROIS]
            midx = tf.stack((tf.range(config.BATCH_SIZE*SPN_NMS_MAX_SIZE, dtype=tf.int64), tf.reshape(midx,[-1])), 1)
            sem_prob_cropped = tf.gather_nd(tf.reshape(sem_prob_cropped, [-1, config.NUM_CATEGORY]), midx)
            sem_prob_cropped = tf.reshape(sem_prob_cropped, [config.BATCH_SIZE, SPN_NMS_MAX_SIZE])
            end_points['sem_prob_cropped'] = sem_prob_cropped

            # Generate detections: [B, DETECTION_MAX_INSTANCES, (center_x, center_y, center_z, l, w, h, class_id, score)]
            # detections = batch_slice(
            #     [rois, mrcnn3d_class, mrcnn3d_bbox, pc],
            #     lambda x, y, w, z: refine_detections(x, y, w, z, config),
            #     config.BATCH_SIZE)
            detections = batch_slice(
                [rois, mrcnn3d_class, mrcnn3d_bbox, pc, fb_prob_cropped, sem_prob_cropped],
                lambda u, v, x, y, w, z: refine_detections(u, v, x, y, w, z, config),
                config.BATCH_SIZE)

            # Re-crop point cloud for mask prediction
            names = ["rois_final", "mask_selection_idx_final"]
            outputs = batch_slice(
                [detections[:,:,:6], pc],
                lambda x, y: mask_selection_gen(x, y, config.DETECTION_MAX_INSTANCES, config, empty_removal=False),
                config.BATCH_SIZE, names=names)
            rois_final, mask_selection_idx_final = outputs
            # Points cropping - pc_fea_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NFEA]
            # pc_center_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
            # pc_coord_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
            ##### sem fea
            # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), 
            #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### pn2
            # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), 
            #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)
            ##### sem fpn fea
            pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4, sem_fea_full_l5), -1), 
                end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION)

            # Mask prediction head
            mrcnn3d_mask = segmentation_head(pc_coord_cropped_final,
                tf.concat((pc_fea_cropped_final, pc_center_cropped_final), -1), config.NUM_CATEGORY,
                [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head')

            # Unmold segmentation
            # mrcnn3d_mask_unmolded = unmold_segmentation(tf.nn.sigmoid(mrcnn3d_mask), rois_final, detections[:,:,6], pc_coord_cropped_final_unnormalized, pc)
            mrcnn3d_mask_selected = select_segmentation(tf.nn.sigmoid(mrcnn3d_mask), detections[:,:,6])

    # Update end_points
    end_points['group_label'] = group_label
    end_points['seg_label'] = seg_label
    end_points['seg_label_per_group'] = seg_label_per_group #[B, NUM_GROUP]
    end_points['bbox_ins'] = bbox_ins #[B, NUM_GROUP, 6]
    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        end_points['selected_indices'] = selected_indices #[B, SPN_NMS_MAX_SIZE]
        end_points['spn_rois'] = spn_rois #[B, SPN_NMS_MAX_SIZE, 6]
        end_points['rois'] = rois #[B, NUM_ROIS, 6]
        end_points['target_class_ids'] = target_class_ids #[B, NUM_ROIS]
        end_points['target_bbox'] = target_bbox #[B, NUM_ROIS, 6]
        end_points['target_mask_selection_idx'] = target_mask_selection_idx #[B, NUM_ROIS, NUM_POINT_PER_ROI]
        end_points['target_mask'] = target_mask #[B, NUM_ROIS, NUM_POINT_PER_ROI]
        end_points['mrcnn3d_class_logits'] = mrcnn3d_class_logits #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_class'] = mrcnn3d_class #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_bbox'] = mrcnn3d_bbox #[B, NUM_ROIS, NUM_CATEGORY, 6]
        end_points['mrcnn3d_mask'] = mrcnn3d_mask #[B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
    elif mode=='inference':
        end_points['selected_indices'] = selected_indices #[B, SPN_NMS_MAX_SIZE]
        end_points['spn_rois'] = spn_rois #[B, SPN_NMS_MAX_SIZE, 6]
        end_points['rois'] = rois #[B, NUM_ROIS, 6]
        end_points['mrcnn3d_class_logits'] = mrcnn3d_class_logits #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_class'] = mrcnn3d_class #[B, NUM_ROIS, NUM_CATEGORY]
        end_points['mrcnn3d_bbox'] = mrcnn3d_bbox #[B, NUM_ROIS, NUM_CATEGORY, 6]
        end_points['detections'] = detections #[B, DETECTION_MAX_INSTANCES, 6+2]
        end_points['mrcnn3d_mask'] = mrcnn3d_mask #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NUM_CATEGORY]
        # end_points['mrcnn3d_mask_unmolded'] = mrcnn3d_mask_unmolded #[B, DETECTION_MAX_INSTANCES, NUM_POINT]
        end_points['mrcnn3d_mask_selected'] = mrcnn3d_mask_selected #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI]
        end_points['pc_coord_cropped_final_unnormalized'] = pc_coord_cropped_final_unnormalized #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3]
    return end_points#, alexnetmodel

def get_sem_class_loss(sem_class_logits, labels, smpw):
    '''
    Inputs:
       sem_class_logits: [B, npoint, num_category]
       label: [B, npoint]
    '''
    labels = tf.cast(labels, tf.int32)
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=sem_class_logits, weights=smpw)
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=sem_class_logits)
    loss = tf.divide(tf.reduce_sum(tf.multiply(loss, smpw)), tf.reduce_sum(smpw)+1e-8)
    return loss

def smooth_l1_loss(y_true, y_pred):
    """Implements Smooth-L1 loss.
    y_true and y_pred are typically: [N, 6], but could be any shape.
    """
    diff = tf.abs(y_true - y_pred)
    less_than_one = tf.cast(tf.less(diff, 1.0), tf.float32)
    loss = (less_than_one * 0.5 * diff**2) + (1 - less_than_one) * (diff - 0.5)
    return loss

def get_spn_class_loss(fb_logits, spn_match):
    '''
    Inputs:
       fb_logits: [B, nsmp, 2]
       spn_match: [B, nsmp]
    '''
    # Only postive and negative contribute to loss but not neutral
    fb_logits = tf.reshape(fb_logits, [-1, 2])
    spn_match = tf.reshape(spn_match, [-1])
    valid_mask = tf.not_equal(spn_match, 0)
    fb_logits = tf.boolean_mask(fb_logits, valid_mask)
    spn_match = tf.cast(tf.equal(spn_match, 1), tf.int32)
    spn_match = tf.boolean_mask(spn_match, valid_mask)
    loss = tf.cond(tf.size(spn_match)>0,
        lambda: tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=spn_match, logits=fb_logits)),
        lambda: tf.constant(0.0))

    return loss

def focal_loss(prediction_tensor, target_tensor, alpha=0.25, gamma=2):
    """Compute focal loss for predictions.
        Multi-labels Focal loss formula:
            FL = -alpha * (z-p)^gamma * log(p) -(1-alpha) * p^gamma * log(1-p)
                 ,which alpha = 0.25, gamma = 2, p = sigmoid(x), z = target_tensor.
    Args:
     prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
     target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
     alpha: A scalar tensor for focal loss alpha hyper-parameter
     gamma: A scalar tensor for focal loss gamma hyper-parameter
    Returns:
        loss: A (scalar) tensor representing the value of the loss function
    """
    sigmoid_p = tf.nn.sigmoid(prediction_tensor)
    zeros = array_ops.zeros_like(sigmoid_p, dtype=sigmoid_p.dtype)
    
    # For poitive prediction, only need consider front part loss, back part is 0;
    # target_tensor > zeros <=> z=1, so poitive coefficient = z - p.
    pos_p_sub = array_ops.where(target_tensor > zeros, target_tensor - sigmoid_p, zeros)
    
    # For negative prediction, only need consider back part loss, front part is 0;
    # target_tensor > zeros <=> z=1, so negative coefficient = 0.
    neg_p_sub = array_ops.where(target_tensor > zeros, zeros, sigmoid_p)
    per_entry_cross_ent = - alpha * (pos_p_sub ** gamma) * tf.log(tf.clip_by_value(sigmoid_p, 1e-8, 1.0)) \
                          - (1 - alpha) * (neg_p_sub ** gamma) * tf.log(tf.clip_by_value(1.0 - sigmoid_p, 1e-8, 1.0))
    return tf.reduce_sum(per_entry_cross_ent, -1)

def get_mrcnn3d_class_focal_loss(mrcnn3d_class_logits, gt_class_ids, roi_valid_mask, config):
    '''
    Inputs:
       mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
       gt_class_ids: [B, NUM_ROIS], zero padded
       roi_valid_mask: [B, NUM_ROIS]
    '''
    gt_class_ids_onehot = tf.one_hot(gt_class_ids, depth=config.NUM_CATEGORY, axis=-1)
    loss = focal_loss(mrcnn3d_class_logits, gt_class_ids_onehot)
    loss = tf.multiply(loss, roi_valid_mask)
    # loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(roi_valid_mask)+1e-8)
    loss = tf.reduce_sum(loss)

    return loss

def get_mrcnn3d_class_loss(mrcnn3d_class_logits, gt_class_ids, roi_valid_mask):
    '''
    Inputs:
       mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]
       gt_class_ids: [B, NUM_ROIS], zero padded
       roi_valid_mask: [B, NUM_ROIS]
    '''
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=gt_class_ids, logits=mrcnn3d_class_logits)
    loss = tf.multiply(loss, roi_valid_mask)
    loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(roi_valid_mask)+1e-8)

    return loss

def get_mrcnn3d_bbox_loss(gt_bbox, gt_class_ids, pred_bbox, roi_valid_mask, num_category):
    '''
    Inputs:
       gt_bbox: [B, NUM_ROIS, 6]
       gt_class_ids: [B, NUM_ROIS], zero padded
       pred_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6]
       roi_valid_mask: [B, NUM_ROIS]
       num_category: scalar
    '''
    # Only foreground box contribute to the loss
    gt_bbox = tf.reshape(gt_bbox, [-1, 6])
    gt_class_ids = tf.reshape(gt_class_ids, [-1])
    pred_bbox = tf.reshape(pred_bbox, [-1, num_category, 6])
    roi_valid_mask = tf.reshape(roi_valid_mask, [-1])

    gt_selected_indices = tf.where(tf.logical_and(tf.greater(roi_valid_mask, 0), tf.greater(gt_class_ids, 0)))[:,0]
    gt_selected_indices = tf.cast(gt_selected_indices, tf.int32)
    pred_selected_indices = tf.concat((tf.reshape(gt_selected_indices, [-1,1]),
        tf.reshape(tf.gather(gt_class_ids, gt_selected_indices), [-1,1])), axis=1)

    gt_bbox = tf.gather(gt_bbox, gt_selected_indices, axis=0)
    pred_bbox = tf.gather_nd(pred_bbox, pred_selected_indices)

    loss = tf.cond(tf.size(gt_bbox)>0,
        lambda: tf.reduce_mean(tf.reduce_sum(smooth_l1_loss(y_true=gt_bbox, y_pred=pred_bbox),1),0),
        lambda: tf.constant(0.0))

    return loss

def get_mrcnn3d_mask_loss(gt_masks, gt_class_ids, pred_masks, roi_valid_mask, num_category, num_point_per_roi):
    '''
    Inputs:
       gt_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI]
       gt_class_ids: [B, NUM_ROIS], zero padded
       pred_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]
       roi_valid_mask: [B, NUM_ROIS]
       num_category: scalar
       num_point_per_roi: scalar
    '''
    # Only foreground box contribute to the loss
    gt_masks = tf.reshape(gt_masks, [-1, num_point_per_roi])
    gt_class_ids = tf.reshape(gt_class_ids, [-1])
    pred_masks = tf.reshape(pred_masks, [-1, num_point_per_roi, num_category])
    pred_masks = tf.transpose(pred_masks, perm=[0,2,1])
    roi_valid_mask = tf.reshape(roi_valid_mask, [-1])

    gt_selected_indices = tf.where(tf.logical_and(tf.greater(roi_valid_mask, 0), tf.greater(gt_class_ids, 0)))[:,0]
    gt_selected_indices = tf.cast(gt_selected_indices, tf.int32)
    pred_selected_indices = tf.concat((tf.reshape(gt_selected_indices, [-1,1]),
        tf.reshape(tf.gather(gt_class_ids, gt_selected_indices), [-1,1])), axis=1)

    gt_masks = tf.gather(gt_masks, gt_selected_indices, axis=0) # [N, NUM_POINT_PER_ROI]
    gt_masks = tf.cast(gt_masks, tf.float32)
    pred_masks = tf.gather_nd(pred_masks, pred_selected_indices) # [N, NUM_POINT_PER_ROI]

    loss = tf.cond(tf.size(gt_masks)>0,
        lambda: tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=gt_masks, logits=pred_masks)),
        lambda: tf.constant(0.0))

    return loss

def get_loss(end_points, config, alpha, smpw, mode='training'):
    batch_size = config.BATCH_SIZE
    nsmp = config.NUM_SAMPLE
    nsmp_ins = config.NUM_POINT_INS
    bbox_size = tf.reduce_max(end_points['pc_ins_centered_seed'], axis=2, keep_dims=True)-tf.reduce_min(end_points['pc_ins_centered_seed'], axis=2, keep_dims=True)
    radius = 1e-8 + tf.sqrt(tf.reduce_sum(tf.square(bbox_size/2), axis=-1, keep_dims=True)) # [B, nsmp, 1, 1]
    # radius = 1e-8 + tf.reduce_max(tf.sqrt(tf.reduce_sum(tf.square(end_points['pc_ins_centered_seed']), axis=-1, keep_dims=True)), axis=2, keep_dims=True) # [B, nsmp, 1, 1]
    shift_gt_seed = end_points['pc_ins_center_seed']-tf.expand_dims(end_points['pc_seed'],2) # [B, nsmp, 1, 3]
    shift_dist = tf.sqrt(tf.reduce_sum(tf.square(shift_gt_seed), 3, keep_dims=True)+1e-8)
    shift_gt_seed_normalized_4d = tf.concat((tf.divide(shift_gt_seed, shift_dist), tf.divide(shift_dist, radius)), -1)
    shift_pred_seed_normalized_4d = tf.concat((tf.expand_dims(end_points['shift_pred_seed_4d'],2)[:,:,:,:3], tf.divide(tf.expand_dims(end_points['shift_pred_seed_4d'],2)[:,:,:,3:], radius)), -1)

    # Fg/Bg loss, spn_match: [B, nsmp]
    fb_score_gt = tf.squeeze(gather_selection(tf.expand_dims(end_points['group_label'],-1), end_points['ind_seed'], end_points['ind_seed'].get_shape()[1].value),-1)
    fb_score_gt = tf.cast(tf.greater(fb_score_gt,0), tf.float32) # [B, nsmp]
    spn_match = batch_slice(
        [end_points['bbox_ins_pred'], fb_score_gt, tf.cast(tf.greater(end_points['seg_label_per_group'],0),tf.float32), end_points['bbox_ins']],
        lambda w, x, y, z: spn_target_gen(w, x, y, z),
        batch_size, names=["spn_match"])
    spn_match = tf.stop_gradient(tf.cast(spn_match, tf.int32))
    end_points['spn_match'] = spn_match
    spn_class_loss = get_spn_class_loss(end_points['fb_logits'], spn_match)
    # spn_class_loss = tf.constant(0.0)

    # # Reconstruction loss
    # pc_ins_pred = end_points['pc_ins_pred']
    # pc_ins_pred_normalized = tf.reshape(tf.div(pc_ins_pred, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    # pc_ins_gt_normalized = tf.reshape(tf.div(end_points['pc_ins_centered_seed']+shift_gt_seed, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    # recon_valid_mask = tf.reshape(fb_score_gt, [-1])
    # recon_valid_mask = tf.stop_gradient(recon_valid_mask)
    # dists_forward,_,dists_backward,_ = tf_nndistance.nn_distance(pc_ins_pred_normalized, pc_ins_gt_normalized)
    # recons_loss = tf.reduce_mean(dists_forward+dists_backward, axis=-1) # B*nsmp
    # recons_loss = tf.divide(tf.reduce_sum(tf.multiply(recons_loss, recon_valid_mask)),
    #     tf.reduce_sum(recon_valid_mask)+1e-8)
    # # recons_loss = tf.constant(0.0)

    # Reconstruction loss & confidence loss
    pc_ins_pred = end_points['pc_ins_pred']
    pc_ins_pred_normalized = tf.reshape(tf.div(pc_ins_pred, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    pc_ins_gt_normalized = tf.reshape(tf.div(end_points['pc_ins_centered_seed']+shift_gt_seed, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3]
    recon_valid_mask = tf.reshape(fb_score_gt, [-1])
    recon_valid_mask = tf.stop_gradient(recon_valid_mask)
    end_points['recon_valid_mask'] = recon_valid_mask
    dists_forward,_,dists_backward,_ = tf_nndistance.nn_distance(pc_ins_pred_normalized, pc_ins_gt_normalized)
    # dists_forward [B*nsmp, nsmp_ins]
    dist_thres = 0.04
    width = 0.03
    dists = tf.reshape(dists_forward, [batch_size, nsmp, nsmp_ins])

    conf_gt = tf.reshape(1.0 / (1.0 + tf.exp(-5*(dists-dist_thres)/width)), [-1, nsmp_ins]) # [B*nsmp, nsmp_ins]
    end_points['conf_gt'] = conf_gt
    conf_pred = tf.reshape(end_points['pc_ins_conf_pred'], [-1, nsmp_ins]) # [B*nsmp, nsmp_ins]
    conf_loss = tf.reduce_mean(smooth_l1_loss(conf_gt, conf_pred), axis=-1) # [B*nsmp]
    conf_loss = tf.divide(tf.reduce_sum(tf.multiply(conf_loss, recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)

    recons_loss = tf.reduce_mean(dists_forward+dists_backward, axis=-1) # B*nsmp
    recons_loss = tf.divide(tf.reduce_sum(tf.multiply(recons_loss, recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)
    # recons_loss = tf.constant(0.0)

    # Shift loss
    shift_loss = tf.reduce_sum(smooth_l1_loss(shift_gt_seed_normalized_4d, shift_pred_seed_normalized_4d), axis=-1)
    shift_loss = tf.reduce_sum(tf.multiply(tf.reshape(shift_loss, [-1]), recon_valid_mask))
    shift_loss = tf.divide(shift_loss, tf.reduce_sum(recon_valid_mask)+1e-8)
    # shift_loss = tf.constant(0.0)

    # Sem loss
    ind_sem = end_points['ind_sem']
    nsmp_sem = ind_sem.get_shape()[1].value
    sem_labels = end_points['seg_label']
    ind_sem_aug = tf.tile(tf.reshape(tf.range(batch_size),[-1, 1]), [1, nsmp_sem])
    ind_sem_aug = tf.concat( (tf.reshape(ind_sem_aug, [-1, 1]), tf.reshape(ind_sem, [-1, 1])), 1 )
    sem_labels = tf.reshape(tf.gather_nd(sem_labels, ind_sem_aug), [batch_size, nsmp_sem])
    smpw = tf.reshape(tf.gather_nd(smpw, ind_sem_aug), [batch_size, nsmp_sem])
    sem_labels = tf.cast(sem_labels, tf.int32)
    sem_loss = tf.losses.sparse_softmax_cross_entropy(labels=sem_labels, logits=end_points['sem_class_logits'], weights=smpw)
    end_points['sem_labels'] = sem_labels
    # sem_loss = tf.constant(0.0)

    # KL loss
    mean = end_points['mean'] # [B, nsmp, 256]
    log_var = end_points['log_var']
    cmean = end_points['cmean']
    clog_var = end_points['clog_var']
    kl_loss = 0.5 * tf.reduce_mean( log_var - clog_var + (tf.exp(clog_var) + (mean-cmean)**2)/tf.exp(log_var) - 1.0, 2) 
    kl_loss = tf.divide(tf.reduce_sum(tf.multiply(tf.reshape(kl_loss, [-1]), recon_valid_mask)),
        tf.reduce_sum(recon_valid_mask)+1e-8)
    # kl_loss = tf.constant(0.0)

    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        # MRCNN3D classification loss
        roi_valid_mask = tf.cast(tf.not_equal(tf.reduce_sum(tf.abs(end_points['rois']), axis=-1),0), tf.float32)
        mrcnn3d_class_loss = get_mrcnn3d_class_loss(end_points['mrcnn3d_class_logits'], end_points['target_class_ids'], roi_valid_mask)
        # mrcnn3d_class_loss = get_mrcnn3d_class_focal_loss(end_points['mrcnn3d_class_logits'], end_points['target_class_ids'], roi_valid_mask, config)

        # MRCNN3D bbox loss
        mrcnn3d_bbox_loss = get_mrcnn3d_bbox_loss(end_points['target_bbox'], end_points['target_class_ids'], end_points['mrcnn3d_bbox'], roi_valid_mask, config.NUM_CATEGORY)

        # MRCNN3D mask loss
        mrcnn3d_mask_loss = get_mrcnn3d_mask_loss(end_points['target_mask'], end_points['target_class_ids'], end_points['mrcnn3d_mask'], roi_valid_mask, config.NUM_CATEGORY, config.NUM_POINT_INS_MASK)
    
    if 'SPN' in config.TRAIN_MODULE:
        loss = kl_loss * alpha + recons_loss + shift_loss + spn_class_loss + sem_loss + conf_loss
        if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
            loss += mrcnn3d_class_loss + mrcnn3d_bbox_loss + mrcnn3d_mask_loss
    elif 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        loss = mrcnn3d_class_loss + mrcnn3d_bbox_loss + mrcnn3d_mask_loss
    else:
        loss = tf.constant(0.0)

    # Store end_points
    end_points['spn_class_loss'] = spn_class_loss
    end_points['recons_loss'] = recons_loss
    end_points['shift_loss'] = shift_loss
    end_points['conf_loss'] = conf_loss
    end_points['sem_loss'] = sem_loss
    end_points['kl_loss'] = kl_loss
    if 'MRCNN' in config.TRAIN_MODULE and mode=='training':
        end_points['mrcnn3d_class_loss'] = mrcnn3d_class_loss
        end_points['mrcnn3d_bbox_loss'] = mrcnn3d_bbox_loss
        end_points['mrcnn3d_mask_loss'] = mrcnn3d_mask_loss
    end_points['loss'] = loss
    
    return loss, end_points


if __name__=='__main__':
    # boxes = np.array([[0,0,0,1.,1.,1.],[0.1,0.1,0.1,1.,1.,1.],[1.,1.,1.,1.,1.,1.],[0.8,0.8,0.8,1.,1.,1.]])
    # scores = np.array([1.0,1.0,0.5,1.0])
    # boxes = np.reshape(boxes,[1,4,6])
    # scores = np.reshape(scores,[1,4])
    # select_idx = nms_3d(boxes, scores, 4, iou_threshold=0.3, score_threshold=0.0)
    # print(select_idx)
    myconfig = config.Config()
    with tf.Graph().as_default():
        pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, bbox_ins_pl = placeholder_inputs(myconfig)
        end_points = mrcnn_3d(pc_pl, pc_ins_pl, group_label_pl, group_indicator_pl, seg_label_pl, bbox_ins_pl, myconfig, tf.constant(True), bn_decay=None)
        loss, end_points = get_loss(end_points, myconfig, 1.0)
        print(end_points)

Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor             npoint: int32 -- #points sampled in farthest point sampling             radius: list of float32 -- search radius in local region             nsample: list of int32 -- how many points in each local region             mlp: list of list of int32 -- output size for MLP on each point             use_xyz: bool, if True concat XYZ with local point features, otherwise just use point features         Return:             new_xyz: (batch_size, npoint, 3) TF tensor             new_points: (batch_size, npoint, mlp2[-1]) TF tensor             shift_pred: (batch_size, npoint, 3) TF tensor             fps_idx: (batch_size, npoint) TF tensor # (batch_size, npoint) # (batch_size, npoint, 3) # [B, nseed, nsmp, 3] # (batch_size, npoint, \sum_k{mlp[k][-1]}) Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor         Return:             pc_seed: (batch_size, npoint_seed, 3) TF tensor             shift_pred_seed_4d: (batch_size, npoint_seed, 4) TF tensor             ind_seed: (batch_size, npoint_seed) TF tensor # (batch_size, npoint_seed) # (batch_size, npoint_seed, 3) # do not use color for shift prediction # Layer 1 # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1') # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2') # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3') # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4') # Feature Propagation layers # FC layers Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor             npoint_sem: int32 -- #points to sample for fast training             num_category: int32 -- #output category             ind_seed: (batch_size, npoint_seed) sampling index of seed points         Return:             sem_fea_seed: (batch_size, npoint_seed, nfea)             sem_fea: (batch_size, npoint_sem, nfea)             sem_fea_full: (batch_size, ndataset, nfea)             sem_class_logits: (batch_size, npoint_sem, num_category)             ind_sem: (batch_size, npoint_sem) sampling index of sem points # (batch_size, npoint_sem) # (batch_size, npoint_sem, 3) # (batch_size, npoint_seed, 3) # Layer 1 # l1_xyz, l1_points, l1_indices = pointnet_sa_module(l0_xyz, l0_points, npoint=1024, radius=0.2, nsample=32, mlp=[32,32,64], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1') # l2_xyz, l2_points, l2_indices = pointnet_sa_module(l1_xyz, l1_points, npoint=256, radius=0.4, nsample=32, mlp=[64,64,128], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2') # l3_xyz, l3_points, l3_indices = pointnet_sa_module(l2_xyz, l2_points, npoint=64, radius=0.8, nsample=32, mlp=[128,128,256], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer3') # l4_xyz, l4_points, l4_indices = pointnet_sa_module(l3_xyz, l3_points, npoint=16, radius=1.6, nsample=32, mlp=[256,256,512], mlp2=None, group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer4') # with FPN # Feature Propagation layers # FC layers Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor             npoint_sem: int32 -- #points to sample for fast training             num_category: int32 -- #output category             ind_seed: (batch_size, npoint_seed) sampling index of seed points         Return:             sem_fea_seed: (batch_size, npoint_seed, nfea)             sem_fea: (batch_size, npoint_sem, nfea)             sem_fea_full: (batch_size, ndataset, nfea)             sem_class_logits: (batch_size, npoint_sem, num_category)             ind_sem: (batch_size, npoint_sem) sampling index of sem points # Layer 1 # Feature Propagation layers # FC layers Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor             npoint_sem: int32 -- #points to sample for fast training             num_category: int32 -- #output category             ind_seed: (batch_size, npoint_seed) sampling index of seed points         Return:             sem_fea_seed: (batch_size, npoint_seed, nfea)             sem_fea: (batch_size, npoint_sem, nfea)             sem_fea_full: (batch_size, ndataset, nfea)             sem_class_logits: (batch_size, npoint_sem, num_category)             ind_sem: (batch_size, npoint_sem) sampling index of sem points # (batch_size, npoint_sem) # (batch_size, npoint_sem, 3) # (batch_size, npoint_seed, 3) # Layer 1 # with FPN # Feature Propagation layers # FC layers # net = end_points['sem_fea_full'] Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor             npoint_sem: int32 -- #points to sample for fast training             num_category: int32 -- #output category             ind_seed: (batch_size, npoint_seed) sampling index of seed points         Return:             sem_fea_seed: (batch_size, npoint_seed, nfea)             sem_fea: (batch_size, npoint_sem, nfea)             sem_fea_full: (batch_size, ndataset, nfea)             sem_class_logits: (batch_size, npoint_sem, num_category)             ind_sem: (batch_size, npoint_sem) sampling index of sem points # Layer 1 # Feature Propagation layers # FC layers Encode multiple context.         Input:             xyz: (batch_size, ndataset, 3) TF tensor             points: (batch_size, ndataset, channel) TF tensor         Return:             new_points: (batch_size, ndataset, channel_out) TF tensor # Layer 1 # Feature Propagation layers The encoding network for instance     Input:         pc: [B, N, 3]     Return:         net: [B, nfea] # Sample z The decoding network for shape generation     Input:         feat: [B, nsmp, nfea]     Return:         pc: [B, nsmp, num_point, 3] # UPCONV Decoder # FC Decoder # Merge The decoding network for shape generation & confidence generation for each generated point     Input:         feat: [B, nsmp, nfea]     Return:         pc: [B, nsmp, num_point, 3] # UPCONV Decoder # FC Decoder # Merge Shape proposal generation     Inputs:         pc: [B, NUM_POINT, 3]         pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys         group_label: [B, NUM_POINT]         group_indicator: [B, NUM_GROUP]         bbox_ins_gt: [B, NUM_GROUP, 6], only used during training     Returns:         fb_logits: [B, NUM_SAMPLE, 2] confidence logits (before softmax)         fb_prob: [B, NUM_SAMPLE, 2] confidence probabilities         bbox_ins: [B, NUM_SAMPLE, (x, y, z, l, w, h)]         entity_fea: [B, NUM_POINT, nfea] entity feature for each point         center_pos: [B, NUM_POINT, 3] center coordinate for each point, in world coord sys # Parameter extraction # Shift prediction, ind_seed [B, nsmp], shift_pred_seed [B, nsmp, 3] # Semantic prediction, sem_fea_seed [B, nsmp, nfea] # end_points = sem_net2(pc, None, 4096, num_category, ind_seed, end_points, 'sem_predictor', is_training, bn_decay=bn_decay, return_fullfea=return_fullfea) # Encode instance, pcfea_ins_centered [B, ngroup, nfea_ins], pc_ins_center [B, ngroup, 1, 3] # [B, ngroup, 1, 3] -> requires random padding for pc_ins generation # [B, ngroup, nfea_ins] # Collect instance feature for seed points [B, nsmp, nfea_seed] # [B, nsmp] # Encode context, pcfea_seed [B, nsmp, nfea_context] # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.25,0.5,1.0,1.5], [128,256,256,512], [[64,64,128], [64,128,256], [64,128,256], [64,128,256]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed) # _, pcfea_seed, _, _ = multi_encoding_net(pc, None, nsmp, [0.4,0.8,1.5], [192,256,384], [[64,128,192], [64,128,192], [64,128,192]], [], is_training, bn_decay, scope='context_encoder', use_xyz=True, output_shift=False, shift_pred=tf.stop_gradient(shift_pred_seed), fps_idx=ind_seed) # Compute foreground/background score [B, nsmp, 2] # fb_logits = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 64, 2], 'fb_logits', is_training, bn_decay) #### CVAE # Compute mu and sigma [B, nsmp, 512] # mu_sigma_c = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed), axis=-1), [256, 512, 512], 'mu_sigma_c', is_training, bn_decay) # mu_sigma_x = fea_trans_net(tf.concat((sem_fea_seed, pcfea_seed, pcfea_ins_seed), axis=-1), [256, 512, 512], 'mu_sigma_x', is_training, bn_decay) # mu_sigma_x = fea_trans_net(pcfea_ins_seed, [256, 512, 512], 'mu_sigma_x', is_training, bn_decay) # Sample z [B, nsmp, 256] # Decode shapes pc [B, nsmp, nsmp_ins, 3] # pc_ins_pred = decoding_net(feat, nsmp_ins, 'decoder', is_training=is_training, bn_decay=bn_decay) # pc_ins_pred = pc_ins_pred + tf.expand_dims(shift_pred_seed, 2) # Collect bbox for reconstructions, pc_ins_conf_pred [B, nsmp, nsmp_ins] # [B, nsmp, 6] -> center + l,w,h # pc_ins_pred_world_coord = pc_ins_pred + tf.expand_dims(pc_seed, 2) # [B, nsmp, nsmp_ins, 3] # inconf_pt_mask = 1000.0*tf.cast(tf.greater(pc_ins_conf_pred,0.5), tf.float32) # inconf_pt_mask = tf.multiply(inconf_pt_mask, 1-tf.cast(tf.less_equal(pc_ins_conf_pred, 1e-3+tf.reduce_min(pc_ins_conf_pred, 2, keep_dims=True)), tf.float32)) # inconf_pt_mask = tf.expand_dims(inconf_pt_mask, -1) # [B, nsmp, nsmp_ins, 1] # max_edge = tf.reduce_max(pc_ins_pred_world_coord-inconf_pt_mask, 2) # min_edge = tf.reduce_min(pc_ins_pred_world_coord+inconf_pt_mask, 2) # bbox_ins_pred = tf.concat(((max_edge+min_edge)/2, #                           max_edge-min_edge), 2) # [B, nsmp, 6] -> center + l,w,h # # Random shift and scale box during training # bbox_ins_pred = tf.concat( (bbox_ins_pred[:,:,:3]+bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=0.0, stddev=0.1), #     bbox_ins_pred[:,:,3:]*tf.random_normal(tf.shape(bbox_ins_pred[:,:,3:]), mean=1.0, stddev=0.1)), axis=-1) # Propagate seed feature and center position # shift_pred = pointnet_fp_module(pc, pc_seed, None, shift_pred_seed, [], is_training=False, bn_decay=None, scope='shift_prop', bn=False) # [B, N, 256] entity feature of each point # [B, N, 3] center location in the world coord sys # Store end_points # [B, nsmp, 3], offset from seed point to ins center # [B, nsmp, 4], offset from seed point to ins center # [B, nsmp, 3], seed point coordinate in world coord sys # [B, nsmp], seed index # [B, nsmp, nsmp_ins, 3], centered gt instance point cloud for each seed # [B, nsmp, 1, 3], gt instance center for each seed in world coord sys # [B, nsmp, 256] # [B, nsmp, 2] foreground/backgroud logits # [B, nsmp, 2] foreground/background probability # [B, nsmp, nsmp_ins, 3], in local sys, needs to add pc_seed # [B, nsmp, nsmp_ins] per-point confidence # [B, nsmp, 6] #, alexnetmodel Non maximum suppression in 3D     Inputs:         boxes: [B, N, 6] center + l,w,h         scores: [B, N] prob between 0 and 1     Outputs:         selected_indices: [B, M] # [B, N] from large to small # [B, M] Inputs:         source: [B, N, C]         selected_idx: [B, M], -1 means not selecting anything     Returns:         target: [B, M, C], 0 padded Often boxes are represented with matrices of shape [N, 4] and     are padded with zeros. This removes zero boxes.      boxes: [N, 6] matrix of boxes.     non_zeros: [N] a 1D boolean mask identifying the rows to keep 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. # 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 Shrink bounding box so that it is tight with respect to pc     Inputs:         box: [B, NUM_SAMPLE, 6]         pc: [B, NUM_POINT, 3]     Returns:         box: [B, NUM_SAMPLE, 6] # [B, 1, NUM_POINT, 3] # [B, NUM_SAMPLE, 1, 6] # [B, NUM_SAMPLE, NUM_POINT, 3] # [B, NUM_SAMPLE, NUM_POINT] # [B, NUM_SAMPLE, NUM_POINT, 1] # a large number for the box estimation trick # [B, NUM_SAMPLE, 3] # [B, NUM_SAMPLE, 6] Compute refinement needed to transform box to gt_box.     box and gt_box are [N, (center_x, center_y, center_z, l, w, h)] Apply bounding box delta to refine box     Inputs:         box: [NUM_SAMPLE, 6]         delta: [NUM_SAMPLE, 6]     Returns:         box_refined: [NUM_SAMPLE, 6] Inputs:         masks: [nmask, npoint]         nsmp: scalar     Returns:         masks_selection_idx: [nmask, nsmp] SPN target generation     Inputs:         proposals: [NUM_SAMPLE, 6]         proposal_seed_class_ids: [NUM_SAMPLE] - 0 is background class and 1 is foreground class         gt_class_ids: [NUM_GROUP] - 0 is background class and 1 is foreground class         gt_boxes: [NUM_GROUP, 6]     Returns:         spn_match: [NUM_SAMPLE], 1 = positive, -1 = negative, 0 = neutral # Remove zero padding # Compute IoU [NUM_SAMPLE, NUM_GROUP_TRIMMED] # Determine positive and negative ROIs # 1. Positive ROIs are those with >=0.5 IoU with a GT box and gt class id is 1 # We also guarantee that the boxes with the largest IoU with a GT box is also treated as positive # 2. Negative ROIs are those with < 0.5 with every GT box and not a postive ROI Generate detection targets for training     Inputs:         proposals: [SPN_NMS_MAX_SIZE, 6], zero padded         gt_class_ids: [NUM_GROUP] - 0 is background class and objects start from 1         gt_boxes: [NUM_GROUP, 6]         gt_masks: [NUM_POINT, NUM_GROUP]         pc: [NUM_POINT, 3]     Returns:         rois: [TRAIN_ROIS_PER_IMAGE, 6], zero padded         roi_gt_class_ids: [TRAIN_ROIS_PER_IMAGE] - 0 is invalid class and used for padding, zero padded         deltas: [TRAIN_ROIS_PER_IMAGE, 6], zero padded         masks_selection_idx: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded         masks: [TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] - binary mask, zero padded # Remove zero padding # Remove empty proposals # [N, 1, 3] # [1, NP, 6] # [N, NP] # [N, NP'] # [NP', 6] # Compute IoU [n_proposal, n_gt_boxes] # Determine positive and negative ROIs # 1. Positive ROIs are those with >= 0.5 IoU with a GT box # 2. Negative ROIs are those with < 0.5 with every GT box # Subsample ROIs. Aim for 33% positive # Positive ROIs # Negative ROIs. Add enough to maintain positive:negative ratio. # Gather selected ROIs [POSITIVE_COUNT/NEGATIVE_COUNT, 6] # Assign positive ROIs to GT boxes. roi_gt_boxes: [POSITIVE_COUNT, 6], roi_gt_class_ids: [POSITIVE_COUNT] # Compute bbox refinement for positive ROIs. delta: [POSITIVE_COUNT, 6] # Compute mask targets # [NUM_POINT, POSITIVE_COUNT] # [POSITIVE_COUNT, NUM_POINT] # [POSITIVE_COUNT, NUM_POINT_INS_MASK] # [POSITIVE_COUNT, NUM_POINT_INS_MASK] # Append negative ROIs and pad bbox deltas and masks that # are not used for negative ROIs with zeros. Generate detection targets for training     Inputs:         proposals: [SPN_NMS_MAX_SIZE, 6], zero padded         pc: [NUM_POINT, 3]     Returns:         proposals: [NUM_ROIS, 6]         masks_selection_idx: [NUM_ROIS, NUM_POINT_INS_MASK] - which points are selected from the initial point cloud, zero padded # Remove zero padding # Remove empty proposals # [N, 1, 3] # [1, NP, 6] # [N, NP] # [N, NP'] # [NP', 6] # Generate mask selection index # [NP', NUM_POINT_INS_MASK] # Append negative ROIs and pad bbox deltas and masks that # are not used for negative ROIs with zeros. Crop points for network heads, in analogy to ROIAlign     Inputs:         pc: [B, NUM_POINT, 3]         pc_fea: [B, NUM_POINT, NFEA]         pc_center: [B, NUM_POINT, 3]         rois: [B, NUM_ROIS, 6], zero padded         masks_selection_idx: [B, NUM_ROIS, NUM_POINT_PER_ROI]     Returns:         pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]         pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]         pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3] # convert world coord to local # scale box to [1,1,1] Refine classified proposals and filter overlaps and return final     detections.     Inputs:         rois: [NUM_ROIS, 6], zero padded, in world coord sys         probs: [NUM_ROIS, NUM_CATEGORY] - 0 is background class and objects start from 1         deltas: [NUM_ROIS, NUM_CATEGORY, 6]         pc: [NUM_POINT, 3]         fb_prob: [NUM_ROIS]         sem_prob: [NUM_ROIS]     Returns:         detections: [NUM_DETECTIONS, (center_x, center_y, center_z, l, w, h, class_id, score)] # Class IDs per ROI # Class probability of the top class of each ROI # [NUM_ROIS] # Class-specific bounding box deltas # [NUM_ROIS, 6] # Apply bounding box deltas # Shape: [NUM_ROIS, (center_x, center_y, center_z, l, w, h)] # Shrink boxes # Filter out background boxes # keep = tf.where(tf.logical_and(tf.logical_and(class_ids > 0, fb_prob>0.01), sem_prob>0.01))[:, 0] # Filter out low confidence boxes # Apply per-class NMS # 1. Prepare variables # [NUM_KEEP] # pre_nms_scores = tf.gather(class_scores*fb_prob*sem_prob, keep) # [NUM_KEEP] Apply Non-Maximum Suppression on ROIs of the given class. # Indices of ROIs of the given class # Apply NMS # [DETECTION_MAX_INSTANCES], -1 padded # [<=DETECTION_MAX_INSTANCES], no padding # Map indices # Pad with -1 so returned tensors have the same shape # Set shape so map_fn() can infer result shape # 2. Map over class IDs # 3. Merge results into one list, and remove -1 padding # 4. Compute intersection between keep and nms_keep # Keep top detections # class_scores_keep = tf.gather(class_scores*fb_prob*sem_prob, keep) # Arrange output as [N, (center_x, center_y, center_z, l, w, h, class_id, score)] # Pad with zeros if detections < DETECTION_MAX_INSTANCES Classification head for both class id prediction and bbox delta regression     Inputs:         pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]         pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]         num_category: scalar     Returns:         logits: [B, NUM_ROIS, NUM_CATEGORY]         probs: [B, NUM_ROIS, NUM_CATEGORY]         bbox_deltas: [B, NUM_ROIS, NUM_CATEGORY, (dz, dy, dx, log(dh), log(dw), log(dl))] Segmentation head     Inputs:         pc: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]         pc_fea: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA]         num_category: scalar     Returns:         masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY] Convert segmentation into point cloud label     Inputs:         mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]         class_ids: [B, NUM_ROIS]     Returns:         mrcnn3d_mask_selected: [B, NUM_ROIS, NUM_POINT_PER_ROI] # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI] #[-1, NUM_POINT_PER_ROI] Convert segmentation into point cloud label     Inputs:         mrcnn3d_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]         rois: [B, NUM_ROIS, 6]         class_ids: [B, NUM_ROIS]         pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3]         pc: [B, NUM_POINT, 3]     Returns:         mrcnn3d_mask_unmolded: [B, NUM_ROIS, NUM_POINT] # [B*NUM_ROIS, NUM_CATEGORY, NUM_POINT_PER_ROI] #[-1, NUM_POINT_PER_ROI] # [B, NUM_ROIS, NUM_POINT, NUM_POINT_PER_ROI] # [B, NUM_ROIS, NUM_POINT] # [-1, NUM_POINT] # Mask out regions outside rois # [B, 1, NUM_POINT, 3] # [B, NUM_ROIS, 1, 6] # [B, NUM_ROIS, NUM_POINT] Shape proposal generation     Inputs:         pc: [B, NUM_POINT, 3]         color: [B, NUM_POINT, 3]         pc_ins: [B, NUM_GROUP, NUM_POINT_INS, 3], in world coord sys         group_label: [B, NUM_POINT]         group_indicator: [B, NUM_GROUP]         seg_label: [B, NUM_POINT]         bbox_ins: [B, NUM_GROUP, 6]     Returns: #[B, NUM_POINT, NUM_GROUP] # seg_label_per_group = tf.multiply(tf.cast(tf.expand_dims(seg_label,-1), tf.float32), group_label_onehot) # seg_label_per_group = tf.cast(tf.round(tf.divide(tf.reduce_sum(seg_label_per_group, 1),tf.reduce_sum(group_label_onehot, 1)+1e-8)), tf.int32) #[B, NUM_GROUP] # 3D non maximum suppression - selected_indices: [B, M], spn_rois: [B, M, 6] # Detection target generation - rois: [B, TRAIN_ROIS_PER_IMAGE, 6], target_class_ids: [B, TRAIN_ROIS_PER_IMAGE] # target_bbox: [B, TRAIN_ROIS_PER_IMAGE, 6], target_mask_selection_idx: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK] # target_mask: [B, TRAIN_ROIS_PER_IMAGE, NUM_POINT_INS_MASK], all zero padded # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA] # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3] # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3] ##### sem fea # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### color # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], color), -1), #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### pn2 # pc_pn2_fea = pn2_fea_extractor(pc, None, 'pn2_fea_extractor', is_training, bn_decay=bn_decay) # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### sem fpn fea # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), #     end_points['center_pos'], rois, target_mask_selection_idx, config.TRAIN_ROIS_PER_IMAGE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) # Classification and bbox refinement head # mrcnn3d_class_logits, mrcnn3d_class, mrcnn3d_bbox = classification_head(pc_coord_cropped, #     pc_fea_cropped, config.NUM_CATEGORY, #     [128, 256, 512], [256, 256], is_training, bn_decay, 'classification_head') # Mask prediction head # mrcnn3d_mask = segmentation_head(pc_coord_cropped, #     pc_fea_cropped, config.NUM_CATEGORY, #     [64, 64], [64, 128, 512], [256, 256], is_training, bn_decay, 'segmentation_head') # rois: [B, NUM_ROIS, 6] # Points cropping - pc_fea_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, NFEA] # pc_center_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3] # pc_coord_cropped: [B, NUM_ROIS, NUM_POINT_PER_ROI, 3] ##### sem fea # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### pn2 # pc_pn2_fea = pn2_fea_extractor(pc, color, 'pn2_fea_extractor', is_training, bn_decay=bn_decay) # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### sem fpn fea # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, pc_coord_cropped_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], sem_fea_full_l1, sem_fea_full_l2, sem_fea_full_l3, sem_fea_full_l4), -1), #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) #### generate fb_conf and sem_conf # pc: [B, N, 3], fb_prob: [B, nsmp, 2], pc_seed: [B, nsmp, 3] -> fb_prob: [B, N] #[B, NUM_POINT, NUM_CATEGORY] ##### pn2 # pc_fea_cropped, pc_center_cropped, pc_coord_cropped, _ = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea, tf.expand_dims(fb_prob, -1), sem_prob), -1), #     end_points['center_pos'], rois, mask_selection_idx, SPN_NMS_MAX_SIZE, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) # pc_fea_cropped, fb_prob_cropped, sem_prob_cropped = tf.split(pc_fea_cropped, [832, 1, config.NUM_CATEGORY], -1) ##### sem fpn fea # [B, NUM_ROIS] # [B, NUM_ROIS, NUM_CATEGORY] # Classification and bbox refinement head - mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY] # mrcnn3d_class: [B, NUM_ROIS, NUM_CATEGORY], # mrcnn3d_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6] # [B, NUM_ROIS] # Generate detections: [B, DETECTION_MAX_INSTANCES, (center_x, center_y, center_z, l, w, h, class_id, score)] # detections = batch_slice( #     [rois, mrcnn3d_class, mrcnn3d_bbox, pc], #     lambda x, y, w, z: refine_detections(x, y, w, z, config), #     config.BATCH_SIZE) # Re-crop point cloud for mask prediction # Points cropping - pc_fea_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NFEA] # pc_center_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3] # pc_coord_cropped_final: [B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3] ##### sem fea # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], end_points['sem_fea_full']), -1), #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### pn2 # pc_fea_cropped_final, pc_center_cropped_final, pc_coord_cropped_final, pc_coord_cropped_final_unnormalized = points_cropping(pc, tf.concat((end_points['entity_fea'], pc_pn2_fea), -1), #     end_points['center_pos'], rois_final, mask_selection_idx_final, config.DETECTION_MAX_INSTANCES, config.NUM_POINT_INS_MASK, config.NORMALIZE_CROP_REGION) ##### sem fpn fea # Mask prediction head # Unmold segmentation # mrcnn3d_mask_unmolded = unmold_segmentation(tf.nn.sigmoid(mrcnn3d_mask), rois_final, detections[:,:,6], pc_coord_cropped_final_unnormalized, pc) # Update end_points #[B, NUM_GROUP] #[B, NUM_GROUP, 6] #[B, SPN_NMS_MAX_SIZE] #[B, SPN_NMS_MAX_SIZE, 6] #[B, NUM_ROIS, 6] #[B, NUM_ROIS] #[B, NUM_ROIS, 6] #[B, NUM_ROIS, NUM_POINT_PER_ROI] #[B, NUM_ROIS, NUM_POINT_PER_ROI] #[B, NUM_ROIS, NUM_CATEGORY] #[B, NUM_ROIS, NUM_CATEGORY] #[B, NUM_ROIS, NUM_CATEGORY, 6] #[B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY] #[B, SPN_NMS_MAX_SIZE] #[B, SPN_NMS_MAX_SIZE, 6] #[B, NUM_ROIS, 6] #[B, NUM_ROIS, NUM_CATEGORY] #[B, NUM_ROIS, NUM_CATEGORY] #[B, NUM_ROIS, NUM_CATEGORY, 6] #[B, DETECTION_MAX_INSTANCES, 6+2] #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, NUM_CATEGORY] # end_points['mrcnn3d_mask_unmolded'] = mrcnn3d_mask_unmolded #[B, DETECTION_MAX_INSTANCES, NUM_POINT] #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI] #[B, DETECTION_MAX_INSTANCES, NUM_POINT_PER_ROI, 3] #, alexnetmodel Inputs:        sem_class_logits: [B, npoint, num_category]        label: [B, npoint] # loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=sem_class_logits, weights=smpw) Implements Smooth-L1 loss.     y_true and y_pred are typically: [N, 6], but could be any shape. Inputs:        fb_logits: [B, nsmp, 2]        spn_match: [B, nsmp] # Only postive and negative contribute to loss but not neutral Compute focal loss for predictions.         Multi-labels Focal loss formula:             FL = -alpha * (z-p)^gamma * log(p) -(1-alpha) * p^gamma * log(1-p)                  ,which alpha = 0.25, gamma = 2, p = sigmoid(x), z = target_tensor.     Args:      prediction_tensor: A float tensor of shape [batch_size, num_anchors,         num_classes] representing the predicted logits for each class      target_tensor: A float tensor of shape [batch_size, num_anchors,         num_classes] representing one-hot encoded classification targets      alpha: A scalar tensor for focal loss alpha hyper-parameter      gamma: A scalar tensor for focal loss gamma hyper-parameter     Returns:         loss: A (scalar) tensor representing the value of the loss function # For poitive prediction, only need consider front part loss, back part is 0; # target_tensor > zeros <=> z=1, so poitive coefficient = z - p. # For negative prediction, only need consider back part loss, front part is 0; # target_tensor > zeros <=> z=1, so negative coefficient = 0. Inputs:        mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]        gt_class_ids: [B, NUM_ROIS], zero padded        roi_valid_mask: [B, NUM_ROIS] # loss = tf.divide(tf.reduce_sum(loss), tf.reduce_sum(roi_valid_mask)+1e-8) Inputs:        mrcnn3d_class_logits: [B, NUM_ROIS, NUM_CATEGORY]        gt_class_ids: [B, NUM_ROIS], zero padded        roi_valid_mask: [B, NUM_ROIS] Inputs:        gt_bbox: [B, NUM_ROIS, 6]        gt_class_ids: [B, NUM_ROIS], zero padded        pred_bbox: [B, NUM_ROIS, NUM_CATEGORY, 6]        roi_valid_mask: [B, NUM_ROIS]        num_category: scalar # Only foreground box contribute to the loss Inputs:        gt_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI]        gt_class_ids: [B, NUM_ROIS], zero padded        pred_masks: [B, NUM_ROIS, NUM_POINT_PER_ROI, NUM_CATEGORY]        roi_valid_mask: [B, NUM_ROIS]        num_category: scalar        num_point_per_roi: scalar # Only foreground box contribute to the loss # [N, NUM_POINT_PER_ROI] # [N, NUM_POINT_PER_ROI] # [B, nsmp, 1, 1] # radius = 1e-8 + tf.reduce_max(tf.sqrt(tf.reduce_sum(tf.square(end_points['pc_ins_centered_seed']), axis=-1, keep_dims=True)), axis=2, keep_dims=True) # [B, nsmp, 1, 1] # [B, nsmp, 1, 3] # Fg/Bg loss, spn_match: [B, nsmp] # [B, nsmp] # spn_class_loss = tf.constant(0.0) # # Reconstruction loss # pc_ins_pred = end_points['pc_ins_pred'] # pc_ins_pred_normalized = tf.reshape(tf.div(pc_ins_pred, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3] # pc_ins_gt_normalized = tf.reshape(tf.div(end_points['pc_ins_centered_seed']+shift_gt_seed, radius), [-1, nsmp_ins, 3]) # [B*nsmp, nsmp_ins, 3] # recon_valid_mask = tf.reshape(fb_score_gt, [-1]) # recon_valid_mask = tf.stop_gradient(recon_valid_mask) # dists_forward,_,dists_backward,_ = tf_nndistance.nn_distance(pc_ins_pred_normalized, pc_ins_gt_normalized) # recons_loss = tf.reduce_mean(dists_forward+dists_backward, axis=-1) # B*nsmp # recons_loss = tf.divide(tf.reduce_sum(tf.multiply(recons_loss, recon_valid_mask)), #     tf.reduce_sum(recon_valid_mask)+1e-8) # # recons_loss = tf.constant(0.0) # Reconstruction loss & confidence loss # [B*nsmp, nsmp_ins, 3] # [B*nsmp, nsmp_ins, 3] # dists_forward [B*nsmp, nsmp_ins] # [B*nsmp, nsmp_ins] # [B*nsmp, nsmp_ins] # [B*nsmp] # B*nsmp # recons_loss = tf.constant(0.0) # Shift loss # shift_loss = tf.constant(0.0) # Sem loss # sem_loss = tf.constant(0.0) # KL loss # [B, nsmp, 256] # kl_loss = tf.constant(0.0) # MRCNN3D classification loss # mrcnn3d_class_loss = get_mrcnn3d_class_focal_loss(end_points['mrcnn3d_class_logits'], end_points['target_class_ids'], roi_valid_mask, config) # MRCNN3D bbox loss # MRCNN3D mask loss # Store end_points # boxes = np.array([[0,0,0,1.,1.,1.],[0.1,0.1,0.1,1.,1.,1.],[1.,1.,1.,1.,1.,1.],[0.8,0.8,0.8,1.,1.,1.]]) # scores = np.array([1.0,1.0,0.5,1.0]) # boxes = np.reshape(boxes,[1,4,6]) # scores = np.reshape(scores,[1,4]) # select_idx = nms_3d(boxes, scores, 4, iou_threshold=0.3, score_threshold=0.0) # print(select_idx)