<filename>scannet/pointnet2_sem_seg.py
import os
import sys
BASE_DIR = os.path.dirname(__file__)
sys.path.append(BASE_DIR)
sys.path.append(os.path.join(BASE_DIR, '../utils'))
import tensorflow as tf
import numpy as np
import tf_util
from pointnet_util import pointnet_sa_module, pointnet_fp_module

def placeholder_inputs(batch_size, num_point):
    pointclouds_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point, 3))
    labels_pl = tf.placeholder(tf.int32, shape=(batch_size, num_point))
    smpws_pl = tf.placeholder(tf.float32, shape=(batch_size, num_point))
    return pointclouds_pl, labels_pl, smpws_pl

def get_model_from_arch(point_cloud, is_training, num_class, arch, bn_decay=None):
    ln_xyz = [tf.slice(point_cloud, [0, 0, 0], [-1, -1, 3])]    # point coordinates
    ln_feat_in = [tf.slice(point_cloud, [0, 0, 3], [-1, -1, -1])]  # point attributes
    ln_feat = [tf.slice(point_cloud, [0, 0, 3], [-1, -1, -1])]  # point attributes

    for depth, step_dict in enumerate(arch):  # set abstraction
        xyz, feat, _ = pointnet_sa_module(xyz=ln_xyz[depth], points=ln_feat[depth], is_training=is_training,
                                       npoint=step_dict['npoint'],
                                       radius=step_dict['radius'],
                                       nsample=step_dict['nsample'],
                                       mlp=step_dict['mlp'],
                                       pooling=step_dict['pooling'],
                                       mlp2=step_dict['mlp2'],
                                       group_all=False,
                                       bn_decay=bn_decay,
                                       scope='sa_layer_%d' % (depth + 1))
        ln_xyz.append(xyz)
        ln_feat.append(feat)
        ln_feat_in.append(feat)

    for depth, step_dict in enumerate(reversed(arch)):  # feature propagation
        depth = len(arch) - depth
        feat = pointnet_fp_module(xyz1=ln_xyz[depth-1], xyz2=ln_xyz[depth], 
                                  points1=ln_feat[depth-1], points2=ln_feat[depth], is_training=is_training,
                                  mlp=step_dict['reverse_mlp'], bn_decay=bn_decay,
                                  scope='fp_layer_%d' % (depth - 1))
        ln_feat[depth - 1] = feat

    l0_feats = ln_feat[0]
    net = tf_util.conv1d(l0_feats, 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')
    net = tf_util.conv1d(net, num_class, 1, padding='VALID', activation_fn=None, scope='fc2')
    return net


def get_model(point_cloud, is_training, num_class, bn_decay=None):
    """ Semantic segmentation PointNet, input is BxNx3, output Bxnum_class """
    l0_xyz = point_cloud
    l0_points = None

    # Layer 1
    l1_xyz, l1_points, _ = pointnet_sa_module(l0_xyz, l0_points, npoint=8192, radius=1, nsample=16, mlp=[64,512,128], mlp2=None, pooling='max_and_avg', group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer1')
    l2_xyz, l2_points, _ = pointnet_sa_module(l1_xyz, l1_points, npoint=4096, radius=5, nsample=64, mlp=[128,512,256], mlp2=None, pooling='max_and_avg', group_all=False, is_training=is_training, bn_decay=bn_decay, scope='layer2')
    l3_xyz, l3_points, _ = pointnet_sa_module(l2_xyz, l2_points, npoint=2048, radius=15, nsample=64, mlp=[128,512,256], mlp2=None, pooling='max_and_avg', 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=8, 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, [128,64], is_training, bn_decay, scope='fa_layer2')
    l1_points = pointnet_fp_module(l1_xyz, l2_xyz, l1_points, l2_points, [256,256], is_training, bn_decay, scope='fa_layer3')
    l0_points = pointnet_fp_module(l0_xyz, l1_xyz, l0_points, l1_points, [256,256], 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.7, is_training=is_training, scope='dp1')
    net = tf_util.conv1d(net, num_class, 1, padding='VALID', activation_fn=None, scope='fc2')

    return net


def get_loss(pred, label):
    """ pred: BxNxC,
        label: BxN """
    return tf.losses.sparse_softmax_cross_entropy(labels=label, logits=pred)

def fp_high_loss(labels, logits, factor=10):
    weights = tf.where(tf.logical_and(labels != 2, tf.argmax(logits) == 2), factor, 1)
    return tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits, scope='loss', weights=weights)

if __name__=='__main__':
    with tf.Graph().as_default():
        inputs = tf.zeros((32,2048,3))
        net, _ = get_model(inputs, tf.constant(True), 10)
        print(net)
