├── sampling-based ├── __init__.py ├── splatnet │ ├── __init__.py │ ├── dataset │ │ ├── __init__.py │ │ ├── __init__.pyc │ │ ├── dataset_lidar.pyc │ │ ├── indoor3d_util.pyc │ │ ├── dataset_modelnet.pyc │ │ ├── dataset_shapenet.pyc │ │ ├── indoor3d_util.py │ │ └── dataset_facade.py │ ├── partseg3d │ │ ├── __init__.py │ │ ├── __init__.pyc │ │ ├── models.pyc │ │ ├── test.py │ │ ├── train_sense.py │ │ └── train.py │ ├── .DS_Store │ ├── util.pyc │ ├── utils.pyc │ ├── __init__.pyc │ ├── configs.pyc │ ├── plot_log.pyc │ ├── create_solver.pyc │ ├── custom_layers.pyc │ ├── create_solver.py │ ├── configs.py │ ├── plot_log.py │ ├── semseg3d │ │ ├── eval_seg.py │ │ ├── train.py │ │ ├── test.py │ │ └── models.py │ ├── utils.py │ └── custom_layers.py ├── .gitignore ├── .DS_Store ├── exp │ ├── .DS_Store │ └── shapenet3d │ │ ├── .DS_Store │ │ ├── test_only.sh │ │ ├── train_test.sh │ │ ├── car_net_deform.prototxt │ │ ├── lamp_net_deform.prototxt │ │ └── mug_net_deform.prototxt ├── data │ ├── README.md │ └── shapenet_prepare.py └── tools │ └── pick_best_model.py ├── .gitignore ├── .DS_Store ├── assets └── intro.png ├── point_based ├── .DS_Store ├── utils │ ├── .DS_Store │ ├── data_prep_util.py │ └── pc_util.py └── part_seg │ ├── download_data.sh │ ├── transform_nets.py │ ├── provider.py │ ├── part_seg_model_deform.py │ └── test_deform.py └── README.md /sampling-based/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /sampling-based/splatnet/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | deformableSplatnet 2 | dgcnn-trans 3 | -------------------------------------------------------------------------------- /sampling-based/.gitignore: -------------------------------------------------------------------------------- 1 | *.caffemodel 2 | *.solverstate 3 | .log 4 | .pyc 5 | __pycache__/ 6 | -------------------------------------------------------------------------------- /.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/.DS_Store -------------------------------------------------------------------------------- /assets/intro.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/assets/intro.png -------------------------------------------------------------------------------- /point_based/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/point_based/.DS_Store -------------------------------------------------------------------------------- /sampling-based/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/.DS_Store -------------------------------------------------------------------------------- /point_based/utils/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/point_based/utils/.DS_Store -------------------------------------------------------------------------------- /sampling-based/exp/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/exp/.DS_Store -------------------------------------------------------------------------------- /sampling-based/splatnet/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/.DS_Store -------------------------------------------------------------------------------- /sampling-based/splatnet/util.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/util.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/utils.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/utils.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/__init__.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/configs.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/configs.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/plot_log.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/plot_log.pyc -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/exp/shapenet3d/.DS_Store -------------------------------------------------------------------------------- /sampling-based/splatnet/create_solver.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/create_solver.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/custom_layers.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/custom_layers.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/dataset/__init__.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/partseg3d/__init__.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/models.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/partseg3d/models.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/dataset_lidar.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/dataset/dataset_lidar.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/indoor3d_util.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/dataset/indoor3d_util.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/dataset_modelnet.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/dataset/dataset_modelnet.pyc -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/dataset_shapenet.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/samaonline/spatial-transformer-for-3d-point-clouds/HEAD/sampling-based/splatnet/dataset/dataset_shapenet.pyc -------------------------------------------------------------------------------- /point_based/part_seg/download_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Download original ShapeNetPart dataset (around 1GB) ['PartAnnotation'] 4 | wget https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_v0.zip --no-check-certificate 5 | unzip shapenetcore_partanno_v0.zip 6 | rm shapenetcore_partanno_v0.zip 7 | 8 | # Download HDF5 for ShapeNet Part segmentation (around 346MB) ['hdf5_data'] 9 | wget https://shapenet.cs.stanford.edu/media/shapenet_part_seg_hdf5_data.zip --no-check-certificate 10 | unzip shapenet_part_seg_hdf5_data.zip 11 | rm shapenet_part_seg_hdf5_data.zip 12 | -------------------------------------------------------------------------------- /sampling-based/data/README.md: -------------------------------------------------------------------------------- 1 | # datasets 2 | 3 | ### ShapeNet Part 4 | 1. Download and uncompress the file from the [PointNet++ repo](https://github.com/charlesq34/pointnet2). 5 | The link is located under the [Object Part Segmentation](https://github.com/charlesq34/pointnet2/blob/master/README.md#object-part-segmentation) 6 | section. 7 | 2. Convert to PLY files: 8 | ```bash 9 | # this step can take 10~15 minutes 10 | python shapenet_prepare.py 11 | ``` 12 | Check that folder `data/shapenet_ericyi_ply` is now generated filled with PLY files. 13 | 14 | 15 | ### References 16 | 17 | ShapeNet point clouds and annotations were originally collected by the authors of: 18 | - L. Yi, et al. A Scalable Active Framework for Region Annotation in 3D Shape Collections. SIGGRAPH Asia, 2016. 19 | 20 | Normal directions of the point clouds are provided by the authors of: 21 | - C. R. Qi, et al. PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space. NIPS, 2017. 22 | -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/test_only.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # enter environment if using conda 4 | # source activate caffe 5 | 6 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" 7 | 8 | if [ -z $EXP_DIR ]; then EXP_DIR=$SCRIPT_DIR; fi 9 | if [ -z $SPLT_CODE ]; then SPLT_CODE="$SCRIPT_DIR/../../splatnet"; fi 10 | if [ -z $SPLT_DATA ]; then SPLT_DATA="$SCRIPT_DIR/../../data"; fi 11 | if [ -z $SKIP_TRAIN ]; then SKIP_TRAIN=1; fi 12 | if [ -z $SKIP_TEST ]; then SKIP_TEST=0; fi 13 | if [ -z "$CATS" ]; then CATS="laptop car skateboard chair earphone motorbike mug airplane \ 14 | pistol lamp bag cap rocket guitar table knife"; fi 15 | 16 | mkdir -p $EXP_DIR 17 | 18 | # train 19 | if [ $SKIP_TRAIN -le 0 ]; then 20 | for CAT in $CATS; do 21 | python $SPLT_CODE/partseg3d/train.py $EXP_DIR \ 22 | --categories $CAT \ 23 | --arch c32_b64_b128_b256_b256_b256_c128 \ 24 | --skips 6_2 6_3 6_4 6_5 \ 25 | --lattice x*64_y*64_z*64 x*32_y*32_z*32 x*16_y*16_z*16 x*8_y*8_z*8 x*4_y*4_z*4 \ 26 | --feat x_y_z \ 27 | --dataset_params jitter_xyz 0.01 jitter_rotation 10 root $SPLT_DATA/shapenet_ericyi_ply \ 28 | --batch_size 32 --sample_size 3000 \ 29 | --base_lr 0.0001 --lr_decay 0.1 --stepsize 2000 --num_iter 2000 --test_interval 20 --snapshot_interval 100 \ 30 | 2>&1 | tee $EXP_DIR/$CAT.log; done ; 31 | python $SCRIPT_DIR/../../tools/pick_best_model.py $EXP_DIR best_loss 32 | fi 33 | 34 | # test & plot 35 | if [ $SKIP_TEST -le 0 ]; then 36 | python $SPLT_CODE/partseg3d/test.py shapenet \ 37 | --dataset_params root $SPLT_DATA/shapenet_ericyi_ply \ 38 | --categories $CATS \ 39 | --sample_size -1 --batch_size 1 \ 40 | --exp_dir $EXP_DIR ; 41 | fi 42 | -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/train_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # enter environment if using conda 4 | # source activate caffe 5 | 6 | SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" 7 | 8 | if [ -z $EXP_DIR ]; then EXP_DIR=$SCRIPT_DIR; fi 9 | if [ -z $SPLT_CODE ]; then SPLT_CODE="$SCRIPT_DIR/../../splatnet"; fi 10 | if [ -z $SPLT_DATA ]; then SPLT_DATA="$SCRIPT_DIR/../../data"; fi 11 | if [ -z $SKIP_TRAIN ]; then SKIP_TRAIN=0; fi 12 | if [ -z $SKIP_TEST ]; then SKIP_TEST=1; fi 13 | if [ -z "$CATS" ]; then CATS="rocket"; fi 14 | 15 | mkdir -p $EXP_DIR 16 | 17 | # train 18 | if [ $SKIP_TRAIN -le 0 ]; then 19 | for CAT in $CATS; do 20 | mkdir -p $CAT 21 | TEMP1="_net_deform.prototxt" 22 | TEMP2=".caffemodel" 23 | python $SPLT_CODE/partseg3d/train.py $EXP_DIR/$CAT \ 24 | --categories $CAT \ 25 | --arch c32_b64_b128_b256_b256_b256_c128 \ 26 | --skips 6_2 6_3 6_4 6_5 \ 27 | --lattice x*64_y*64_z*64 x*32_y*32_z*32 x*16_y*16_z*16 x*8_y*8_z*8 x*4_y*4_z*4 \ 28 | --feat x_y_z \ 29 | --gpus 1 \ 30 | --dataset_params jitter_xyz 0.01 jitter_rotation 10 root $SPLT_DATA/shapenet_ericyi_ply \ 31 | --batch_size 32 --sample_size 3000 \ 32 | --network $CAT$TEMP1 \ 33 | --prefix net_deform \ 34 | --base_lr 0.0001 --lr_decay 0.1 --stepsize 2000 --num_iter 15000 --test_interval 20 --snapshot_interval 100 \ 35 | 2>&1 | tee $EXP_DIR/$CAT/train_net_deform.log; done ; 36 | fi 37 | 38 | # test & plot 39 | if [ $SKIP_TEST -le 0 ]; then 40 | python $SPLT_CODE/partseg3d/test.py shapenet \ 41 | --dataset_params root $SPLT_DATA/shapenet_ericyi_ply \ 42 | --categories $CATS \ 43 | --snapshot best_loss \ 44 | --sample_size -1 --batch_size 1 \ 45 | --exp_dir $CAT/$EXP_DIR ; 46 | fi 47 | -------------------------------------------------------------------------------- /sampling-based/data/shapenet_prepare.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import json 4 | import numpy as np 5 | 6 | parser = argparse.ArgumentParser() 7 | parser.add_argument('data_dir') 8 | parser.add_argument('--save_dir', default='shapenet_ericyi_ply', required=False) 9 | args = parser.parse_args() 10 | data_dir = args.data_dir 11 | save_dir = args.save_dir 12 | 13 | SHAPE_CATEGORIES = ('03642806', '02958343', '04225987', '03001627', '03261776', '03790512', '03797390', '02691156', 14 | '03948459', '03636649', '02773838', '02954340', '04099429', '03467517', '04379243', '03624134') 15 | off = [28, 8, 44, 12, 16, 30, 36, 0, 38, 24, 4, 6, 41, 19, 47, 22] 16 | off_map = dict(zip(SHAPE_CATEGORIES, off)) 17 | 18 | CMAP = ((255, 255, 0), 19 | (128, 255, 255), 20 | (128, 0, 255), 21 | (255, 0, 0), 22 | (255, 128, 0), 23 | (0, 255, 0), 24 | (0, 0, 255)) 25 | 26 | header = '''ply 27 | format ascii 1.0 28 | element vertex {} 29 | property float x 30 | property float y 31 | property float z 32 | property float nx 33 | property float ny 34 | property float nz 35 | property uchar red 36 | property uchar green 37 | property uchar blue 38 | property uchar label 39 | end_header''' 40 | 41 | fmt = '%.6f %.6f %.6f %.6f %.6f %.6f %d %d %d %d' 42 | 43 | for subset in ('val', 'test', 'train'): 44 | sample_list_path = os.path.join(data_dir, 'train_test_split', 'shuffled_{}_file_list.json'.format(subset)) 45 | os.makedirs(os.path.join(save_dir, subset), exist_ok=True) 46 | sample_list = json.load(open(sample_list_path)) 47 | print('processing {} {} samples ... '.format(len(sample_list), subset), flush=True, end='') 48 | for sample in sample_list: 49 | category, sample_id = str(sample).split('/')[1:] 50 | os.makedirs(os.path.join(save_dir, subset, category), exist_ok=True) 51 | data = np.loadtxt(os.path.join(data_dir, category, '{}.txt'.format(sample_id))) 52 | data[:, -1] -= off_map[category] 53 | data = np.vstack([np.concatenate((d[:6], CMAP[int(d[-1])], [int(d[-1])])) for d in data]) 54 | np.savetxt(os.path.join(save_dir, subset, category, '{}.ply'.format(sample_id)), data, 55 | fmt=fmt, 56 | header=header.format(len(data)), 57 | comments='') 58 | print('done!', flush=True) 59 | -------------------------------------------------------------------------------- /sampling-based/splatnet/create_solver.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import tempfile 6 | from caffe.proto import caffe_pb2 7 | 8 | 9 | def get_prototxt(solver_proto, save_path=None): 10 | if save_path: 11 | f = open(save_path, mode='w+') 12 | else: 13 | f = tempfile.NamedTemporaryFile(mode='w+', delete=False) 14 | f.write(str(solver_proto)) 15 | f.close() 16 | 17 | return f.name 18 | 19 | 20 | def standard_solver(train_net, 21 | test_net, 22 | prefix, 23 | solver_type='SGD', 24 | weight_decay=0.001, 25 | base_lr=0.01, 26 | gamma=0.1, 27 | stepsize=100, 28 | test_iter=100, 29 | test_interval=1000, 30 | max_iter=1e5, 31 | iter_size=1, 32 | snapshot=1000, 33 | display=1, 34 | random_seed=0, 35 | debug_info=False, 36 | create_prototxt=True, 37 | save_path=None): 38 | 39 | solver = caffe_pb2.SolverParameter() 40 | solver.train_net = train_net 41 | solver.test_net.extend([test_net]) 42 | 43 | solver.test_iter.extend([test_iter]) 44 | solver.test_interval = test_interval 45 | 46 | solver.base_lr = base_lr 47 | solver.lr_policy = 'step' # "fixed" 48 | solver.gamma = gamma 49 | solver.stepsize = stepsize 50 | 51 | solver.display = display 52 | solver.max_iter = max_iter 53 | solver.iter_size = iter_size 54 | solver.snapshot = snapshot 55 | solver.snapshot_prefix = prefix 56 | solver.random_seed = random_seed 57 | 58 | solver.solver_mode = caffe_pb2.SolverParameter.GPU 59 | if solver_type is 'SGD': 60 | solver.solver_type = caffe_pb2.SolverParameter.SGD 61 | elif solver_type is 'ADAM': 62 | solver.solver_type = caffe_pb2.SolverParameter.ADAM 63 | solver.momentum = 0.9 64 | solver.momentum2 = 0.999 65 | 66 | solver.weight_decay = weight_decay 67 | 68 | solver.debug_info = debug_info 69 | 70 | if create_prototxt: 71 | solver = get_prototxt(solver, save_path) 72 | 73 | return solver 74 | 75 | -------------------------------------------------------------------------------- /point_based/part_seg/transform_nets.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import sys 4 | import os 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | sys.path.append(BASE_DIR) 7 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 8 | import tf_util 9 | 10 | def input_transform_net(edge_feature, is_training, bn_decay=None, K=3, is_dist=False): 11 | """ Input (XYZ) Transform Net, input is BxNx3 gray image 12 | Return: 13 | Transformation matrix of size 3xK """ 14 | batch_size = edge_feature.get_shape()[0].value 15 | num_point = edge_feature.get_shape()[1].value 16 | 17 | # input_image = tf.expand_dims(point_cloud, -1) 18 | net = tf_util.conv2d(edge_feature, 64, [1,1], 19 | padding='VALID', stride=[1,1], 20 | bn=True, is_training=is_training, 21 | scope='tconv1', bn_decay=bn_decay, is_dist=is_dist) 22 | net = tf_util.conv2d(net, 128, [1,1], 23 | padding='VALID', stride=[1,1], 24 | bn=True, is_training=is_training, 25 | scope='tconv2', bn_decay=bn_decay, is_dist=is_dist) 26 | 27 | net = tf.reduce_max(net, axis=-2, keep_dims=True) 28 | 29 | net = tf_util.conv2d(net, 1024, [1,1], 30 | padding='VALID', stride=[1,1], 31 | bn=True, is_training=is_training, 32 | scope='tconv3', bn_decay=bn_decay, is_dist=is_dist) 33 | net = tf_util.max_pool2d(net, [num_point,1], 34 | padding='VALID', scope='tmaxpool') 35 | 36 | net = tf.reshape(net, [batch_size, -1]) 37 | net = tf_util.fully_connected(net, 512, bn=True, is_training=is_training, 38 | scope='tfc1', bn_decay=bn_decay,is_dist=is_dist) 39 | net = tf_util.fully_connected(net, 256, bn=True, is_training=is_training, 40 | scope='tfc2', bn_decay=bn_decay,is_dist=is_dist) 41 | 42 | with tf.variable_scope('transform_XYZ') as sc: 43 | # assert(K==3) 44 | with tf.device('/cpu:0'): 45 | weights = tf.get_variable('weights', [256, K*K], 46 | initializer=tf.constant_initializer(0.0), 47 | dtype=tf.float32) 48 | biases = tf.get_variable('biases', [K*K], 49 | initializer=tf.constant_initializer(0.0), 50 | dtype=tf.float32) 51 | biases += tf.constant(np.eye(K).flatten(), dtype=tf.float32) 52 | transform = tf.matmul(net, weights) 53 | transform = tf.nn.bias_add(transform, biases) 54 | 55 | transform = tf.reshape(transform, [batch_size, K, K]) 56 | return transform -------------------------------------------------------------------------------- /sampling-based/splatnet/configs.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import sys 7 | 8 | # project root 9 | ROOT_DIR = os.path.abspath(os.path.join(__file__, '..', '..')) 10 | 11 | # required for custom layers 12 | sys.path.append(os.path.join(ROOT_DIR, 'splatnet')) 13 | sys.path.append(os.path.join(ROOT_DIR, 'splatnet', 'dataset')) 14 | 15 | # modify these if you put data in non-default locations 16 | FACADE_DATA_DIR = os.path.join(ROOT_DIR, 'data', 'ruemonge428') 17 | SHAPENET3D_DATA_DIR = os.path.join(ROOT_DIR, 'data', 'shapenet_ericyi_ply') 18 | SHAPENET2D3D_DATA_DIR = os.path.join(ROOT_DIR, 'data', 'shapenet_2d3d_h5') 19 | LIDAR3D_DATA_DIR = os.path.join(ROOT_DIR, 'data', 'stanford_indoor3d') 20 | MODELNET_DATA_DIR = os.path.join(ROOT_DIR, 'data', 'modelnet40_ply_hdf5_2048') 21 | # facade global variables 22 | 23 | FACADE_CATEGORIES = ('wall', 'sky', 'balcony', 'window', 'door', 'shop', 'roof') 24 | 25 | FACADE_CMAP = ((255, 255, 0), 26 | (128, 255, 255), 27 | (128, 0, 255), 28 | (255, 0, 0), 29 | (255, 128, 0), 30 | (0, 255, 0), 31 | (0, 0, 255)) 32 | 33 | # shapenet global variables 34 | 35 | SN_CATEGORIES = ('03642806', '02958343', '04225987', '03001627', '03261776', '03790512', '03797390', '02691156', 36 | '03948459', '03636649', '02773838', '02954340', '04099429', '03467517', '04379243', '03624134') 37 | 38 | SN_CATEGORY_NAMES = ('laptop', 'car', 'skateboard', 'chair', 'earphone', 'motorbike', 'mug', 'airplane', 39 | 'pistol', 'lamp', 'bag', 'cap', 'rocket', 'guitar', 'table', 'knife') 40 | 41 | SN_NUM_PART_CATEGORIES = (2, 4, 3, 4, 3, 6, 2, 4, 42 | 3, 4, 2, 2, 3, 3, 3, 2) 43 | 44 | SN_CMAP = ((255, 255, 0), 45 | (128, 255, 255), 46 | (128, 0, 255), 47 | (255, 0, 0), 48 | (255, 128, 0), 49 | (0, 255, 0), 50 | (0, 0, 255)) 51 | 52 | # stanford indoor 3d variables 53 | 54 | g_classes = [x.rstrip() for x in open(os.path.join(LIDAR3D_DATA_DIR, 'class_names.txt'))] 55 | g_class2label = {cls: i for i,cls in enumerate(g_classes)} 56 | g_class2color = {'ceiling': [0,255,0], 57 | 'floor': [0,0,255], 58 | 'wall': [0,255,255], 59 | 'beam': [255,255,0], 60 | 'column': [255,0,255], 61 | 'window': [100,100,255], 62 | 'door': [200,200,100], 63 | 'table': [170,120,200], 64 | 'chair': [255,0,0], 65 | 'sofa': [200,100,100], 66 | 'bookcase': [10,200,100], 67 | 'board': [200,200,200], 68 | 'clutter': [50,50,50]} 69 | g_easy_view_labels = [7,8,9,10,11,1] 70 | g_label2color = {g_classes.index(cls): g_class2color[cls] for cls in g_classes} 71 | 72 | # LIDAR_CATEGORIES -------------------------------------------------------------------------------- /sampling-based/tools/pick_best_model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import sys 7 | import glob 8 | import shutil 9 | import subprocess 10 | import numpy as np 11 | 12 | 13 | def parse_and_plot(path, skip_train=100, skip_test=10, caffe_root=None): 14 | 15 | if not caffe_root: 16 | import caffe 17 | caffe_root = os.path.join(os.path.dirname(caffe.__file__), '..', '..') 18 | 19 | parser_path = os.path.join(caffe_root, 'tools/extra/parse_log.py') 20 | subprocess.run('python2 {} {} {}'.format(parser_path, path, os.path.dirname(path)).split(' ')) 21 | 22 | stats_train = np.loadtxt(path + '.train', delimiter=',', skiprows=1) 23 | stats_test = np.loadtxt(path + '.test', delimiter=',', skiprows=1) 24 | 25 | return stats_train, stats_test 26 | 27 | 28 | def copy_best_model(exp_dir, categories, pick_rule='best_loss'): 29 | if os.path.exists(os.path.join(exp_dir, 'train.log')): 30 | raise NotImplementedError() 31 | 32 | for category in categories: 33 | log = os.path.join(exp_dir, category + '.log') 34 | snapshot_prefix = category 35 | iters_avail = [int(os.path.split(v)[1].split('_')[-1][:-11]) for v in 36 | glob.glob(os.path.join(exp_dir, snapshot_prefix + '_iter_*.caffemodel'))] 37 | if pick_rule == 'last': 38 | snap_iter = max(iters_avail) 39 | elif pick_rule == 'best_acc': 40 | _, stats_val = parse_and_plot(log, skip_train=0, skip_test=0) 41 | iter_acc = dict(zip(stats_val[:, 0], stats_val[:, -2])) 42 | snap_iter = iters_avail[np.argmax([iter_acc[v] for v in iters_avail])] 43 | elif pick_rule == 'best_loss': 44 | _, stats_val = parse_and_plot(log, skip_train=0, skip_test=0) 45 | iter_loss = dict(zip(stats_val[:, 0], stats_val[:, -1])) 46 | snap_iter = iters_avail[np.argmin([iter_loss[v] for v in iters_avail])] 47 | elif pick_rule.isdigit(): 48 | snap_iter = int(pick_rule) 49 | else: 50 | raise ValueError('Unknown snapshot rule: {}'.format(pick_rule)) 51 | weights = os.path.join(exp_dir, '{}_iter_{}.caffemodel'.format(snapshot_prefix, snap_iter)) 52 | target = os.path.join(exp_dir, '{}.caffemodel'.format(snapshot_prefix)) 53 | print('copying {} to {} ...'.format(weights, target), end='', flush=True) 54 | shutil.copyfile(weights, target) 55 | print(' done!', flush=True) 56 | 57 | 58 | SHAPE_CATEGORY_NAMES = ('laptop', 'car', 'skateboard', 'chair', 'earphone', 'motorbike', 'mug', 'airplane', 59 | 'pistol', 'lamp', 'bag', 'cap', 'rocket', 'guitar', 'table', 'knife') 60 | 61 | 62 | if __name__ == "__main__": 63 | if len(sys.argv) < 3: 64 | print('Usage: {} '.format(os.path.basename(sys.argv[0])) + ' []') 65 | else: 66 | categories = SHAPE_CATEGORY_NAMES if len(sys.argv) < 4 else (sys.argv[3],) 67 | copy_best_model(sys.argv[1], categories, sys.argv[2]) 68 | 69 | -------------------------------------------------------------------------------- /sampling-based/splatnet/plot_log.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import argparse 7 | import subprocess 8 | import numpy as np 9 | import matplotlib.pyplot as plt 10 | 11 | 12 | def parse_and_plot(path, subplot_size=5, n_col=3, skip_train=100, skip_test=10, 13 | caffe_root=None, show_wo_save=False): 14 | 15 | if not caffe_root: 16 | import caffe 17 | caffe_root = os.path.join(os.path.dirname(caffe.__file__), '..', '..') 18 | 19 | parser_path = os.path.join(caffe_root, 'tools/extra/parse_log.py') 20 | subprocess.run('python2 {} {} {}'.format(parser_path, path, os.path.dirname(path)).split(' ')) 21 | 22 | with open(path + '.train') as f: 23 | l = f.readline() 24 | labels_train = l.strip().split(',')[2:] 25 | with open(path + '.test') as f: 26 | l = f.readline() 27 | labels_test = l.strip().split(',')[2:] 28 | 29 | stats_train = np.loadtxt(path + '.train', delimiter=',', skiprows=1) 30 | stats_test = np.loadtxt(path + '.test', delimiter=',', skiprows=1) 31 | 32 | n_row = (len(labels_train) - 1) // n_col + 1, (len(labels_test) - 1) // n_col + 1 33 | 34 | if not show_wo_save: 35 | plt.switch_backend('agg') 36 | 37 | plt.figure(figsize=(n_col * subplot_size, sum(n_row) * subplot_size)) 38 | for i, l in enumerate(labels_train): 39 | plt.subplot(sum(n_row), n_col, i + 1) 40 | plt.plot(stats_train[skip_train:, 0], stats_train[skip_train:, i + 2]) 41 | plt.title('train_' + l), plt.grid('on') 42 | for i, l in enumerate(labels_test): 43 | plt.subplot(sum(n_row), n_col, n_row[0] * n_col + i + 1) 44 | plt.plot(stats_test[skip_test:, 0], stats_test[skip_test:, i + 2]) 45 | plt.title('test_' + l), plt.grid('on') 46 | 47 | if show_wo_save: 48 | plt.show() 49 | else: 50 | plt.savefig(path + '.png') 51 | 52 | return stats_train, stats_test 53 | 54 | 55 | if __name__ == '__main__': 56 | parser = argparse.ArgumentParser(description='Parse caffe training log and plot stats', 57 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 58 | parser.add_argument('path') 59 | parser.add_argument('--subplot_size', default=5, type=int, help='width and height of each subplot') 60 | parser.add_argument('--num_column', default=3, type=int, help='number of columns') 61 | parser.add_argument('--skip_train', default=100, type=int, help='skip first such training iterations') 62 | parser.add_argument('--skip_test', default=10, type=int, help='skip first such testing iterations') 63 | parser.add_argument('--show_wo_save', action='store_true', help='show figure instead of saving it') 64 | parser.add_argument('--caffe_root', default=None, type=str, help='path to caffe installation directory') 65 | 66 | args = parser.parse_args() 67 | 68 | parse_and_plot(args.path, args.subplot_size, args.num_column, args.skip_train, args.skip_test, 69 | args.caffe_root, args.show_wo_save) 70 | 71 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Spatial Transformer for 3D Point Clouds 2 | 3 | [[Project]](http://pwang.pw/spn.html) [[Paper]](https://arxiv.org/abs/1906.10887) 4 | 5 | ## Quick Start 6 | For quick addition of the spatial transformer to your network, refer to [network architecture file](point_based/part_seg/part_seg_model_deform.py#L53) of how transformer can be added, and [offset_deform](point_based/utils/tf_util.py#L120-L160) for the transformer implementation. 7 | 8 | ## Overview 9 | This is the author's re-implementation of 10 | "[Spatial Transformer for 3D Point Clouds](https://arxiv.org/abs/1906.10887)", by 11 | [Jiayun Wang](http://pwang.pw/),  [Rudrasis Chakraborty](https://rudra1988.github.io/),  [Stella X. Yu](https://www1.icsi.berkeley.edu/~stellayu/)  (UC Berkeley / ICSI)  12 | in IEEE Transactions on Pattern Analysis and Machine Intelligence. 13 | 14 | 15 | 16 | Further information please contact [Jiayun Wang](mailto:peterwg@berkeley.edu). 17 | 18 | ## Update notifications 19 | * 03/09/2019: Uploaded point-based methods for ShapeNet part segmentation. 20 | * 10/07/2019: Uploaded sampling-based methods for ShapeNet part segmentation. 21 | 22 | ## Requirements 23 | * [Tensorflow](https://www.tensorflow.org/get_started/os_setup) (for the point-based method, version >= 1.13.1) 24 | * [CAFFE](https://github.com/samaonline/caffe-deform) (for the sampling-based method, please use our version as we rewrite some source codes.) 25 | * [NCCL](https://github.com/NVIDIA/nccl) (for multi-gpu in the sampling-based method) 26 | 27 | ## Point-based Methods 28 | 29 | Please navigate to the specific folder first. 30 | 31 | ```bash 32 | cd point_based 33 | ``` 34 | 35 | ### Install Tensorflow and h5py 36 | 37 | Install TensorFlow. You may also need to install h5py. 38 | 39 | To install h5py for Python: 40 | ```bash 41 | sudo apt-get install libhdf5-dev 42 | pip install h5py 43 | ``` 44 | 45 | ### Data Preparation 46 | 47 | Download the data for part segmentation. 48 | 49 | ``` 50 | sh +x download_data.sh 51 | ``` 52 | 53 | ### Running Examples 54 | 55 | #### Train 56 | 57 | Train the deformable spatial transformer. Specify number of gpus used with '--num_gpu'. Specify number of graphs and number of feature dimensions by '--graphnum' and '--featnum', respectively. 58 | 59 | ``` 60 | cd part_seg 61 | python train_deform.py --graphnum 2 --featnum 64 62 | ``` 63 | 64 | Model parameters are saved every 10 epochs in "train_results/trained_models/". 65 | 66 | #### Evaluation 67 | 68 | To evaluate the model saved after epoch n, 69 | 70 | ``` 71 | python test.py --model_path train_results/trained_models/epoch_n.ckpt --graphnum 2 --featnum 64 72 | ``` 73 | 74 | ## Sampling-based Methods 75 | 76 | ### Install Caffe 77 | 78 | Please use [our version of CAFFE](https://github.com/samaonline/caffe-deform), as we provide the implementation of spatial transformers for bilateralNN, as described in the paper. A guide to CAFFE installation can be found [here](https://caffe.berkeleyvision.org/installation.html). 79 | 80 | ### Data Preparation 81 | 82 | Please navigate to the specific folder first. 83 | 84 | ```bash 85 | cd sampling-based 86 | ``` 87 | 88 | See instructions in [data/README.md](https://github.com/samaonline/spatial-transformer-for-3d-point-clouds/blob/master/sampling-based/data/README.md). 89 | 90 | ### Running Examples 91 | 92 | * ShapeNet Part segmentation 93 | * train and evaluate 94 | ```bash 95 | cd sampling-based/exp/shapenet3d 96 | ./train_test.sh 97 | * test trained model 98 | ```bash 99 | cd sampling-based/exp/shapenet3d 100 | ./test_only.sh 101 | ``` 102 | Predictions are under `pred/`, with evaluation results in `test.log`. 103 | 104 | ## Benchmarks and Model Zoo 105 | 106 | Please refer to Section 4 of the [paper](https://arxiv.org/abs/1906.10887). 107 | 108 | ## Additional Notes 109 | The code is implemented based on [Dynamic Graph CNN](https://github.com/WangYueFt/dgcnn), [BilateralNN](https://github.com/MPI-IS/bilateralNN) and [SplatNet](https://github.com/NVlabs/splatnet). 110 | 111 | ## License and Citation 112 | The use of this software is RESTRICTED to **non-commercial research and educational purposes**. 113 | ``` 114 | @article{spn3dpointclouds, 115 | author = {Jiayun Wang and 116 | Rudrasis Chakraborty and 117 | Stella X. Yu}, 118 | title = {Spatial Transformer for 3D Points}, 119 | journal = {CoRR}, 120 | volume = {abs/1906.10887}, 121 | year = {2019}, 122 | url = {http://arxiv.org/abs/1906.10887}, 123 | } 124 | ``` 125 | -------------------------------------------------------------------------------- /sampling-based/splatnet/semseg3d/eval_seg.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import sys 6 | import argparse 7 | import numpy as np 8 | from sklearn.metrics import confusion_matrix 9 | import splatnet.configs 10 | from splatnet.utils import TimedBlock 11 | 12 | 13 | def get_label(ply_path, column=3, from_rgb=True, cmap=None): 14 | with open(ply_path) as f: 15 | header_size = np.where([l.strip() == 'end_header' for l in f.readlines(1000)])[0][0] + 1 16 | ply_data = np.loadtxt(ply_path, skiprows=header_size) 17 | if from_rgb: 18 | return [np.where(np.prod(v == cmap, axis=1))[0][0] for v in ply_data[:, column:column+3]] 19 | else: 20 | return ply_data[:, column] - 1 21 | 22 | 23 | def compute_scores(pred, gt): 24 | pred, gt = np.array(pred), np.array(gt) 25 | scores = dict() 26 | labels = np.unique(gt) 27 | assert np.all([(v in labels) for v in np.unique(pred)]) 28 | 29 | TPs, FPs, FNs, Total = [], [], [], [] 30 | for l in labels: 31 | TPs.append(sum((gt == l) * (pred == l))) 32 | FPs.append(sum((gt != l) * (pred == l))) 33 | FNs.append(sum((gt == l) * (pred != l))) 34 | Total.append(sum(gt == l)) 35 | 36 | scores['accuracy'] = sum(gt == pred) / len(gt) 37 | scores['confusion'] = confusion_matrix(gt, pred) 38 | scores['class_accuracy'] = [TPs[i] / (TPs[i] + FNs[i]) for i in range(len(labels))] 39 | scores['avg_class_accuracy'] = sum(scores['class_accuracy']) / len(labels) 40 | scores['class_iou'] = [TPs[i] / (TPs[i] + FNs[i] + FPs[i]) for i in range(len(labels))] 41 | scores['avg_class_iou'] = sum(scores['class_iou']) / len(labels) 42 | scores['num_points'] = Total 43 | 44 | return scores 45 | 46 | 47 | if __name__ == '__main__': 48 | parser = argparse.ArgumentParser(description='Compute evaluation metrics for segmentation results', 49 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 50 | parser.add_argument('pred', type=str, help='path to predictions') 51 | parser.add_argument('gt', type=str, help='path to ground-truth') 52 | parser.add_argument('--dataset', default='facade', choices=('facade', 'stanford3d'), 53 | help='specify a dataset for the evaluation') 54 | parser.add_argument('--pred_rgb', action='store_true', help='turn this on if pred is encoded with rgb values') 55 | parser.add_argument('--gt_rgb', action='store_true', help='turn this on if gt is encoded with rgb values') 56 | parser.add_argument('--pred_column', type=int, default=4, help='(starting) column of label in pred (1-index)') 57 | parser.add_argument('--gt_column', type=int, default=4, help='(starting) column of label in gt (1-index)') 58 | parser.add_argument('--log', type=str, default=None, help='redirect output to log if specified') 59 | parser.add_argument('-q', '--quiet', action='store_true', help='silent intermediate information') 60 | args = parser.parse_args() 61 | 62 | if args.dataset == 'facade': 63 | classes = splatnet.configs.FACADE_CATEGORIES 64 | cmap = splatnet.configs.FACADE_CMAP 65 | 66 | with TimedBlock('Loading predictions from {}'.format(args.pred), not args.quiet): 67 | pred = get_label(args.pred, args.pred_column - 1, args.pred_rgb, cmap=cmap) 68 | 69 | with TimedBlock('Loading ground-truth from {}'.format(args.gt), not args.quiet): 70 | gt = get_label(args.gt, args.gt_column - 1, args.gt_rgb, cmap=cmap) 71 | 72 | elif args.dataset == 'stanford3d': 73 | raise NotImplementedError() 74 | else: 75 | raise ValueError('Dataset {} is not supported'.format(args.dataset)) 76 | 77 | with TimedBlock('Computing scores', not args.quiet): 78 | scores = compute_scores(pred, gt) 79 | 80 | if not args.quiet: 81 | print('Evaluation done!') 82 | 83 | if args.log: 84 | sys.stdout = open(args.log, 'a') 85 | 86 | print('-------------------- Summary --------------------') 87 | print(' Overall accuracy: {:.4f}'.format(scores['accuracy'])) 88 | print('Avg. class accuracy: {:.4f}'.format(scores['avg_class_accuracy'])) 89 | print(' IoU: {:.4f}'.format(scores['avg_class_iou'])) 90 | print('-------------------- Breakdown --------------------') 91 | print(' class count(ratio) accuracy IoU') 92 | total_points = sum(scores['num_points']) 93 | for i in range(len(classes)): 94 | print('{:10} {:7d}({:4.1f}%) {:.4f} {:.4f}'.format(classes[i], scores['num_points'][i], 95 | 100 * scores['num_points'][i] / total_points, 96 | scores['class_accuracy'][i], scores['class_iou'][i])) 97 | print('-------------------- Confusion --------------------') 98 | print(' {}'.format(' '.join(['{:>7}'.format(v) for v in classes]))) 99 | for i, c in enumerate(classes): 100 | print('{:7} {}'.format(c, ' '.join(['{:7d}'.format(v) for v in scores['confusion'][i]]))) 101 | 102 | -------------------------------------------------------------------------------- /point_based/utils/data_prep_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 4 | sys.path.append(BASE_DIR) 5 | from plyfile import (PlyData, PlyElement, make2d, PlyParseError, PlyProperty) 6 | import numpy as np 7 | import h5py 8 | 9 | SAMPLING_BIN = os.path.join(BASE_DIR, 'third_party/mesh_sampling/build/pcsample') 10 | 11 | SAMPLING_POINT_NUM = 2048 12 | SAMPLING_LEAF_SIZE = 0.005 13 | 14 | MODELNET40_PATH = '../datasets/modelnet40' 15 | def export_ply(pc, filename): 16 | vertex = np.zeros(pc.shape[0], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 17 | for i in range(pc.shape[0]): 18 | vertex[i] = (pc[i][0], pc[i][1], pc[i][2]) 19 | ply_out = PlyData([PlyElement.describe(vertex, 'vertex', comments=['vertices'])]) 20 | ply_out.write(filename) 21 | 22 | # Sample points on the obj shape 23 | def get_sampling_command(obj_filename, ply_filename): 24 | cmd = SAMPLING_BIN + ' ' + obj_filename 25 | cmd += ' ' + ply_filename 26 | cmd += ' -n_samples %d ' % SAMPLING_POINT_NUM 27 | cmd += ' -leaf_size %f ' % SAMPLING_LEAF_SIZE 28 | return cmd 29 | 30 | # -------------------------------------------------------------- 31 | # Following are the helper functions to load MODELNET40 shapes 32 | # -------------------------------------------------------------- 33 | 34 | # Read in the list of categories in MODELNET40 35 | def get_category_names(): 36 | shape_names_file = os.path.join(MODELNET40_PATH, 'shape_names.txt') 37 | shape_names = [line.rstrip() for line in open(shape_names_file)] 38 | return shape_names 39 | 40 | # Return all the filepaths for the shapes in MODELNET40 41 | def get_obj_filenames(): 42 | obj_filelist_file = os.path.join(MODELNET40_PATH, 'filelist.txt') 43 | obj_filenames = [os.path.join(MODELNET40_PATH, line.rstrip()) for line in open(obj_filelist_file)] 44 | print('Got %d obj files in modelnet40.' % len(obj_filenames)) 45 | return obj_filenames 46 | 47 | # Helper function to create the father folder and all subdir folders if not exist 48 | def batch_mkdir(output_folder, subdir_list): 49 | if not os.path.exists(output_folder): 50 | os.mkdir(output_folder) 51 | for subdir in subdir_list: 52 | if not os.path.exists(os.path.join(output_folder, subdir)): 53 | os.mkdir(os.path.join(output_folder, subdir)) 54 | 55 | # ---------------------------------------------------------------- 56 | # Following are the helper functions to load save/load HDF5 files 57 | # ---------------------------------------------------------------- 58 | 59 | # Write numpy array data and label to h5_filename 60 | def save_h5_data_label_normal(h5_filename, data, label, normal, 61 | data_dtype='float32', label_dtype='uint8', noral_dtype='float32'): 62 | h5_fout = h5py.File(h5_filename, 'r') 63 | h5_fout.create_dataset( 64 | 'data', data=data, 65 | compression='gzip', compression_opts=4, 66 | dtype=data_dtype) 67 | h5_fout.create_dataset( 68 | 'normal', data=normal, 69 | compression='gzip', compression_opts=4, 70 | dtype=normal_dtype) 71 | h5_fout.create_dataset( 72 | 'label', data=label, 73 | compression='gzip', compression_opts=1, 74 | dtype=label_dtype) 75 | h5_fout.close() 76 | 77 | 78 | # Write numpy array data and label to h5_filename 79 | def save_h5(h5_filename, data, label, data_dtype='uint8', label_dtype='uint8'): 80 | h5_fout = h5py.File(h5_filename, 'r') 81 | h5_fout.create_dataset( 82 | 'data', data=data, 83 | compression='gzip', compression_opts=4, 84 | dtype=data_dtype) 85 | h5_fout.create_dataset( 86 | 'label', data=label, 87 | compression='gzip', compression_opts=1, 88 | dtype=label_dtype) 89 | h5_fout.close() 90 | 91 | # Read numpy array data and label from h5_filename 92 | def load_h5_data_label_normal(h5_filename): 93 | f = h5py.File(h5_filename, 'r') 94 | data = f['data'][:] 95 | label = f['label'][:] 96 | normal = f['normal'][:] 97 | return (data, label, normal) 98 | 99 | # Read numpy array data and label from h5_filename 100 | def load_h5_data_label_seg(h5_filename): 101 | f = h5py.File(h5_filename, 'r') 102 | data = f['data'][:] 103 | label = f['label'][:] 104 | seg = f['pid'][:] 105 | return (data, label, seg) 106 | 107 | # Read numpy array data and label from h5_filename 108 | def load_h5(h5_filename): 109 | f = h5py.File(h5_filename, 'r') 110 | data = f['data'][:] 111 | label = f['label'][:] 112 | return (data, label) 113 | 114 | # ---------------------------------------------------------------- 115 | # Following are the helper functions to load save/load PLY files 116 | # ---------------------------------------------------------------- 117 | 118 | # Load PLY file 119 | def load_ply_data(filename, point_num): 120 | plydata = PlyData.read(filename) 121 | pc = plydata['vertex'].data[:point_num] 122 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 123 | return pc_array 124 | 125 | # Load PLY file 126 | def load_ply_normal(filename, point_num): 127 | plydata = PlyData.read(filename) 128 | pc = plydata['normal'].data[:point_num] 129 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 130 | return pc_array 131 | 132 | # Make up rows for Nxk array 133 | # Input Pad is 'edge' or 'constant' 134 | def pad_arr_rows(arr, row, pad='edge'): 135 | assert(len(arr.shape) == 2) 136 | assert(arr.shape[0] <= row) 137 | assert(pad == 'edge' or pad == 'constant') 138 | if arr.shape[0] == row: 139 | return arr 140 | if pad == 'edge': 141 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'edge') 142 | if pad == 'constant': 143 | return np.lib.pad(arr, ((0, row-arr.shape[0]), (0, 0)), 'constant', (0, 0)) 144 | 145 | 146 | -------------------------------------------------------------------------------- /point_based/part_seg/provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import numpy as np 4 | import h5py 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | sys.path.append(BASE_DIR) 7 | 8 | # Download dataset for point cloud classification 9 | DATA_DIR = os.path.join(BASE_DIR, 'data') 10 | if not os.path.exists(DATA_DIR): 11 | os.mkdir(DATA_DIR) 12 | if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')): 13 | www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip' 14 | zipfile = os.path.basename(www) 15 | os.system('wget %s; unzip %s' % (www, zipfile)) 16 | os.system('mv %s %s' % (zipfile[:-4], DATA_DIR)) 17 | os.system('rm %s' % (zipfile)) 18 | 19 | 20 | def shuffle_data(data, labels): 21 | """ Shuffle data and labels. 22 | Input: 23 | data: B,N,... numpy array 24 | label: B,... numpy array 25 | Return: 26 | shuffled data, label and shuffle indices 27 | """ 28 | idx = np.arange(len(labels)) 29 | np.random.shuffle(idx) 30 | return data[idx, ...], labels[idx], idx 31 | 32 | 33 | def rotate_point_cloud(batch_data): 34 | """ Randomly rotate the point clouds to augument the dataset 35 | rotation is per shape based along up direction 36 | Input: 37 | BxNx3 array, original batch of point clouds 38 | Return: 39 | BxNx3 array, rotated batch of point clouds 40 | """ 41 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 42 | for k in xrange(batch_data.shape[0]): 43 | rotation_angle = np.random.uniform() * 2 * np.pi 44 | cosval = np.cos(rotation_angle) 45 | sinval = np.sin(rotation_angle) 46 | rotation_matrix = np.array([[cosval, 0, sinval], 47 | [0, 1, 0], 48 | [-sinval, 0, cosval]]) 49 | shape_pc = batch_data[k, ...] 50 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 51 | return rotated_data 52 | 53 | 54 | def rotate_point_cloud_by_angle(batch_data, rotation_angle): 55 | """ Rotate the point cloud along up direction with certain angle. 56 | Input: 57 | BxNx3 array, original batch of point clouds 58 | Return: 59 | BxNx3 array, rotated batch of point clouds 60 | """ 61 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 62 | for k in xrange(batch_data.shape[0]): 63 | #rotation_angle = np.random.uniform() * 2 * np.pi 64 | cosval = np.cos(rotation_angle) 65 | sinval = np.sin(rotation_angle) 66 | rotation_matrix = np.array([[cosval, 0, sinval], 67 | [0, 1, 0], 68 | [-sinval, 0, cosval]]) 69 | shape_pc = batch_data[k, ...] 70 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), rotation_matrix) 71 | return rotated_data 72 | 73 | 74 | def rotate_perturbation_point_cloud(batch_data, angle_sigma=0.06, angle_clip=0.18): 75 | """ Randomly perturb the point clouds by small rotations 76 | Input: 77 | BxNx3 array, original batch of point clouds 78 | Return: 79 | BxNx3 array, rotated batch of point clouds 80 | """ 81 | rotated_data = np.zeros(batch_data.shape, dtype=np.float32) 82 | for k in xrange(batch_data.shape[0]): 83 | angles = np.clip(angle_sigma*np.random.randn(3), -angle_clip, angle_clip) 84 | Rx = np.array([[1,0,0], 85 | [0,np.cos(angles[0]),-np.sin(angles[0])], 86 | [0,np.sin(angles[0]),np.cos(angles[0])]]) 87 | Ry = np.array([[np.cos(angles[1]),0,np.sin(angles[1])], 88 | [0,1,0], 89 | [-np.sin(angles[1]),0,np.cos(angles[1])]]) 90 | Rz = np.array([[np.cos(angles[2]),-np.sin(angles[2]),0], 91 | [np.sin(angles[2]),np.cos(angles[2]),0], 92 | [0,0,1]]) 93 | R = np.dot(Rz, np.dot(Ry,Rx)) 94 | shape_pc = batch_data[k, ...] 95 | rotated_data[k, ...] = np.dot(shape_pc.reshape((-1, 3)), R) 96 | return rotated_data 97 | 98 | 99 | def jitter_point_cloud(batch_data, sigma=0.01, clip=0.05): 100 | """ Randomly jitter points. jittering is per point. 101 | Input: 102 | BxNx3 array, original batch of point clouds 103 | Return: 104 | BxNx3 array, jittered batch of point clouds 105 | """ 106 | B, N, C = batch_data.shape 107 | assert(clip > 0) 108 | jittered_data = np.clip(sigma * np.random.randn(B, N, C), -1*clip, clip) 109 | jittered_data += batch_data 110 | return jittered_data 111 | 112 | def shift_point_cloud(batch_data, shift_range=0.1): 113 | """ Randomly shift point cloud. Shift is per point cloud. 114 | Input: 115 | BxNx3 array, original batch of point clouds 116 | Return: 117 | BxNx3 array, shifted batch of point clouds 118 | """ 119 | B, N, C = batch_data.shape 120 | shifts = np.random.uniform(-shift_range, shift_range, (B,3)) 121 | for batch_index in range(B): 122 | batch_data[batch_index,:,:] += shifts[batch_index,:] 123 | return batch_data 124 | 125 | 126 | def random_scale_point_cloud(batch_data, scale_low=0.8, scale_high=1.25): 127 | """ Randomly scale the point cloud. Scale is per point cloud. 128 | Input: 129 | BxNx3 array, original batch of point clouds 130 | Return: 131 | BxNx3 array, scaled batch of point clouds 132 | """ 133 | B, N, C = batch_data.shape 134 | scales = np.random.uniform(scale_low, scale_high, B) 135 | for batch_index in range(B): 136 | batch_data[batch_index,:,:] *= scales[batch_index] 137 | return batch_data 138 | 139 | def getDataFiles(list_filename): 140 | return [line.rstrip() for line in open(list_filename)] 141 | 142 | def load_h5(h5_filename): 143 | f = h5py.File(h5_filename) 144 | data = f['data'][:] 145 | label = f['label'][:] 146 | return (data, label) 147 | 148 | def loadDataFile(filename): 149 | return load_h5(filename) 150 | 151 | 152 | def load_h5_data_label_seg(h5_filename): 153 | f = h5py.File(h5_filename) 154 | data = f['data'][:] # (2048, 2048, 3) 155 | label = f['label'][:] # (2048, 1) 156 | seg = f['pid'][:] # (2048, 2048) 157 | return (data, label, seg) -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/indoor3d_util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import glob 3 | import os 4 | import sys 5 | 6 | def sample_data(data, num_sample): 7 | """ data is in N x ... 8 | we want to keep num_samplexC of them. 9 | if N > num_sample, we will randomly keep num_sample of them. 10 | if N < num_sample, we will randomly duplicate samples. 11 | """ 12 | N = data.shape[0] 13 | if (N == num_sample): 14 | return data, range(N) 15 | elif (N > num_sample): 16 | sample = np.random.choice(N, num_sample) 17 | return data[sample, ...], sample 18 | else: 19 | sample = np.random.choice(N, num_sample-N) 20 | dup_data = data[sample, ...] 21 | return np.concatenate([data, dup_data], 0), range(N)+list(sample) 22 | 23 | def sample_data_label(data, label, num_sample): 24 | new_data, sample_indices = sample_data(data, num_sample) 25 | new_label = label[sample_indices] 26 | return new_data, new_label 27 | 28 | def room2blocks(data, label, num_point, block_size=1.0, stride=1.0, 29 | random_sample=False, sample_num=None, sample_aug=1): 30 | """ Prepare block training data. 31 | Args: 32 | data: N x 6 numpy array, 012 are XYZ in meters, 345 are RGB in [0,1] 33 | assumes the data is shifted (min point is origin) and aligned 34 | (aligned with XYZ axis) 35 | label: N size uint8 numpy array from 0-12 36 | num_point: int, how many points to sample in each block 37 | block_size: float, physical size of the block in meters 38 | stride: float, stride for block sweeping 39 | random_sample: bool, if True, we will randomly sample blocks in the room 40 | sample_num: int, if random sample, how many blocks to sample 41 | [default: room area] 42 | sample_aug: if random sample, how much aug 43 | Returns: 44 | block_datas: K x num_point x 6 np array of XYZRGB, RGB is in [0,1] 45 | block_labels: K x num_point x 1 np array of uint8 labels 46 | 47 | TODO: for this version, blocking is in fixed, non-overlapping pattern. 48 | """ 49 | #assert(stride<=block_size) 50 | 51 | limit = np.amax(data, 0)[0:3] 52 | 53 | # Get the corner location for our sampling blocks 54 | xbeg_list = [] 55 | ybeg_list = [] 56 | if not random_sample: 57 | num_block_x = int(np.ceil((limit[0] - block_size) / stride)) + 1 58 | num_block_y = int(np.ceil((limit[1] - block_size) / stride)) + 1 59 | for i in range(num_block_x): 60 | for j in range(num_block_y): 61 | xbeg_list.append(i*stride) 62 | ybeg_list.append(j*stride) 63 | else: 64 | num_block_x = int(np.ceil(limit[0] / block_size)) 65 | num_block_y = int(np.ceil(limit[1] / block_size)) 66 | if sample_num is None: 67 | sample_num = num_block_x * num_block_y * sample_aug 68 | for _ in range(sample_num): 69 | xbeg = np.random.uniform(-block_size, limit[0]) 70 | ybeg = np.random.uniform(-block_size, limit[1]) 71 | xbeg_list.append(xbeg) 72 | ybeg_list.append(ybeg) 73 | 74 | # Collect blocks 75 | block_data_list = [] 76 | block_label_list = [] 77 | idx = 0 78 | for idx in range(len(xbeg_list)): 79 | xbeg = xbeg_list[idx] 80 | ybeg = ybeg_list[idx] 81 | xcond = (data[:,0]<=xbeg+block_size) & (data[:,0]>=xbeg) 82 | ycond = (data[:,1]<=ybeg+block_size) & (data[:,1]>=ybeg) 83 | cond = xcond & ycond 84 | if np.sum(cond) < 100: # discard block if there are less than 100 pts. 85 | continue 86 | 87 | block_data = data[cond, :] 88 | block_label = label[cond] 89 | 90 | # randomly subsample data 91 | block_data_sampled, block_label_sampled = \ 92 | sample_data_label(block_data, block_label, num_point) 93 | block_data_list.append(np.expand_dims(block_data_sampled, 0)) 94 | block_label_list.append(np.expand_dims(block_label_sampled, 0)) 95 | 96 | return np.concatenate(block_data_list, 0), \ 97 | np.concatenate(block_label_list, 0) 98 | 99 | def room2blocks_wrapper_normalized(data, label, num_point, block_size=1.0, stride=1.0, 100 | random_sample=False, sample_num=None, sample_aug=1): 101 | data_all = np.zeros((0, num_point, 9)) 102 | label_all = np.zeros((0, num_point, )) 103 | for data_now, label_now in zip(data, label): 104 | data_temp, label_temp = room2blocks_plus_normalized(data_now, label_now , num_point, block_size, stride, 105 | random_sample, sample_num, sample_aug) 106 | data_all = np.concatenate((data_all, data_temp), axis = 0) 107 | label_all = np.concatenate((label_all, label_temp), axis = 0) 108 | return data_all, label_all 109 | 110 | def room2blocks_plus_normalized(data, label, num_point, block_size, stride, 111 | random_sample, sample_num, sample_aug): 112 | """ room2block, with input filename and RGB preprocessing. 113 | for each block centralize XYZ, add normalized XYZ as 678 channels 114 | """ 115 | data = data[:,0:6] 116 | data[:,3:6] /= 255.0 117 | label = label.astype(np.uint8) 118 | max_room_x = max(data[:,0]) 119 | max_room_y = max(data[:,1]) 120 | max_room_z = max(data[:,2]) 121 | 122 | data_batch, label_batch = room2blocks(data, label, num_point, block_size, stride, 123 | random_sample, sample_num, sample_aug) 124 | new_data_batch = np.zeros((data_batch.shape[0], num_point, 9)) 125 | for b in range(data_batch.shape[0]): 126 | new_data_batch[b, :, 6] = data_batch[b, :, 0]/max_room_x 127 | new_data_batch[b, :, 7] = data_batch[b, :, 1]/max_room_y 128 | new_data_batch[b, :, 8] = data_batch[b, :, 2]/max_room_z 129 | minx = min(data_batch[b, :, 0]) 130 | miny = min(data_batch[b, :, 1]) 131 | data_batch[b, :, 0] = 2* (data_batch[b, :, 0] - (minx+block_size/2) ) 132 | data_batch[b, :, 1] = 2* (data_batch[b, :, 1] - (miny+block_size/2) ) 133 | 134 | new_data_batch[:, :, 0:6] = data_batch 135 | return new_data_batch, label_batch -------------------------------------------------------------------------------- /sampling-based/splatnet/utils.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import time 6 | import tempfile 7 | import numpy as np 8 | from numpy import sin, cos 9 | 10 | 11 | def get_prototxt(net_proto, save_path=None): 12 | if save_path: 13 | f = open(save_path, mode='w+') 14 | else: 15 | f = tempfile.NamedTemporaryFile(mode='w+', delete=False) 16 | f.write(str(net_proto)) 17 | f.close() 18 | 19 | return f.name 20 | 21 | 22 | def parse_channel_scale(channel_scale_str, channel_str=False, scale_str=False): 23 | channels, scales = [], [] 24 | for f in channel_scale_str.split('_'): 25 | if f.find('*') >= 0: 26 | channels.append(f[:f.find('*')] if channel_str else int(f[:f.find('*')])) 27 | scales.append(f[f.find('*'):] if scale_str else float(f[f.find('*') + 1:])) 28 | elif f.find('/') >= 0: 29 | channels.append(f[:f.find('/')] if channel_str else int(f[:f.find('/')])) 30 | scales.append(f[f.find('/'):] if scale_str else (1.0 / float(f[f.find('/') + 1:]))) 31 | else: 32 | channels.append(f if channel_str else int(f)) 33 | scales.append('' if scale_str else 1.0) 34 | return channels, scales 35 | 36 | 37 | def map_channel_scale(channel_scale_str, refs): 38 | channels, scales = parse_channel_scale(channel_scale_str, channel_str=True, scale_str=True) 39 | return '_'.join([str(f) + s for f, s in zip([np.where([i == v for v in refs])[0][0] for i in channels], scales)]) 40 | 41 | 42 | def rotate_3d(xyz, rotations, center=(0, 0, 0)): 43 | """ 44 | Apply rotations to 3d points. 45 | :param xyz: N x 3 ndarray 46 | :param rotations: a list of rotations, each as (axis, angle) 47 | :param center: optionally, rotate around a non-origin center 48 | :return: rotated N x 3 ndarray (a copy -- original xyz is untouched) 49 | """ 50 | rot = np.eye(3) 51 | xyz = xyz.copy() 52 | 53 | # translate center to origin 54 | if np.any(center): 55 | xyz -= center 56 | 57 | # rotation matrix 58 | for axis, theta in rotations: 59 | if axis in {'x', 'X'}: 60 | rot_axis = np.array([[1, 0, 0], 61 | [0, cos(theta), -sin(theta)], 62 | [0, sin(theta), cos(theta)]]) 63 | elif axis in {'y', 'Y'}: 64 | rot_axis = np.array([[cos(theta), 0, sin(theta)], 65 | [0, 1, 0], 66 | [-sin(theta), 0, cos(theta)]]) 67 | elif axis in {'z', 'Z'}: 68 | rot_axis = np.array([[cos(theta), -sin(theta), 0], 69 | [sin(theta), cos(theta), 0], 70 | [0, 0, 1]]) 71 | else: 72 | raise ValueError('Unknown axis: ' + axis) 73 | rot = rot.dot(rot_axis) 74 | 75 | # apply rotation to row vectors 76 | if len(rotations) > 0: 77 | xyz = xyz.dot(rot.T) 78 | 79 | # translate back to origin 80 | if np.any(center): 81 | xyz += center 82 | 83 | return xyz 84 | 85 | 86 | def modify_blob_shape(net_path, tops, blob_size): 87 | """ 88 | TODO: make this more generic 89 | Modify .prototxt network specs 90 | :return: path to a temporary network definition with modified blob size 91 | """ 92 | from caffe.proto import caffe_pb2 93 | import google.protobuf.text_format as txtf 94 | import tempfile 95 | 96 | net = caffe_pb2.NetParameter() 97 | 98 | with open(net_path) as f: 99 | txtf.Merge(f.read(), net) 100 | 101 | for top in tops: 102 | for i in blob_size: 103 | net.layer[[l.name for l in net.layer].index(top)].input_param.shape[0].dim[i] = blob_size[i] 104 | 105 | f = tempfile.NamedTemporaryFile(mode='w+', delete=False) 106 | f.write(str(net)) 107 | f.close() 108 | return f.name 109 | 110 | 111 | class TimedBlock: 112 | """ 113 | Context manager that times the execution of a block of code. 114 | """ 115 | def __init__(self, msg='', verbose=False): 116 | self.msg = msg 117 | self.verbose = verbose 118 | 119 | def __enter__(self): 120 | if self.verbose and self.msg: 121 | #print('{} ...'.format(self.msg), end='', flush=True) 122 | print(str(self.msg)) 123 | self.tic = time.time() 124 | 125 | def __exit__(self, exc_type, exc_val, exc_tb): 126 | toc = time.time() 127 | if self.verbose and self.msg: 128 | print(' done! ' + str(toc - self.tic)) 129 | 130 | 131 | def seg_scores(pred_list, gt_list, nclasses=-1): 132 | 133 | eps = 0.0001 134 | 135 | labels = np.unique(np.concatenate([np.unique(gt) for gt in gt_list])) 136 | if nclasses != -1: 137 | assert len(labels) == nclasses 138 | 139 | acc, avg_class_acc, avg_class_iou = [], [], [] 140 | perclass_iou = np.zeros((0, len(labels) )) 141 | for pred, gt in zip(pred_list, gt_list): 142 | pred, gt = np.array(pred), np.array(gt) 143 | assert np.all([(v in labels) for v in np.unique(pred)]) 144 | 145 | tp, fp, fn = [], [], [] 146 | for l in labels: 147 | tp.append(sum((gt == l) * (pred == l))) 148 | fp.append(sum((gt != l) * (pred == l))) 149 | fn.append(sum((gt == l) * (pred != l))) 150 | 151 | class_acc = [(tp[i] + eps) / (tp[i] + fn[i] + eps) for i in range(len(labels))] 152 | class_iou = [(tp[i] + eps) / (tp[i] + fn[i] + fp[i] + eps) for i in range(len(labels))] 153 | perclass_iou = np.concatenate( (perclass_iou, np.asarray(class_iou).reshape((1, len(labels) )) ), axis =0) 154 | 155 | acc.append(sum(gt == pred)*1.0 / len(gt)) 156 | avg_class_acc.append(sum(class_acc) / len(labels)) 157 | avg_class_iou.append(sum(class_iou) / len(labels)) 158 | perclass_iou = np.mean( perclass_iou,axis =0 ) 159 | np.save("perclass_iou", perclass_iou) 160 | for i in perclass_iou: 161 | print(i) 162 | return acc, avg_class_acc, avg_class_iou 163 | 164 | def pairwise_distance(point_cloud): 165 | 166 | og_batch_size = point_cloud.shape[0] 167 | if og_batch_size == 1: 168 | point_cloud = np.expand_dims(point_cloud, 0) 169 | 170 | point_cloud_transpose = np.transpose(point_cloud, axes=[0, 2, 1]) 171 | point_cloud_inner = np.matmul(point_cloud, point_cloud_transpose) 172 | point_cloud_inner = -2*point_cloud_inner 173 | point_cloud_square = np.sum(np.square(point_cloud), axis=-1, keepdims=True) 174 | point_cloud_square_tranpose = np.transpose(point_cloud_square, axes=[0, 2, 1]) 175 | return point_cloud_square + point_cloud_inner + point_cloud_square_tranpose 176 | 177 | def knn(adj_matrix, k=20): 178 | 179 | """neg_adj = -adj_matrix 180 | _, nn_idx = tf.nn.top_k(neg_adj, k=k)""" 181 | nn_idx = np.argpartition(adj_matrix, (1, k+1) )[...,1:k+1] 182 | return nn_idx -------------------------------------------------------------------------------- /sampling-based/splatnet/semseg3d/train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import argparse 7 | import caffe 8 | import splatnet.semseg3d.models as models 9 | from splatnet import create_solver 10 | import splatnet.configs 11 | 12 | 13 | def semseg_train(network, exp_dir, exp_prefix, args): 14 | 15 | if args.cpu: 16 | caffe.set_mode_cpu() 17 | else: 18 | caffe.set_mode_gpu() 19 | caffe.set_device(0) 20 | 21 | if exp_prefix: 22 | snapshot_prefix = os.path.join(exp_dir, exp_prefix) 23 | exp_prefix = exp_prefix + '_' 24 | else: 25 | snapshot_prefix = os.path.join(exp_dir, 'snapshot') 26 | exp_prefix = '' 27 | 28 | if network == 'seq': 29 | batch_norm = True 30 | conv_weight_filler = 'xavier' 31 | network = models.semseg_seq(arch_str=args.arch, 32 | skip_str=args.skips, 33 | dataset=args.dataset, 34 | dataset_params=args.dataset_params, 35 | feat_dims_str=args.feat, 36 | lattice_dims_str=args.lattice, 37 | sample_size=args.sample_size, 38 | batch_size=args.batch_size, 39 | batchnorm=batch_norm, 40 | conv_weight_filler=conv_weight_filler, 41 | save_path=os.path.join(exp_dir, exp_prefix + 'net.prototxt')) 42 | 43 | models.semseg_seq(deploy=True, 44 | arch_str=args.arch, 45 | skip_str=args.skips, 46 | dataset=args.dataset, 47 | dataset_params=args.dataset_params, 48 | feat_dims_str=args.feat, 49 | lattice_dims_str=args.lattice, 50 | sample_size=args.sample_size, 51 | batchnorm=batch_norm, 52 | save_path=os.path.join(exp_dir, exp_prefix + 'net_deploy.prototxt')) 53 | else: 54 | assert network.endswith('.prototxt'), 'Please provide a valid prototxt file' 55 | print('Using network defined at {}'.format(network)) 56 | 57 | random_seed = 0 58 | debug_info = False 59 | solver = create_solver.standard_solver(network, 60 | network, 61 | snapshot_prefix, 62 | base_lr=args.base_lr, 63 | gamma=args.lr_decay, 64 | stepsize=args.stepsize, 65 | test_iter=args.test_iter, 66 | test_interval=args.test_interval, 67 | max_iter=args.num_iter, 68 | snapshot=args.snapshot_interval, 69 | solver_type=args.solver_type, 70 | weight_decay=args.weight_decay, 71 | iter_size=args.iter_size, 72 | debug_info=debug_info, 73 | random_seed=random_seed, 74 | save_path=os.path.join(exp_dir, exp_prefix + 'solver.prototxt')) 75 | solver = caffe.get_solver(solver) 76 | 77 | if args.init_model: 78 | if args.init_model.endswith('.solverstate'): 79 | solver.restore(args.init_model) 80 | elif args.init_model.endswith('.caffemodel'): 81 | solver.net.copy_from(args.init_model) 82 | else: 83 | raise ValueError('Invalid file: {}'.format(args.init_model)) 84 | 85 | solver.solve() 86 | 87 | 88 | if __name__ == '__main__': 89 | parser = argparse.ArgumentParser(add_help=False) 90 | 91 | group = parser.add_argument_group('network options') 92 | group.add_argument('--network', default='seq', help='network type (\'seq\') or path to a .prototxt file') 93 | group.add_argument('--arch', default='64_128_256_256', help='network architecture') 94 | group.add_argument('--dataset', default='facade', choices=('facade', 'stanford3d', 'stanford3dh5', 'shapenet'), help='dataset') 95 | group.add_argument('--dataset_params', nargs='+', help='dataset-specific parameters (key value pairs)') 96 | group.add_argument('--category', type=str, default=None, help='for shapenet tasks') 97 | group.add_argument('--skips', nargs='+', help='skip connections') 98 | group.add_argument('--sample_size', type=int, default=10000, help='number of points in a sample') 99 | group.add_argument('--batch_size', type=int, default=10, help='number of samples in a batch') 100 | group.add_argument('--feat', default='nx_ny_nz_r_g_b_h', help='features to use as input') 101 | group.add_argument('--lattice', nargs='+', help='bnn lattice features and scales') 102 | 103 | group = parser.add_argument_group('solver options') 104 | group.add_argument('--base_lr', type=float, default=0.01, help='starting learning rate') 105 | group.add_argument('--lr_decay', type=float, default=0.5, help='learning rate decay rate') 106 | group.add_argument('--stepsize', type=int, default=10000, help='learning rate decay interval') 107 | group.add_argument('--num_iter', type=int, default=10000, help='number of iterations to train') 108 | group.add_argument('--iter_size', type=int, default=1, help='number of mini-batches per iteration') 109 | group.add_argument('--test_iter', type=int, default=10, help='number of iterations to use at each testing phase') 110 | group.add_argument('--test_interval', type=int, default=100, help='test every such iterations') 111 | group.add_argument('--snapshot_interval', type=int, default=2000, help='snapshot every such iterations') 112 | group.add_argument('--solver_type', default='ADAM', help='optimizer type') 113 | group.add_argument('--weight_decay', type=float, default=0.001, help='weight decay') 114 | 115 | parser = argparse.ArgumentParser(description='Semantic segmentation training', 116 | parents=[parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter) 117 | parser.add_argument('exp_dir') 118 | parser.add_argument('--exp_prefix', default='', help='optional prefix for the experiment') 119 | parser.add_argument('--cpu', action='store_true', help='use cpu') 120 | parser.add_argument('--init_model', default=None, help='a .caffemodel/.solverstate file to start/resume training') 121 | 122 | args = parser.parse_args() 123 | 124 | if not args.dataset_params: 125 | args.dataset_params = {} 126 | else: 127 | args.dataset_params = dict(zip(args.dataset_params[::2], args.dataset_params[1::2])) 128 | 129 | os.makedirs(args.exp_dir, exist_ok=True) 130 | 131 | semseg_train(args.network, args.exp_dir, args.exp_prefix, args) 132 | 133 | -------------------------------------------------------------------------------- /point_based/utils/pc_util.py: -------------------------------------------------------------------------------- 1 | """ Utility functions for processing point clouds. 2 | 3 | Author: Charles R. Qi, Hao Su 4 | Date: November 2016 5 | """ 6 | 7 | import os 8 | import sys 9 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 10 | sys.path.append(BASE_DIR) 11 | 12 | # Draw point cloud 13 | from eulerangles import euler2mat 14 | 15 | # Point cloud IO 16 | import numpy as np 17 | from plyfile import PlyData, PlyElement 18 | 19 | 20 | # ---------------------------------------- 21 | # Point Cloud/Volume Conversions 22 | # ---------------------------------------- 23 | 24 | def point_cloud_to_volume_batch(point_clouds, vsize=12, radius=1.0, flatten=True): 25 | """ Input is BxNx3 batch of point cloud 26 | Output is Bx(vsize^3) 27 | """ 28 | vol_list = [] 29 | for b in range(point_clouds.shape[0]): 30 | vol = point_cloud_to_volume(np.squeeze(point_clouds[b,:,:]), vsize, radius) 31 | if flatten: 32 | vol_list.append(vol.flatten()) 33 | else: 34 | vol_list.append(np.expand_dims(np.expand_dims(vol, -1), 0)) 35 | if flatten: 36 | return np.vstack(vol_list) 37 | else: 38 | return np.concatenate(vol_list, 0) 39 | 40 | 41 | def point_cloud_to_volume(points, vsize, radius=1.0): 42 | """ input is Nx3 points. 43 | output is vsize*vsize*vsize 44 | assumes points are in range [-radius, radius] 45 | """ 46 | vol = np.zeros((vsize,vsize,vsize)) 47 | voxel = 2*radius/float(vsize) 48 | locations = (points + radius)/voxel 49 | locations = locations.astype(int) 50 | vol[locations[:,0],locations[:,1],locations[:,2]] = 1.0 51 | return vol 52 | 53 | #a = np.zeros((16,1024,3)) 54 | #print point_cloud_to_volume_batch(a, 12, 1.0, False).shape 55 | 56 | def volume_to_point_cloud(vol): 57 | """ vol is occupancy grid (value = 0 or 1) of size vsize*vsize*vsize 58 | return Nx3 numpy array. 59 | """ 60 | vsize = vol.shape[0] 61 | assert(vol.shape[1] == vsize and vol.shape[1] == vsize) 62 | points = [] 63 | for a in range(vsize): 64 | for b in range(vsize): 65 | for c in range(vsize): 66 | if vol[a,b,c] == 1: 67 | points.append(np.array([a,b,c])) 68 | if len(points) == 0: 69 | return np.zeros((0,3)) 70 | points = np.vstack(points) 71 | return points 72 | 73 | # ---------------------------------------- 74 | # Point cloud IO 75 | # ---------------------------------------- 76 | 77 | def read_ply(filename): 78 | """ read XYZ point cloud from filename PLY file """ 79 | plydata = PlyData.read(filename) 80 | pc = plydata['vertex'].data 81 | pc_array = np.array([[x, y, z] for x,y,z in pc]) 82 | return pc_array 83 | 84 | 85 | def write_ply(points, filename, text=True): 86 | """ input: Nx3, write points to filename as PLY format. """ 87 | points = [(points[i,0], points[i,1], points[i,2]) for i in range(points.shape[0])] 88 | vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4')]) 89 | el = PlyElement.describe(vertex, 'vertex', comments=['vertices']) 90 | PlyData([el], text=text).write(filename) 91 | 92 | 93 | # ---------------------------------------- 94 | # Simple Point cloud and Volume Renderers 95 | # ---------------------------------------- 96 | 97 | def draw_point_cloud(input_points, canvasSize=500, space=200, diameter=25, 98 | xrot=0, yrot=0, zrot=0, switch_xyz=[0,1,2], normalize=True): 99 | """ Render point cloud to image with alpha channel. 100 | Input: 101 | points: Nx3 numpy array (+y is up direction) 102 | Output: 103 | gray image as numpy array of size canvasSizexcanvasSize 104 | """ 105 | image = np.zeros((canvasSize, canvasSize)) 106 | if input_points is None or input_points.shape[0] == 0: 107 | return image 108 | 109 | points = input_points[:, switch_xyz] 110 | M = euler2mat(zrot, yrot, xrot) 111 | points = (np.dot(M, points.transpose())).transpose() 112 | 113 | # Normalize the point cloud 114 | # We normalize scale to fit points in a unit sphere 115 | if normalize: 116 | centroid = np.mean(points, axis=0) 117 | points -= centroid 118 | furthest_distance = np.max(np.sqrt(np.sum(abs(points)**2,axis=-1))) 119 | points /= furthest_distance 120 | 121 | # Pre-compute the Gaussian disk 122 | radius = (diameter-1)/2.0 123 | disk = np.zeros((diameter, diameter)) 124 | for i in range(diameter): 125 | for j in range(diameter): 126 | if (i - radius) * (i-radius) + (j-radius) * (j-radius) <= radius * radius: 127 | disk[i, j] = np.exp((-(i-radius)**2 - (j-radius)**2)/(radius**2)) 128 | mask = np.argwhere(disk > 0) 129 | dx = mask[:, 0] 130 | dy = mask[:, 1] 131 | dv = disk[disk > 0] 132 | 133 | # Order points by z-buffer 134 | zorder = np.argsort(points[:, 2]) 135 | points = points[zorder, :] 136 | points[:, 2] = (points[:, 2] - np.min(points[:, 2])) / (np.max(points[:, 2] - np.min(points[:, 2]))) 137 | max_depth = np.max(points[:, 2]) 138 | 139 | for i in range(points.shape[0]): 140 | j = points.shape[0] - i - 1 141 | x = points[j, 0] 142 | y = points[j, 1] 143 | xc = canvasSize/2 + (x*space) 144 | yc = canvasSize/2 + (y*space) 145 | xc = int(np.round(xc)) 146 | yc = int(np.round(yc)) 147 | 148 | px = dx + xc 149 | py = dy + yc 150 | 151 | image[px, py] = image[px, py] * 0.7 + dv * (max_depth - points[j, 2]) * 0.3 152 | 153 | image = image / np.max(image) 154 | return image 155 | 156 | def point_cloud_three_views(points): 157 | """ input points Nx3 numpy array (+y is up direction). 158 | return an numpy array gray image of size 500x1500. """ 159 | # +y is up direction 160 | # xrot is azimuth 161 | # yrot is in-plane 162 | # zrot is elevation 163 | img1 = draw_point_cloud(points, zrot=110/180.0*np.pi, xrot=45/180.0*np.pi, yrot=0/180.0*np.pi) 164 | img2 = draw_point_cloud(points, zrot=70/180.0*np.pi, xrot=135/180.0*np.pi, yrot=0/180.0*np.pi) 165 | img3 = draw_point_cloud(points, zrot=180.0/180.0*np.pi, xrot=90/180.0*np.pi, yrot=0/180.0*np.pi) 166 | image_large = np.concatenate([img1, img2, img3], 1) 167 | return image_large 168 | 169 | 170 | from PIL import Image 171 | def point_cloud_three_views_demo(): 172 | """ Demo for draw_point_cloud function """ 173 | points = read_ply('../third_party/mesh_sampling/piano.ply') 174 | im_array = point_cloud_three_views(points) 175 | img = Image.fromarray(np.uint8(im_array*255.0)) 176 | img.save('piano.jpg') 177 | 178 | if __name__=="__main__": 179 | point_cloud_three_views_demo() 180 | 181 | 182 | import matplotlib.pyplot as plt 183 | def pyplot_draw_point_cloud(points, output_filename): 184 | """ points is a Nx3 numpy array """ 185 | fig = plt.figure() 186 | ax = fig.add_subplot(111, projection='3d') 187 | ax.scatter(points[:,0], points[:,1], points[:,2]) 188 | ax.set_xlabel('x') 189 | ax.set_ylabel('y') 190 | ax.set_zlabel('z') 191 | #savefig(output_filename) 192 | 193 | def pyplot_draw_volume(vol, output_filename): 194 | """ vol is of size vsize*vsize*vsize 195 | output an image to output_filename 196 | """ 197 | points = volume_to_point_cloud(vol) 198 | pyplot_draw_point_cloud(points, output_filename) 199 | -------------------------------------------------------------------------------- /point_based/part_seg/part_seg_model_deform.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import numpy as np 3 | import math 4 | import os 5 | import sys 6 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 7 | sys.path.append(os.path.dirname(BASE_DIR)) 8 | sys.path.append(os.path.join(BASE_DIR, '../utils')) 9 | sys.path.append(os.path.join(BASE_DIR, '../models')) 10 | sys.path.append(os.path.join(BASE_DIR, '../')) 11 | import tf_util 12 | from pdb import set_trace as st 13 | from transform_nets import input_transform_net 14 | 15 | def get_model(point_cloud, input_label, is_training, cat_num, part_num, \ 16 | batch_size, num_point, weight_decay, graphnum, featnum, bn_decay=None): 17 | batch_size = point_cloud.get_shape()[0].value 18 | num_point = point_cloud.get_shape()[1].value 19 | input_image = tf.expand_dims(point_cloud, -1) 20 | 21 | k = 30 22 | 23 | adj = tf_util.pairwise_distance(point_cloud) 24 | nn_idx = tf_util.knn(adj, k=k) 25 | edge_feature = tf_util.get_edge_feature(input_image, nn_idx=nn_idx, k=k) 26 | 27 | with tf.variable_scope('transform_net1') as sc: 28 | transform = input_transform_net(edge_feature, is_training, bn_decay, K=3, is_dist=True) 29 | point_cloud_transformed = tf.matmul(point_cloud, transform) 30 | input_image = tf.expand_dims(point_cloud_transformed, -2) 31 | adj = tf_util.pairwise_distance(point_cloud_transformed) 32 | nn_idx = tf_util.knn(adj, k=k) 33 | edge_feature = tf_util.get_edge_feature(input_image, nn_idx=nn_idx, k=k) 34 | 35 | out1 = tf_util.conv2d(edge_feature, 64, [1,1], 36 | padding='VALID', stride=[1,1], 37 | bn=True, is_training=is_training, weight_decay=weight_decay, 38 | scope='adj_conv1', bn_decay=bn_decay, is_dist=True) 39 | 40 | out2 = tf_util.conv2d(out1, 64, [1,1], 41 | padding='VALID', stride=[1,1], 42 | bn=True, is_training=is_training, weight_decay=weight_decay, 43 | scope='adj_conv2', bn_decay=bn_decay, is_dist=True) 44 | 45 | net_max_1 = tf.reduce_max(out2, axis=-2, keep_dims=True) 46 | net_mean_1 = tf.reduce_mean(out2, axis=-2, keep_dims=True) 47 | 48 | out3 = tf_util.conv2d(tf.concat([net_max_1, net_mean_1], axis=-1), 64, [1,1], 49 | padding='VALID', stride=[1,1], 50 | bn=True, is_training=is_training, weight_decay=weight_decay, 51 | scope='adj_conv3', bn_decay=bn_decay, is_dist=True) 52 | 53 | out3A, net_max_1A, net_mean_1A = tf_util.offset_deform(input_image, out3, scope="trans_conv0", num_neighbor = k, num_graph=graphnum[0], num_feat=featnum[0], weight_decay=weight_decay, is_training=is_training, bn_decay=bn_decay) 54 | 55 | out5, net_max_2, net_mean_2 = tf_util.offset_deform(input_image, out3A, scope="trans_conv1", num_neighbor = k, num_graph=graphnum[0], num_feat=featnum[0], weight_decay=weight_decay, is_training=is_training, bn_decay=bn_decay) 56 | 57 | out7, net_max_3, net_mean_3 = tf_util.offset_deform(input_image, out5, scope="trans_conv2", num_neighbor = k, num_graph=graphnum[1], num_feat=featnum[1], weight_decay=weight_decay, is_training=is_training, bn_decay=bn_decay) 58 | '''adj = tf_util.pairwise_distance(tf.squeeze(trans2, axis=-2)) 59 | nn_idx = tf_util.knn(adj, k=k) 60 | edge_feature = tf_util.get_edge_feature(tf.concat([out5,trans2], axis = -1), nn_idx=nn_idx, k=k) 61 | 62 | out6 = tf_util.conv2d(edge_feature, 64, [1,1], 63 | padding='VALID', stride=[1,1], 64 | bn=True, is_training=is_training, weight_decay=weight_decay, 65 | scope='adj_conv6', bn_decay=bn_decay, is_dist=True) 66 | 67 | net_max_3 = tf.reduce_max(out6, axis=-2, keep_dims=True) 68 | net_mean_3 = tf.reduce_mean(out6, axis=-2, keep_dims=True) 69 | 70 | out7 = tf_util.conv2d(tf.concat([net_max_3, net_mean_3], axis=-1), 64, [1,1], 71 | padding='VALID', stride=[1,1], 72 | bn=True, is_training=is_training, weight_decay=weight_decay, 73 | scope='adj_conv7', bn_decay=bn_decay, is_dist=True)''' 74 | 75 | out8 = tf_util.conv2d(tf.concat([out3, out5, out7], axis=-1), 1024, [1, 1], 76 | padding='VALID', stride=[1,1], 77 | bn=True, is_training=is_training, 78 | scope='adj_conv13', bn_decay=bn_decay, is_dist=True) 79 | 80 | out_max = tf_util.max_pool2d(out8, [num_point, 1], padding='VALID', scope='maxpool') 81 | 82 | one_hot_label_expand = tf.reshape(input_label, [batch_size, 1, 1, cat_num]) 83 | one_hot_label_expand = tf_util.conv2d(one_hot_label_expand, 128, [1, 1], 84 | padding='VALID', stride=[1,1], 85 | bn=True, is_training=is_training, 86 | scope='one_hot_label_expand', bn_decay=bn_decay, is_dist=True) 87 | out_max = tf.concat(axis=3, values=[out_max, one_hot_label_expand]) 88 | expand = tf.tile(out_max, [1, num_point, 1, 1]) 89 | 90 | concat = tf.concat(axis=3, values=[expand, 91 | net_max_1, 92 | net_mean_1, 93 | out3, 94 | net_max_2, 95 | net_mean_2, 96 | out5, 97 | net_max_3, 98 | net_mean_3, 99 | out7, 100 | out8]) 101 | 102 | net2 = tf_util.conv2d(concat, 256, [1,1], padding='VALID', stride=[1,1], bn_decay=bn_decay, 103 | bn=True, is_training=is_training, scope='seg/conv1', weight_decay=weight_decay, is_dist=True) 104 | net2 = tf_util.dropout(net2, keep_prob=0.6, is_training=is_training, scope='seg/dp1') 105 | net2 = tf_util.conv2d(net2, 256, [1,1], padding='VALID', stride=[1,1], bn_decay=bn_decay, 106 | bn=True, is_training=is_training, scope='seg/conv2', weight_decay=weight_decay, is_dist=True) 107 | net2 = tf_util.dropout(net2, keep_prob=0.6, is_training=is_training, scope='seg/dp2') 108 | net2 = tf_util.conv2d(net2, 128, [1,1], padding='VALID', stride=[1,1], bn_decay=bn_decay, 109 | bn=True, is_training=is_training, scope='seg/conv3', weight_decay=weight_decay, is_dist=True) 110 | net2 = tf_util.conv2d(net2, part_num, [1,1], padding='VALID', stride=[1,1], activation_fn=None, 111 | bn=False, scope='seg/conv4', weight_decay=weight_decay, is_dist=True) 112 | 113 | net2 = tf.reshape(net2, [batch_size, num_point, part_num]) 114 | 115 | return net2 116 | 117 | 118 | def get_loss(seg_pred, seg): 119 | per_instance_seg_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=seg_pred, labels=seg), axis=1) 120 | seg_loss = tf.reduce_mean(per_instance_seg_loss) 121 | per_instance_seg_pred_res = tf.argmax(seg_pred, 2) 122 | 123 | return seg_loss, per_instance_seg_loss, per_instance_seg_pred_res 124 | -------------------------------------------------------------------------------- /sampling-based/splatnet/custom_layers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import numpy as np 6 | import caffe 7 | import utils 8 | class GlobalPooling(caffe.Layer): 9 | def setup(self, bottom, top): 10 | pass 11 | 12 | def reshape(self, bottom, top): 13 | top[0].reshape(*bottom[0].data.shape) 14 | 15 | def forward(self, bottom, top): 16 | n, c, h, w = bottom[0].data.shape 17 | self.max_loc = bottom[0].data.reshape(n, c, h*w).argmax(axis=2) 18 | top[0].data[...] = bottom[0].data.max(axis=(2, 3), keepdims=True) 19 | 20 | def backward(self, top, propagate_down, bottom): 21 | n, c, h, w = top[0].diff.shape 22 | nn, cc = np.ix_(np.arange(n), np.arange(c)) 23 | bottom[0].diff[...] = 0 24 | bottom[0].diff.reshape(n, c, -1)[nn, cc, self.max_loc] = top[0].diff.sum(axis=(2, 3)) 25 | 26 | 27 | class ProbRenorm(caffe.Layer): 28 | def setup(self, bottom, top): 29 | pass 30 | 31 | def reshape(self, bottom, top): 32 | top[0].reshape(*bottom[0].data.shape) 33 | 34 | def forward(self, bottom, top): 35 | clipped = bottom[0].data * bottom[1].data 36 | self.sc = 1.0 / (np.sum(clipped, axis=1, keepdims=True) + 1e-10) 37 | top[0].data[...] = clipped * self.sc 38 | 39 | def backward(self, top, propagate_down, bottom): 40 | bottom[0].diff[...] = top[0].diff * bottom[1].data * self.sc 41 | 42 | 43 | class Permute(caffe.Layer): 44 | def setup(self, bottom, top): 45 | self.dims = [int(v) for v in self.param_str.split('_')] 46 | self.dims_ind = list(np.argsort(self.dims)) 47 | 48 | def reshape(self, bottom, top): 49 | old_shape = bottom[0].data.shape 50 | new_shape = [old_shape[d] for d in self.dims] 51 | top[0].reshape(*new_shape) 52 | 53 | def forward(self, bottom, top): 54 | top[0].data[...] = bottom[0].data.transpose(*self.dims) 55 | 56 | def backward(self, top, propagate_down, bottom): 57 | bottom[0].diff[...] = top[0].diff.transpose(*self.dims_ind) 58 | 59 | 60 | class LossHelper(caffe.Layer): 61 | def setup(self, bottom, top): 62 | self.old_shape = bottom[0].data.shape 63 | 64 | def reshape(self, bottom, top): 65 | new_shape = (self.old_shape[0] * self.old_shape[3], self.old_shape[1], 1, 1) 66 | top[0].reshape(*new_shape) 67 | 68 | def forward(self, bottom, top): 69 | top[0].data[...] = bottom[0].data.transpose(0, 3, 1, 2).reshape(*top[0].data.shape) 70 | 71 | def backward(self, top, propagate_down, bottom): 72 | bottom[0].diff[...] = top[0].diff.reshape(self.old_shape[0], self.old_shape[3], self.old_shape[1], 1 73 | ).transpose(0, 2, 3, 1) 74 | 75 | 76 | class LogLoss(caffe.Layer): 77 | def setup(self, bottom, top): 78 | self.n, self.c, _, self.s = bottom[0].data.shape 79 | self.inds = np.ix_(np.arange(self.n), np.arange(self.c), np.arange(1), np.arange(self.s)) 80 | 81 | def reshape(self, bottom, top): 82 | top[0].reshape(1, 1, 1, 1) 83 | 84 | def forward(self, bottom, top): 85 | self.valid = bottom[0].data[self.inds[0], bottom[1].data.astype(int), self.inds[2], self.inds[3]] 86 | top[0].data[:] = -np.mean(np.log(self.valid + 1e-10)) 87 | 88 | def backward(self, top, propagate_down, bottom): 89 | bottom[0].diff[:] = 0.0 90 | bottom[0].diff[self.inds[0], bottom[1].data.astype(int), self.inds[2], self.inds[3]] = \ 91 | -1.0 / ((self.valid + 1e-10) * (self.n * self.s)) 92 | 93 | 94 | class PickAndScale(caffe.Layer): 95 | def setup(self, bottom, top): 96 | self.nch_out = len(self.param_str.split('_')) 97 | self.dims = [] 98 | for f in self.param_str.split('_'): 99 | if f.find('*') >= 0: 100 | self.dims.append((int(f[:f.find('*')]), float(f[f.find('*') + 1:]))) 101 | elif f.find('/') >= 0: 102 | self.dims.append((int(f[:f.find('/')]), 1.0 / float(f[f.find('/') + 1:]))) 103 | else: 104 | self.dims.append((int(f), 1.0)) 105 | 106 | def reshape(self, bottom, top): 107 | top[0].reshape(bottom[0].data.shape[0], self.nch_out, bottom[0].data.shape[2], bottom[0].data.shape[3]) 108 | 109 | def forward(self, bottom, top): 110 | for i, (j, s) in enumerate(self.dims): 111 | top[0].data[:, i, :, :] = bottom[0].data[:, j, :, :] * s 112 | 113 | def backward(self, top, propagate_down, bottom): 114 | for i, (j, s) in enumerate(self.dims): 115 | bottom[0].diff[:, i, :, :] = top[0].diff[:, j, :, :] *s 116 | 117 | class NormAndScale(caffe.Layer): 118 | def setup(self, bottom, top): 119 | self.nch_out = len(self.param_str.split('_')) 120 | self.dims = [] 121 | for f in self.param_str.split('_'): 122 | if f.find('*') >= 0: 123 | self.dims.append((int(f[:f.find('*')]), float(f[f.find('*') + 1:]))) 124 | elif f.find('/') >= 0: 125 | self.dims.append((int(f[:f.find('/')]), 1.0 / float(f[f.find('/') + 1:]))) 126 | else: 127 | self.dims.append((int(f), 1.0)) 128 | 129 | def reshape(self, bottom, top): 130 | top[0].reshape(bottom[0].data.shape[0], self.nch_out, bottom[0].data.shape[2], bottom[0].data.shape[3]) 131 | 132 | def forward(self, bottom, top): 133 | norm = np.linalg.norm(bottom[1].data) 134 | for i, (j, s) in enumerate(self.dims): 135 | top[0].data[:, i, :, :] = bottom[0].data[:, j, :, :] * s/ norm 136 | 137 | def backward(self, top, propagate_down, bottom): 138 | norm = np.linalg.norm(bottom[1].data) 139 | for i, (j, s) in enumerate(self.dims): 140 | bottom[0].diff[:, i, :, :] = top[0].diff[:, j, :, :] *s /norm 141 | 142 | class PickAndScale_noback(caffe.Layer): 143 | def setup(self, bottom, top): 144 | self.nch_out = len(self.param_str.split('_')) 145 | self.dims = [] 146 | for f in self.param_str.split('_'): 147 | if f.find('*') >= 0: 148 | self.dims.append((int(f[:f.find('*')]), float(f[f.find('*') + 1:]))) 149 | elif f.find('/') >= 0: 150 | self.dims.append((int(f[:f.find('/')]), 1.0 / float(f[f.find('/') + 1:]))) 151 | else: 152 | self.dims.append((int(f), 1.0)) 153 | 154 | def reshape(self, bottom, top): 155 | top[0].reshape(bottom[0].data.shape[0], self.nch_out, bottom[0].data.shape[2], bottom[0].data.shape[3]) 156 | 157 | def forward(self, bottom, top): 158 | for i, (j, s) in enumerate(self.dims): 159 | top[0].data[:, i, :, :] = bottom[0].data[:, j, :, :] * s 160 | 161 | def backward(self, top, propagate_down, bottom): 162 | pass 163 | '''for i, (j, s) in enumerate(self.dims): 164 | bottom[0].diff[:, i, :, :] = top[0].diff[:, j, :, :] *s''' 165 | 166 | class GraphFeature(caffe.Layer): 167 | def setup(self, bottom, top): 168 | self.numNN = int(self.param_str) 169 | self.numpts = int(bottom[0].data.shape[3]) 170 | self.point_cloud = np.squeeze(bottom[0].data).transpose(0, 2, 1) 171 | adj = utils.pairwise_distance(self.point_cloud) 172 | self.nn_idx = utils.knn(adj, k=self.numNN) 173 | 174 | 175 | def reshape(self, bottom, top): 176 | top[0].reshape(bottom[0].data.shape[0], bottom[0].data.shape[1], bottom[0].data.shape[2]*(self.numNN + 1) , bottom[0].data.shape[3]) 177 | 178 | def forward(self, bottom, top): 179 | for ind in range(bottom[0].data.shape[0]): 180 | for i in range(self.numpts): 181 | top[0].data[ind, :, 0, i] = bottom[0].data[ind, :, 0, i] 182 | for j in range(self.numNN ): 183 | top[0].data[ind, :, j+1, i] = bottom[0].data[ind, :, 0, self.nn_idx[ind , i, j]] - bottom[0].data[ind, :, 0, i] 184 | 185 | def backward(self, top, propagate_down, bottom): 186 | for ind in range(bottom[0].data.shape[0]): 187 | for i in range(self.numpts): 188 | top[0].diff[ind, :, 0, i] = bottom[0].diff[ind, :, 0, i] 189 | for j in range(self.numNN ): 190 | top[0].diff[ind, :, j+1, i] = bottom[0].diff[ind, :, :, self.nn_idx[ind , i, j]] - bottom[0].diff[ind, :, :, i] -------------------------------------------------------------------------------- /sampling-based/splatnet/semseg3d/test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import glob 7 | import argparse 8 | import time 9 | import numpy as np 10 | import caffe 11 | import splatnet.configs 12 | from splatnet.utils import modify_blob_shape 13 | from splatnet import plot_log 14 | import splatnet.configs 15 | 16 | 17 | EVAL_SCRIPT_PATH = os.path.join(splatnet.configs.ROOT_DIR, 'splatnet', 'semseg3d', 'eval_seg.py') 18 | 19 | 20 | def extract_feat_scene(network_path, weights_path, feed, out_names, batch_size=1, sample_size=-1): 21 | net = caffe.Net(network_path, weights_path, caffe.TEST) 22 | net_bs, _, _, net_ss = net.blobs[list(feed.keys())[0]].data.shape 23 | 24 | npt = list(feed.values())[0].shape[-1] 25 | if sample_size == -1: 26 | sample_size = npt 27 | elif sample_size == 0: 28 | sample_size = net_ss 29 | else: 30 | assert sample_size * batch_size <= npt 31 | 32 | if net_bs != batch_size or net_ss != sample_size: 33 | network_path = modify_blob_shape(network_path, feed.keys(), {0: batch_size, 3: sample_size}) 34 | net = caffe.Net(network_path, weights_path, caffe.TEST) 35 | 36 | if type(out_names) == str: 37 | out_names = (out_names,) 38 | single_target = True 39 | else: 40 | single_target = False 41 | 42 | outs = {v: [] for v in out_names} 43 | pts_per_batch = batch_size * sample_size 44 | for b in range(int(np.ceil(npt / pts_per_batch))): 45 | b_end = min(pts_per_batch * (b + 1), npt) 46 | b_slice = slice(b_end - pts_per_batch, b_end) 47 | bs = min(pts_per_batch, npt - pts_per_batch * b) 48 | 49 | for in_key in feed: 50 | net.blobs[in_key].data[...] \ 51 | = feed[in_key][:, :, :, b_slice].reshape(-1, batch_size, sample_size, 1).transpose(1, 0, 3, 2) 52 | net.forward() 53 | for out_key in out_names: 54 | out_sz = net.blobs[out_key].data.shape 55 | out = net.blobs[out_key].data.transpose(1, 2, 0, 3).reshape(1, out_sz[1], out_sz[2], -1)[:, :, :, -bs:] 56 | outs[out_key].append(out.copy()) 57 | 58 | result = {v: np.concatenate(outs[v], axis=3) for v in out_names} 59 | if single_target: 60 | result = result[out_names[0]] 61 | 62 | return result 63 | 64 | 65 | def semseg_test(dataset, network, weights, input_dims='nx_ny_nz_r_g_b_h', sample_size=-1, 66 | dataset_params=None, save_dir='', save_prefix='', use_cpu=False): 67 | """ 68 | Testing trained semantic segmentation network 69 | :param dataset: choices: 'facade', 'stanford3d' 70 | :param network: path to a .prototxt file 71 | :param weights: path to a .caffemodel file 72 | :param input_dims: feat dims and scales 73 | :param sample_size: -1 -- use all points in a single sample, 0 -- use the size in network 74 | :param dataset_params: a dict with optional dataset parameters 75 | :param save_dir: default '' 76 | :param save_prefix: default '' 77 | :param use_cpu: default False 78 | :return: 79 | """ 80 | 81 | if use_cpu: 82 | caffe.set_mode_cpu() 83 | else: 84 | caffe.set_mode_gpu() 85 | caffe.set_device(0) 86 | 87 | # dataset specific: data, xyz, cmap 88 | if dataset == 'facade': 89 | from splatnet.dataset import dataset_facade 90 | dataset_params_new = {} if not dataset_params else dataset_params 91 | dataset_params = dict(subset='test', val_ratio=0.0) # default values 92 | dataset_params.update(dataset_params_new) 93 | for v in {'val_ratio'}: 94 | if v in dataset_params: 95 | dataset_params[v] = float(dataset_params[v]) 96 | data, xyz, norms = dataset_facade.points(dims=input_dims+',x_y_z,nx_ny_nz', **dataset_params) 97 | cmap = splatnet.configs.FACADE_CMAP 98 | elif dataset == 'stanford3d': 99 | norms = None 100 | pass # TODO set cmap, data, xyz 101 | else: 102 | raise ValueError('Unsupported dataset: {}'.format(dataset)) 103 | 104 | tic = time.time() 105 | prob = extract_feat_scene(network, weights, 106 | feed=dict(data=data.transpose().reshape(1, -1, 1, len(data))), 107 | out_names='prob', 108 | sample_size=sample_size) 109 | elapsed = time.time() - tic 110 | 111 | pred = prob.argmax(axis=1).squeeze() 112 | 113 | if norms is None: 114 | out = np.array([np.concatenate((x, cmap[int(c)]), axis=0) for (x, c) in zip(xyz, pred)]) 115 | header = '''ply 116 | format ascii 1.0 117 | element vertex {} 118 | property float x 119 | property float y 120 | property float z 121 | property uchar diffuse_red 122 | property uchar diffuse_green 123 | property uchar diffuse_blue 124 | end_header'''.format(len(data)) 125 | fmt = '%.6f %.6f %.6f %d %d %d' 126 | else: 127 | out = np.array([np.concatenate((x, n, cmap[int(c)]), axis=0) for (x, n, c) in zip(xyz, norms, pred)]) 128 | header = '''ply 129 | format ascii 1.0 130 | element vertex {} 131 | property float x 132 | property float y 133 | property float z 134 | property float nx 135 | property float ny 136 | property float nz 137 | property uchar diffuse_red 138 | property uchar diffuse_green 139 | property uchar diffuse_blue 140 | end_header'''.format(len(data)) 141 | fmt = '%.6f %.6f %.6f %.6f %.6f %.6f %d %d %d' 142 | 143 | save_path = os.path.join(save_dir, '{}pred_{}.ply'.format(save_prefix, dataset_params['subset'])) 144 | np.savetxt(save_path, out, fmt=fmt, header=header, comments='') 145 | 146 | return save_path, elapsed, len(data) 147 | 148 | 149 | if __name__ == '__main__': 150 | parser = argparse.ArgumentParser(add_help=False) 151 | 152 | group = parser.add_argument_group('testing options') 153 | group.add_argument('dataset') 154 | group.add_argument('--dataset_params', nargs='+', help='dataset-specific parameters (key value pairs)') 155 | group.add_argument('--input', default='nx_ny_nz_r_g_b_h', help='features to use as input') 156 | group.add_argument('--cpu', action='store_true', help='use cpu') 157 | group.add_argument('--exp_dir', default=None, type=str, help='together with exp_prefix, set defaults to args below') 158 | group.add_argument('--exp_prefix', default='', type=str, help='together with exp_dir, set defaults to args below') 159 | group.add_argument('--network', default=None, type=str, help='a .prototxt file') 160 | group.add_argument('--weights', default=None, type=str, help='a .caffemodel file') 161 | group.add_argument('--sample_size', default=-1, type=int, help='testing sample size') 162 | group.add_argument('--log', default=None, type=str, help='a .log file with training logs') 163 | group.add_argument('--log_eval', default=None, type=str, help='path to write evaluation logs') 164 | group.add_argument('--save_dir', default=None, type=str, help='together with save_prefix, a place for predictions') 165 | group.add_argument('--save_prefix', default=None, type=str, help='together with save_dir, a place for predictions') 166 | 167 | group = parser.add_argument_group('evaluation options') 168 | group.add_argument('--gt', default=None, type=str, help='path to ground-truth') 169 | group.add_argument('--gt_rgb', action='store_true', help='turn this on if gt is encoded with rgb values') 170 | group.add_argument('--gt_column', type=int, default=11, help='(starting) column of label in gt (1-index)') 171 | 172 | parser = argparse.ArgumentParser(description='Testing trained semantic segmentation network', 173 | parents=[parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter) 174 | args = parser.parse_args() 175 | 176 | network, weights = args.network, args.weights 177 | save_dir, save_prefix = args.save_dir, args.save_prefix 178 | log_train = args.log 179 | log_eval = args.log_eval 180 | 181 | if args.exp_dir is not None: 182 | if args.exp_prefix: 183 | exp_prefix = args.exp_prefix + '_' 184 | snapshot_prefix = exp_prefix 185 | else: 186 | exp_prefix = '' 187 | snapshot_prefix = 'snapshot_' 188 | if network is None: 189 | network = os.path.join(args.exp_dir, exp_prefix + 'net_deploy.prototxt') 190 | if weights is None: 191 | last_iter = max([int(os.path.split(v)[1].split('_')[-1][:-11]) for v in 192 | glob.glob(os.path.join(args.exp_dir, snapshot_prefix + 'iter_*.caffemodel'))]) 193 | weights = os.path.join(args.exp_dir, '{}iter_{}.caffemodel'.format(snapshot_prefix, last_iter)) 194 | if log_train is None: 195 | log_train = os.path.join(args.exp_dir, exp_prefix + 'train.log') 196 | if not os.path.exists(log_train): 197 | log_train = '' 198 | if save_dir is None: 199 | save_dir = args.exp_dir 200 | if save_prefix is None: 201 | save_prefix = exp_prefix 202 | if log_eval is None: 203 | log_eval = os.path.join(args.exp_dir, exp_prefix + 'test.log') 204 | 205 | if not args.dataset_params: 206 | args.dataset_params = {} 207 | else: 208 | args.dataset_params = dict(zip(args.dataset_params[::2], args.dataset_params[1::2])) 209 | 210 | pred_path, elapsed, num_pts = semseg_test(args.dataset, network, weights, args.input, args.sample_size, 211 | args.dataset_params, save_dir, save_prefix, args.cpu) 212 | 213 | if log_eval: 214 | with open(log_eval, 'a') as f: 215 | f.write('Predictions saved to {}.\n'.format(pred_path)) 216 | f.write('{} points evaluated in {:.2f} secs.\n'.format(num_pts, elapsed)) 217 | 218 | if args.gt is not None: 219 | import subprocess 220 | subprocess.run('python {} --dataset {} {}' 221 | '{} --pred_rgb --pred_column 7 ' 222 | '{} {}--gt_column {}'.format(EVAL_SCRIPT_PATH, 223 | args.dataset, 224 | '--log {} '.format(log_eval) if log_eval else '', 225 | os.path.abspath(pred_path), 226 | os.path.abspath(args.gt), 227 | '--gt_rgb ' if args.gt_rgb else '', 228 | args.gt_column).split(' ')) 229 | 230 | if log_train: 231 | plot_log.parse_and_plot(log_train) 232 | 233 | -------------------------------------------------------------------------------- /point_based/part_seg/test_deform.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import tensorflow as tf 3 | import json 4 | import numpy as np 5 | import os 6 | import sys 7 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 8 | sys.path.append(BASE_DIR) 9 | sys.path.append(os.path.dirname(BASE_DIR)) 10 | import provider 11 | import part_seg_model_deformF as model 12 | from pdb import set_trace as st 13 | 14 | parser = argparse.ArgumentParser() 15 | parser.add_argument('--model_path', default='train_results/trained_models/epoch_160.ckpt', help='Model checkpoint path') 16 | parser.add_argument('--graphnum', type=str ) 17 | parser.add_argument('--featnum', type=str ) 18 | parser.add_argument('--output_dir', type=str, default='affineF', help='Directory that stores all training logs and trained models') 19 | FLAGS = parser.parse_args() 20 | 21 | graphnum = np.array(FLAGS.graphnum.split(",")).astype(int) 22 | featnum = np.array(FLAGS.featnum.split(",")).astype(int) 23 | 24 | # DEFAULT SETTINGS 25 | pretrained_model_path = FLAGS.model_path 26 | hdf5_data_dir = os.path.join(BASE_DIR, './hdf5_data') 27 | ply_data_dir = os.path.join(BASE_DIR, './PartAnnotation') 28 | gpu_to_use = 0 29 | output_dir = os.path.join(BASE_DIR, FLAGS.output_dir ) 30 | output_verbose = False 31 | 32 | # MAIN SCRIPT 33 | point_num = 3000 34 | batch_size = 1 35 | 36 | test_file_list = os.path.join(BASE_DIR, 'testing_ply_file_list.txt') 37 | 38 | oid2cpid = json.load(open(os.path.join(hdf5_data_dir, 'overallid_to_catid_partid.json'), 'r')) 39 | 40 | object2setofoid = {} 41 | for idx in range(len(oid2cpid)): 42 | objid, pid = oid2cpid[idx] 43 | if not objid in object2setofoid.keys(): 44 | object2setofoid[objid] = [] 45 | object2setofoid[objid].append(idx) 46 | 47 | all_obj_cat_file = os.path.join(hdf5_data_dir, 'all_object_categories.txt') 48 | fin = open(all_obj_cat_file, 'r') 49 | lines = [line.rstrip() for line in fin.readlines()] 50 | objcats = [line.split()[1] for line in lines] 51 | objnames = [line.split()[0] for line in lines] 52 | on2oid = {objcats[i]:i for i in range(len(objcats))} 53 | fin.close() 54 | 55 | color_map_file = os.path.join(hdf5_data_dir, 'part_color_mapping.json') 56 | color_map = json.load(open(color_map_file, 'r')) 57 | 58 | NUM_OBJ_CATS = 16 59 | NUM_PART_CATS = 50 60 | 61 | cpid2oid = json.load(open(os.path.join(hdf5_data_dir, 'catid_partid_to_overallid.json'), 'r')) 62 | 63 | def printout(flog, data): 64 | print(data) 65 | flog.write(data + '\n') 66 | 67 | def output_color_point_cloud(data, seg, out_file): 68 | with open(out_file, 'w') as f: 69 | l = len(seg) 70 | for i in range(l): 71 | color = color_map[seg[i]] 72 | f.write('v %f %f %f %f %f %f\n' % (data[i][0], data[i][1], data[i][2], color[0], color[1], color[2])) 73 | 74 | def output_color_point_cloud_red_blue(data, seg, out_file): 75 | with open(out_file, 'w') as f: 76 | l = len(seg) 77 | for i in range(l): 78 | if seg[i] == 1: 79 | color = [0, 0, 1] 80 | elif seg[i] == 0: 81 | color = [1, 0, 0] 82 | else: 83 | color = [0, 0, 0] 84 | 85 | f.write('v %f %f %f %f %f %f\n' % (data[i][0], data[i][1], data[i][2], color[0], color[1], color[2])) 86 | 87 | 88 | def pc_normalize(pc): 89 | l = pc.shape[0] 90 | centroid = np.mean(pc, axis=0) 91 | pc = pc - centroid 92 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 93 | pc = pc / m 94 | return pc 95 | 96 | def placeholder_inputs(): 97 | pointclouds_ph = tf.placeholder(tf.float32, shape=(batch_size, point_num, 3)) 98 | input_label_ph = tf.placeholder(tf.float32, shape=(batch_size, NUM_OBJ_CATS)) 99 | return pointclouds_ph, input_label_ph 100 | 101 | def output_color_point_cloud(data, seg, out_file): 102 | with open(out_file, 'w') as f: 103 | l = len(seg) 104 | for i in range(l): 105 | color = color_map[seg[i]] 106 | f.write('v %f %f %f %f %f %f\n' % (data[i][0], data[i][1], data[i][2], color[0], color[1], color[2])) 107 | 108 | def load_pts_seg_files(pts_file, seg_file, catid): 109 | with open(pts_file, 'r') as f: 110 | pts_str = [item.rstrip() for item in f.readlines()] 111 | pts = np.array([np.float32(s.split()) for s in pts_str], dtype=np.float32) 112 | with open(seg_file, 'r') as f: 113 | part_ids = np.array([int(item.rstrip()) for item in f.readlines()], dtype=np.uint8) 114 | seg = np.array([cpid2oid[catid+'_'+str(x)] for x in part_ids]) 115 | return pts, seg 116 | 117 | def pc_augment_to_point_num(pts, pn): 118 | assert(pts.shape[0] <= pn) 119 | cur_len = pts.shape[0] 120 | res = np.array(pts) 121 | while cur_len < pn: 122 | res = np.concatenate((res, pts)) 123 | cur_len += pts.shape[0] 124 | return res[:pn, :] 125 | 126 | def convert_label_to_one_hot(labels): 127 | label_one_hot = np.zeros((labels.shape[0], NUM_OBJ_CATS)) 128 | for idx in range(labels.shape[0]): 129 | label_one_hot[idx, labels[idx]] = 1 130 | return label_one_hot 131 | 132 | def predict(): 133 | is_training = False 134 | 135 | with tf.device('/gpu:'+str(gpu_to_use)): 136 | pointclouds_ph, input_label_ph = placeholder_inputs() 137 | is_training_ph = tf.placeholder(tf.bool, shape=()) 138 | 139 | seg_pred = model.get_model(pointclouds_ph, input_label_ph, \ 140 | cat_num=NUM_OBJ_CATS, part_num=NUM_PART_CATS, is_training=is_training_ph, \ 141 | batch_size=batch_size, num_point=point_num, weight_decay=0.0, bn_decay=None, graphnum=graphnum, featnum = featnum) 142 | 143 | saver = tf.train.Saver() 144 | 145 | config = tf.ConfigProto() 146 | config.gpu_options.allow_growth = True 147 | config.allow_soft_placement = True 148 | 149 | with tf.Session(config=config) as sess: 150 | if not os.path.exists(output_dir): 151 | os.mkdir(output_dir) 152 | 153 | flog = open(os.path.join(output_dir, 'log.txt'), 'a') 154 | 155 | printout(flog, 'Loading model %s' % pretrained_model_path) 156 | saver.restore(sess, pretrained_model_path) 157 | printout(flog, 'Model restored.') 158 | 159 | batch_data = np.zeros([batch_size, point_num, 3]).astype(np.float32) 160 | 161 | total_acc = 0.0 162 | total_seen = 0 163 | total_acc_iou = 0.0 164 | 165 | total_per_cat_acc = np.zeros((NUM_OBJ_CATS)).astype(np.float32) 166 | total_per_cat_iou = np.zeros((NUM_OBJ_CATS)).astype(np.float32) 167 | total_per_cat_seen = np.zeros((NUM_OBJ_CATS)).astype(np.int32) 168 | 169 | ffiles = open(test_file_list, 'r') 170 | lines = [line.rstrip() for line in ffiles.readlines()] 171 | pts_files = [line.split()[0] for line in lines] 172 | seg_files = [line.split()[1] for line in lines] 173 | labels = [line.split()[2] for line in lines] 174 | ffiles.close() 175 | 176 | len_pts_files = len(pts_files) 177 | for shape_idx in range(len_pts_files): 178 | if shape_idx % 100 == 0: 179 | printout(flog, '%d/%d ...' % (shape_idx, len_pts_files)) 180 | 181 | cur_gt_label = on2oid[labels[shape_idx]] # 0/1/.../15 182 | 183 | cur_label_one_hot = np.zeros((1, NUM_OBJ_CATS), dtype=np.float32) 184 | cur_label_one_hot[0, cur_gt_label] = 1 185 | 186 | pts_file_to_load = os.path.join(ply_data_dir, pts_files[shape_idx]) 187 | seg_file_to_load = os.path.join(ply_data_dir, seg_files[shape_idx]) 188 | 189 | pts, seg = load_pts_seg_files(pts_file_to_load, seg_file_to_load, objcats[cur_gt_label]) 190 | ori_point_num = len(seg) 191 | 192 | batch_data[0, ...] = pc_augment_to_point_num(pc_normalize(pts), point_num) 193 | 194 | seg_pred_res = sess.run(seg_pred, feed_dict={ 195 | pointclouds_ph: batch_data, 196 | input_label_ph: cur_label_one_hot, 197 | is_training_ph: is_training}) 198 | 199 | """np.save(os.path.join("saved_shape", str(cur_gt_label), str(shape_idx)), pts) 200 | for i, tran in enumerate(trans): 201 | for j, tra in enumerate(tran): 202 | np.save(os.path.join("saved_shape", str(cur_gt_label), str(shape_idx)+"_"+str(i)+"_"+str(j)) , np.squeeze(tra))""" 203 | 204 | seg_pred_res = seg_pred_res[0, ...] 205 | 206 | iou_oids = object2setofoid[objcats[cur_gt_label]] 207 | non_cat_labels = list(set(np.arange(NUM_PART_CATS)).difference(set(iou_oids))) 208 | 209 | mini = np.min(seg_pred_res) 210 | seg_pred_res[:, non_cat_labels] = mini - 1000 211 | 212 | seg_pred_val = np.argmax(seg_pred_res, axis=1)[:ori_point_num] 213 | 214 | seg_acc = np.mean(seg_pred_val == seg) 215 | 216 | total_acc += seg_acc 217 | total_seen += 1 218 | 219 | total_per_cat_seen[cur_gt_label] += 1 220 | total_per_cat_acc[cur_gt_label] += seg_acc 221 | 222 | mask = np.int32(seg_pred_val == seg) 223 | 224 | total_iou = 0.0 225 | iou_log = '' 226 | for oid in iou_oids: 227 | n_pred = np.sum(seg_pred_val == oid) 228 | n_gt = np.sum(seg == oid) 229 | n_intersect = np.sum(np.int32(seg == oid) * mask) 230 | n_union = n_pred + n_gt - n_intersect 231 | iou_log += '_' + str(n_pred)+'_'+str(n_gt)+'_'+str(n_intersect)+'_'+str(n_union)+'_' 232 | if n_union == 0: 233 | total_iou += 1 234 | iou_log += '_1\n' 235 | else: 236 | total_iou += n_intersect * 1.0 / n_union 237 | iou_log += '_'+str(n_intersect * 1.0 / n_union)+'\n' 238 | 239 | avg_iou = total_iou / len(iou_oids) 240 | total_acc_iou += avg_iou 241 | total_per_cat_iou[cur_gt_label] += avg_iou 242 | 243 | if output_verbose: 244 | output_color_point_cloud(pts, seg, os.path.join(output_dir, str(shape_idx)+'_gt.obj')) 245 | output_color_point_cloud(pts, seg_pred_val, os.path.join(output_dir, str(shape_idx)+'_pred.obj')) 246 | output_color_point_cloud_red_blue(pts, np.int32(seg == seg_pred_val), 247 | os.path.join(output_dir, str(shape_idx)+'_diff.obj')) 248 | 249 | with open(os.path.join(output_dir, str(shape_idx)+'.log'), 'w') as fout: 250 | fout.write('Total Point: %d\n\n' % ori_point_num) 251 | fout.write('Ground Truth: %s\n' % objnames[cur_gt_label]) 252 | fout.write('Accuracy: %f\n' % seg_acc) 253 | fout.write('IoU: %f\n\n' % avg_iou) 254 | fout.write('IoU details: %s\n' % iou_log) 255 | 256 | printout(flog, 'Accuracy: %f' % (total_acc / total_seen)) 257 | printout(flog, 'IoU: %f' % (total_acc_iou / total_seen)) 258 | 259 | for cat_idx in range(NUM_OBJ_CATS): 260 | printout(flog, '\t ' + objcats[cat_idx] + ' Total Number: ' + str(total_per_cat_seen[cat_idx])) 261 | if total_per_cat_seen[cat_idx] > 0: 262 | printout(flog, '\t ' + objcats[cat_idx] + ' Accuracy: ' + \ 263 | str(total_per_cat_acc[cat_idx] / total_per_cat_seen[cat_idx])) 264 | printout(flog, '\t ' + objcats[cat_idx] + ' IoU: '+ \ 265 | str(total_per_cat_iou[cat_idx] / total_per_cat_seen[cat_idx])) 266 | 267 | with tf.Graph().as_default(): 268 | predict() 269 | -------------------------------------------------------------------------------- /sampling-based/splatnet/dataset/dataset_facade.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import numpy as np 7 | from numpy.linalg import eig 8 | import caffe 9 | from splatnet.utils import rotate_3d 10 | from splatnet.configs import FACADE_DATA_DIR 11 | 12 | 13 | def ordered_points(subset, dims='x_y_z_nx_ny_nz_r_g_b_h,l', order_dim='z', val_ratio=0.0, root=FACADE_DATA_DIR): 14 | pcl_train_path = os.path.join(root, 'pcl_train.ply') 15 | pcl_test_path = os.path.join(root, 'pcl_test.ply') 16 | 17 | # data files have 11 columns: (x, y, z, nx, ny, nz, r, g, b, height, label) 18 | feat_dict = dict(zip('x_y_z_nx_ny_nz_r_g_b_h_l'.split('_'), range(11))) 19 | feat_idxs, feat_scales = [], [] 20 | for g in dims.split(','): 21 | feat_idxs.append([]) 22 | feat_scales.append([]) 23 | for f in g.split('_'): 24 | if f.find('*') >= 0: 25 | feat_idxs[-1].append(feat_dict[f[:f.find('*')]]) 26 | feat_scales[-1].append(float(f[f.find('*') + 1:])) 27 | elif f.find('/') >= 0: 28 | feat_idxs[-1].append(feat_dict[f[:f.find('/')]]) 29 | feat_scales[-1].append(1.0 / float(f[f.find('/') + 1:])) 30 | else: 31 | feat_idxs[-1].append(feat_dict[f]) 32 | feat_scales[-1].append(1.0) 33 | 34 | if subset == 'train': 35 | pcl_data = np.loadtxt(pcl_train_path, skiprows=15) 36 | num_val = int(val_ratio * len(pcl_data)) 37 | order_idx = np.argsort(pcl_data[:, 2])[num_val:] 38 | pcl_data = pcl_data[order_idx, :] 39 | elif subset == 'val': 40 | pcl_data = np.loadtxt(pcl_train_path, skiprows=15) 41 | num_val = int(val_ratio * len(pcl_data)) 42 | order_idx = np.argsort(pcl_data[:, 2])[:num_val] 43 | pcl_data = pcl_data[order_idx, :] 44 | elif subset == 'test': 45 | pcl_data = np.loadtxt(pcl_test_path, skiprows=15) 46 | else: 47 | raise ValueError('Unknown subset: ' + subset) 48 | 49 | order_dim = feat_dict[order_dim] 50 | if subset not in {'train', 'val'} or order_dim != 2: 51 | order_idx = np.argsort(pcl_data[:, order_dim]) 52 | pcl_data = pcl_data[order_idx, :] 53 | 54 | return tuple([pcl_data[:, idx] * sc for (idx, sc) in zip(feat_idxs, feat_scales)]) 55 | 56 | 57 | def points(subset, dims='x_y_z_nx_ny_nz_r_g_b_h,l', shuffle=False, val_ratio=0.0, root=FACADE_DATA_DIR): 58 | pcl_train_path = os.path.join(root, 'pcl_train.ply') 59 | pcl_test_path = os.path.join(root, 'pcl_test.ply') 60 | 61 | # data files have 11 columns: (x, y, z, nx, ny, nz, r, g, b, height, label) 62 | feat_dict = dict(zip('x_y_z_nx_ny_nz_r_g_b_h_l'.split('_'), range(11))) 63 | feat_idxs, feat_scales = [], [] 64 | for g in dims.split(','): 65 | feat_idxs.append([]) 66 | feat_scales.append([]) 67 | for f in g.split('_'): 68 | if f.find('*') >= 0: 69 | feat_idxs[-1].append(feat_dict[f[:f.find('*')]]) 70 | feat_scales[-1].append(float(f[f.find('*') + 1:])) 71 | elif f.find('/') >= 0: 72 | feat_idxs[-1].append(feat_dict[f[:f.find('/')]]) 73 | feat_scales[-1].append(1.0 / float(f[f.find('/') + 1:])) 74 | else: 75 | feat_idxs[-1].append(feat_dict[f]) 76 | feat_scales[-1].append(1.0) 77 | 78 | order_dim = 2 79 | if subset == 'train': 80 | pcl_data = np.loadtxt(pcl_train_path, skiprows=15) 81 | num_val = int(val_ratio * len(pcl_data)) 82 | order_idx = np.sort(np.argsort(pcl_data[:, order_dim])[num_val:]) 83 | pcl_data = pcl_data[order_idx, :] 84 | elif subset == 'val': 85 | pcl_data = np.loadtxt(pcl_train_path, skiprows=15) 86 | num_val = int(val_ratio * len(pcl_data)) 87 | order_idx = np.sort(np.argsort(pcl_data[:, order_dim])[:num_val]) 88 | pcl_data = pcl_data[order_idx, :] 89 | elif subset == 'test': 90 | pcl_data = np.loadtxt(pcl_test_path, skiprows=15) 91 | else: 92 | raise ValueError('Unknown subset: ' + subset) 93 | 94 | if shuffle: 95 | order_idx = np.random.permutation(len(pcl_data)) 96 | pcl_data = pcl_data[order_idx] 97 | 98 | return tuple([pcl_data[:, idx] * sc for (idx, sc) in zip(feat_idxs, feat_scales)]) 99 | 100 | 101 | class InputFacade(caffe.Layer): 102 | def _restart(self): 103 | self.data[...] = self.data_copy 104 | num_points = len(self.data) 105 | points_per_batch = len(self.data) 106 | if self.mode == 'random': 107 | idx = np.random.permutation(num_points) 108 | self.data, self.label = self.data[idx], self.label[idx] 109 | self.idx = 0 110 | elif self.mode == 'ordered': 111 | # pick a random starting index as a form of data augmentation 112 | # self.idx = np.random.randint(0, num_points - points_per_batch + 1) 113 | self.idx = np.random.randint(0, min(points_per_batch, num_points - points_per_batch + 1)) 114 | 115 | def setup(self, bottom, top): 116 | params = dict(mode='ordered', batch_size=1, sample_size=-1, 117 | jitter_color=0.5, jitter_h=0.001, jitter_rotation=True, 118 | subset='train', val_ratio=0.0, feat_dims='nx_ny_nz_r_g_b_h', 119 | root=FACADE_DATA_DIR) 120 | params.update(eval(self.param_str)) 121 | self.mode = params['mode'] 122 | self.batch_size = params['batch_size'] 123 | self.sample_size = params['sample_size'] 124 | 125 | feat_dims = [] 126 | for f in params['feat_dims'].split('_'): 127 | if f.find('*') >= 0: 128 | feat_dims.append((f[:f.find('*')], float(f[f.find('*')+1:]))) 129 | elif f.find('/') >= 0: 130 | feat_dims.append((f[:f.find('/')], 1.0 / float(f[f.find('/') + 1:]))) 131 | else: 132 | feat_dims.append((f, 1.0)) 133 | self.raw_dims = [] 134 | for feat_group in [['x', 'y', 'z'], ['nx', 'ny', 'nz'], ['r', 'g', 'b'], ['h']]: 135 | if np.any([f in feat_group for f, _ in feat_dims]): 136 | self.raw_dims.extend(feat_group) 137 | self.feat_scales = [(self.raw_dims.index(f), s) for f, s in feat_dims] 138 | 139 | self.data, self.label = ordered_points(params['subset'], 140 | dims='_'.join(self.raw_dims) + ',l', 141 | val_ratio=params['val_ratio'], 142 | root=params['root']) 143 | self.data_copy = self.data.copy() 144 | self.label -= 1 # label starts from 0 145 | self.top_names = ['data', 'label'] 146 | self.top_channels = [len(self.feat_scales), 1] 147 | 148 | if self.sample_size == -1: 149 | self.sample_size = self.data.shape[0] 150 | 151 | if len(self.data) < self.sample_size * self.batch_size: 152 | raise Exception('Too few samples ({}). Is batch size too large?'.format(len(self.data))) 153 | 154 | if len(top) != len(self.top_names): 155 | raise Exception('Incorrect number of outputs (expected %d, got %d)' % 156 | (len(self.top_names), len(top))) 157 | 158 | # prepare for jittering 159 | max_points = 100000 160 | part_idx = np.random.permutation(len(self.data))[:min(len(self.data), max_points)] 161 | if 'r' in self.raw_dims and params['jitter_color'] != 0: 162 | feat_idx = self.raw_dims.index('r') 163 | eigw, eigv = eig(np.cov(self.data[np.ix_(part_idx, range(feat_idx, feat_idx + 3))].T)) 164 | self.jitter_color = params['jitter_color'] * eigv * np.sqrt(eigw) 165 | else: 166 | self.jitter_color = None 167 | if 'h' in self.raw_dims and params['jitter_h'] != 0: 168 | feat_idx = self.raw_dims.index('h') 169 | std = np.std(self.data[part_idx, feat_idx]) 170 | self.jitter_h = params['jitter_h'] * std 171 | else: 172 | self.jitter_h = None 173 | self.jitter_rotation = params['jitter_rotation'] 174 | 175 | self._restart() 176 | 177 | def reshape(self, bottom, top): 178 | for top_index, name in enumerate(self.top_names): 179 | shape = (self.batch_size, self.top_channels[top_index], 1, self.sample_size) 180 | top[top_index].reshape(*shape) 181 | 182 | def forward(self, bottom, top): 183 | points_per_batch = self.sample_size * self.batch_size 184 | data, label = self.data[self.idx:self.idx+points_per_batch].reshape(self.batch_size, self.sample_size, -1), \ 185 | self.label[self.idx:self.idx+points_per_batch] 186 | 187 | # jittering TODO: improve performance (e.g. move some computation to done at once for all samples?) 188 | for i in range(self.batch_size): 189 | if self.jitter_color is not None: 190 | feat_idx = self.raw_dims.index('r') 191 | data[i, :, feat_idx:feat_idx+3] += np.random.randn(3).dot(self.jitter_color.T) 192 | # clipping to [0, 255] 193 | data[i, :, feat_idx:feat_idx+3] = np.maximum(0.0, np.minimum(255.0, data[i, :, feat_idx:feat_idx+3])) 194 | if self.jitter_h is not None: 195 | feat_idx = self.raw_dims.index('h') 196 | data[i, :, feat_idx] += np.random.randn(self.sample_size) * self.jitter_h 197 | if self.jitter_rotation: 198 | rotations = (('z', np.random.rand() * np.pi / 8 - np.pi / 16), 199 | ('x', np.random.rand() * np.pi / 8 - np.pi / 16), 200 | ('y', np.random.rand() * np.pi * 2)) 201 | if 'x' in self.raw_dims: 202 | feat_idx = self.raw_dims.index('x') 203 | center = (np.mean(data[i, :, feat_idx]), 204 | np.max(data[i, :, feat_idx + 1]), 205 | np.mean(data[i, :, feat_idx + 2])) 206 | data[i, :, feat_idx:feat_idx + 3] = rotate_3d(data[i, :, feat_idx:feat_idx + 3], rotations, center) 207 | if 'nx' in self.raw_dims: 208 | feat_idx = self.raw_dims.index('nx') 209 | data[i, :, feat_idx:feat_idx + 3] = rotate_3d(data[i, :, feat_idx:feat_idx + 3], rotations) 210 | 211 | # slicing and scaling 212 | idxs, scs = [v[0] for v in self.feat_scales], [v[1] for v in self.feat_scales] 213 | data = data[:, :, idxs] * np.array(scs) 214 | 215 | top[0].data[...] = data.reshape(self.batch_size, self.sample_size, -1, 1).transpose(0, 2, 3, 1) 216 | top[1].data[...] = label.reshape(self.batch_size, self.sample_size, -1, 1).transpose(0, 2, 3, 1) 217 | 218 | self.idx += points_per_batch 219 | if self.idx + points_per_batch > len(self.data): 220 | self._restart() 221 | 222 | def backward(self, top, propagate_down, bottom): 223 | pass 224 | 225 | -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/test.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import argparse 7 | import glob 8 | import time 9 | import numpy as np 10 | import caffe 11 | 12 | import splatnet.configs 13 | from splatnet import plot_log 14 | from splatnet.utils import modify_blob_shape, seg_scores 15 | 16 | 17 | def extract_feat_shapes(network_path, weights_path, feed, out_names, batch_size=64, sample_size=3000): 18 | if sample_size == -1: 19 | assert batch_size == 1 20 | 21 | ori_sample_sizes = [len(sample) for sample in list(feed.values())[0]] 22 | nsamples = len(ori_sample_sizes) 23 | 24 | batch_size = nsamples if nsamples < batch_size else batch_size 25 | 26 | net = caffe.Net(network_path, weights_path, caffe.TEST) 27 | net_bs, _, _, net_ss = net.blobs[list(feed.keys())[0]].data.shape 28 | 29 | if sample_size != -1 and (net_bs != batch_size or net_ss != sample_size): 30 | network_path = modify_blob_shape(network_path, feed.keys(), {0: batch_size, 3: sample_size}) 31 | net = caffe.Net(network_path, weights_path, caffe.TEST) 32 | 33 | # pad samples to fixed length 34 | if sample_size != -1: 35 | for i in range(nsamples): 36 | k = ori_sample_sizes[i] 37 | idx = np.concatenate((np.tile(np.arange(k), (sample_size // k, )), 38 | np.random.permutation(k)[:(sample_size % k)]), axis=0) 39 | feed = feed.copy() 40 | for in_key in feed: 41 | feed[in_key][i] = feed[in_key][i][idx] 42 | 43 | outs = {v: [] for v in out_names} 44 | for b in range(int(np.ceil(nsamples / batch_size))): 45 | b_end = min(batch_size * (b + 1), nsamples) 46 | b_slice = slice(b_end - batch_size, b_end) 47 | bs = min(batch_size, nsamples - batch_size * b) 48 | 49 | if sample_size == -1: 50 | ss = list(feed.values())[0][b_end - 1].shape[0] 51 | if net_ss != ss or net_bs != 1: 52 | network_path = modify_blob_shape(network_path, feed.keys(), {0: 1, 3: ss}) 53 | net = caffe.Net(network_path, weights_path, caffe.TEST) 54 | net_ss, net_bs = ss, 1 55 | else: 56 | ss = sample_size 57 | 58 | for in_key in feed: 59 | net.blobs[in_key].data[...] \ 60 | = np.concatenate(feed[in_key][b_slice], axis=0).reshape(batch_size, ss, -1, 1).transpose(0, 2, 3, 1) 61 | net.forward() 62 | for out_key in out_names: 63 | out = net.blobs[out_key].data.transpose(0, 3, 1, 2)[-bs:].reshape(bs * ss, -1) 64 | outs[out_key].append(out.copy()) 65 | 66 | if sample_size != -1: 67 | outs = {v: np.array_split(np.concatenate(outs[v], axis=0), nsamples) for v in out_names} 68 | # remove padding 69 | for v in out_names: 70 | outs[v] = [sample[:k] for (sample, k) in zip(outs[v], ori_sample_sizes)] 71 | 72 | return outs 73 | 74 | 75 | def partseg_test(dataset, network, weights, input_dims='x_y_z', sample_size=3000, batch_size=64, 76 | category='airplane', dataset_params=None, 77 | save_dir='', skip_ply=False, use_cpu=False): 78 | """ 79 | Testing trained segmentation network 80 | :param dataset: choices: 'shapenet' 81 | :param network: path to a .prototxt file 82 | :param weights: path to a .caffemodel file 83 | :param input_dims: network input featurse 84 | :param sample_size: 85 | :param batch_size: 86 | :param dataset_params: a dict with optional dataset parameters 87 | :param save_dir: default '' 88 | :param skip_ply: default False 89 | :param use_cpu: default False 90 | :return: 91 | """ 92 | 93 | if use_cpu: 94 | caffe.set_mode_cpu() 95 | else: 96 | caffe.set_mode_gpu() 97 | caffe.set_device(0) 98 | 99 | # dataset specific: data, xyz_norm_list, cmap, num_part_categories, category_offset 100 | if dataset == 'shapenet': 101 | import splatnet.dataset.dataset_shapenet as shapenet 102 | dataset_params_new = {} if not dataset_params else dataset_params 103 | dataset_params = dict(subset='test') # default values 104 | dataset_params.update(dataset_params_new) 105 | 106 | data, _, part_label, names = shapenet.points_single_category(dims=input_dims, category=category, **dataset_params) 107 | xyz_norm_list, _, _, _ = shapenet.points_single_category(dims='x_y_z_nx_ny_nz', category=category, **dataset_params) 108 | cmap = splatnet.configs.SN_CMAP 109 | if category.startswith('0'): 110 | category_id = splatnet.configs.SN_CATEGORIES.index(category) 111 | else: 112 | category_id = splatnet.configs.SN_CATEGORY_NAMES.index(category) 113 | num_part_categories = splatnet.configs.SN_NUM_PART_CATEGORIES[category_id] 114 | category_offset = sum(splatnet.configs.SN_NUM_PART_CATEGORIES[:category_id]) 115 | else: 116 | raise ValueError('Unsupported dataset: {}'.format(dataset)) 117 | 118 | tic = time.time() 119 | probs = extract_feat_shapes(network, weights, 120 | feed=dict(data=data), 121 | out_names=('prob',), 122 | sample_size=sample_size, batch_size=batch_size)['prob'] 123 | elapsed = time.time() - tic 124 | 125 | if probs[0].shape[1] > num_part_categories: 126 | preds = [np.argmax(prob[:, category_offset:category_offset+num_part_categories], axis=1) for prob in probs] 127 | else: 128 | preds = [np.argmax(prob, axis=1) for prob in probs] 129 | 130 | acc, avgacc, avgiou = seg_scores(preds, part_label, nclasses=num_part_categories) 131 | 132 | if not skip_ply: 133 | os.makedirs(save_dir, exist_ok=True) 134 | for xyz_norm, pred, name in zip(xyz_norm_list, preds, names): 135 | out = np.array([np.concatenate((x, cmap[int(c)]), axis=0) for (x, c) in zip(xyz_norm, pred)]) 136 | header = '''ply 137 | format ascii 1.0 138 | element vertex {} 139 | property float x 140 | property float y 141 | property float z 142 | property float nx 143 | property float ny 144 | property float nz 145 | property uchar diffuse_red 146 | property uchar diffuse_green 147 | property uchar diffuse_blue 148 | end_header'''.format(len(pred)) 149 | fmt = '%.6f %.6f %.6f %.6f %.6f %.6f %d %d %d' 150 | save_path = os.path.join(save_dir, '{}.ply'.format(name)) 151 | np.savetxt(save_path, out, fmt=fmt, header=header, comments='') 152 | 153 | return acc, avgacc, avgiou, elapsed 154 | 155 | 156 | if __name__ == '__main__': 157 | parser = argparse.ArgumentParser(description='Testing trained part segmentation network', 158 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) 159 | parser.add_argument('dataset') 160 | parser.add_argument('--dataset_params', nargs='+', help='dataset-specific parameters (key value pairs)') 161 | parser.add_argument('--categories', nargs='+', help='pick some categories, otherwise evaluate on all') 162 | parser.add_argument('--single_model', action='store_true', help='if True, use a single model across all categories') 163 | parser.add_argument('--input', default='x_y_z', help='features to use as input') 164 | parser.add_argument('--cpu', action='store_true', help='use cpu') 165 | parser.add_argument('--skip_ply', action='store_true', help='if True, do not output .ply prediction results') 166 | parser.add_argument('--sample_size', default=3000, type=int, help='testing sample size') 167 | parser.add_argument('--batch_size', default=64, type=int, help='testing sample size') 168 | parser.add_argument('--snapshot', default=None, type=str, help='snapshot rule - last|best_acc|best_loss|ITER') 169 | parser.add_argument('--exp_dir', default=None, type=str, help='if present, assigns values to options below') 170 | parser.add_argument('--network', default=None, type=str, help='a .prototxt file') 171 | parser.add_argument('--weights', default=None, type=str, help='a .caffemodel file; overwrites \'--snapshot\'') 172 | parser.add_argument('--log', default=None, type=str, help='a .log file with training logs') 173 | parser.add_argument('--log_eval', default=None, type=str, help='path to write evaluation logs') 174 | parser.add_argument('--save_dir', default=None, type=str) 175 | 176 | args = parser.parse_args() 177 | 178 | if not args.categories: 179 | if args.dataset == 'shapenet': 180 | args.categories = splatnet.configs.SN_CATEGORY_NAMES 181 | else: 182 | raise ValueError('Unsupported dataset: {}'.format(args.dataset)) 183 | 184 | dataset_params = dict(zip(args.dataset_params[::2], args.dataset_params[1::2])) if args.dataset_params else {} 185 | exp_dir = args.exp_dir 186 | save_dir = args.save_dir if args.save_dir else os.path.join(exp_dir, 'pred') 187 | log_eval = args.log_eval if args.log_eval else os.path.join(exp_dir, 'test.log') 188 | 189 | with open(log_eval, 'a') as f: 190 | f.write('category | #samples | acc | class-acc | iou | snapshot | predictions\n') 191 | f.write('(batch_size={}, sample_size={})\n'.format(args.batch_size, args.sample_size)) 192 | 193 | tic = time.time() 194 | ious = [] 195 | 196 | for category_cnt, category in enumerate(args.categories): 197 | if category_cnt == 0 or not args.single_model: 198 | if args.network: 199 | network = args.network 200 | else: 201 | network_prefix = 'net' if args.single_model else '{}_net'.format(category) 202 | network = os.path.join(exp_dir, '{}_deploy.prototxt'.format(network_prefix)) 203 | if args.log: 204 | log = args.log 205 | else: 206 | log = os.path.join(exp_dir, 'train.log' if args.single_model else category+'.log') 207 | 208 | if args.weights: 209 | assert args.weights.endswith('.caffemodel') 210 | weights = args.weights 211 | else: 212 | snapshot_prefix = 'snapshot' if args.single_model else category 213 | if not args.snapshot: 214 | weights = os.path.join(exp_dir, '{}.caffemodel'.format(snapshot_prefix)) 215 | else: 216 | iters_avail = [int(os.path.split(v)[1].split('_')[-1][:-11]) for v in 217 | glob.glob(os.path.join(exp_dir, snapshot_prefix + '_iter_*.caffemodel'))] 218 | if args.snapshot == 'last': 219 | snap_iter = max(iters_avail) 220 | elif args.snapshot == 'best_acc': 221 | _, stats_val = plot_log.parse_and_plot(log, skip_train=0, skip_test=0) 222 | iter_acc = dict(zip(stats_val[:, 0], stats_val[:, -2])) 223 | snap_iter = iters_avail[np.argmax([iter_acc[v] for v in iters_avail])] 224 | elif args.snapshot == 'best_loss': 225 | _, stats_val = plot_log.parse_and_plot(log, skip_train=0, skip_test=0) 226 | iter_loss = dict(zip(stats_val[:, 0], stats_val[:, -1])) 227 | snap_iter = iters_avail[np.argmin([iter_loss[v] for v in iters_avail])] 228 | elif args.snapshot.isdigit(): 229 | snap_iter = int(args.snapshot) 230 | else: 231 | raise ValueError('Unknown snapshot rule: {}'.format(args.snapshot)) 232 | weights = os.path.join(exp_dir, '{}_iter_{}.caffemodel'.format(snapshot_prefix, snap_iter)) 233 | 234 | acc, avgacc, avgiou, _ = partseg_test(args.dataset, network, weights, args.input, 235 | args.sample_size, args.batch_size, category, dataset_params, 236 | os.path.join(save_dir, category), args.skip_ply, args.cpu) 237 | ious.append((np.mean(avgiou), len(avgiou))) 238 | 239 | with open(log_eval, 'a') as f: 240 | f.write('{} {} {} {} {} {} {}\n'.format(category, len(avgiou), 241 | np.mean(acc), np.mean(avgacc), np.mean(avgiou), 242 | os.path.basename(weights), 243 | '-' if args.skip_ply else os.path.join(save_dir, category))) 244 | 245 | elapsed = time.time() - tic 246 | 247 | with open(log_eval, 'a') as f: 248 | f.write('\nOverall weighted average mean IOU: {}\n'.format( 249 | sum([iou * cnt for (iou, cnt) in ious]) / sum([cnt for (_, cnt) in ious]))) 250 | f.write('Time elapsed: {} seconds\n\n'.format(elapsed)) 251 | 252 | -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/train_sense.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import argparse 7 | import sys 8 | sys.path.insert(0, '/mnt/lustre/wangjiayun/repos/splatnet') 9 | import splatnet.configs 10 | import pdb 11 | from splatnet.partseg3d import models 12 | from splatnet import create_solver 13 | from multiprocessing import Process 14 | os.environ["GLOG_minloglevel"] = "0" 15 | import caffe 16 | 17 | 18 | def partseg_train(network, exp_dir, category, args): 19 | if args.cpu: 20 | caffe.set_mode_cpu() 21 | else: 22 | caffe.set_mode_gpu() 23 | caffe.mpi_init() 24 | 25 | if network == 'seq': 26 | batch_norm = True 27 | conv_weight_filler = 'xavier' 28 | network = models.partseg_seq(arch_str=args.arch, 29 | skip_str=args.skips, 30 | dataset=args.dataset, 31 | dataset_params=args.dataset_params, 32 | category=category, 33 | feat_dims_str=args.feat, 34 | lattice_dims_str=args.lattice, 35 | sample_size=args.sample_size, 36 | batch_size=args.batch_size, 37 | batchnorm=batch_norm, 38 | conv_weight_filler=conv_weight_filler, 39 | save_path=os.path.join(exp_dir, category + '_net.prototxt')) 40 | 41 | models.partseg_seq(deploy=True, 42 | arch_str=args.arch, 43 | skip_str=args.skips, 44 | dataset=args.dataset, 45 | dataset_params=args.dataset_params, 46 | category=category, 47 | feat_dims_str=args.feat, 48 | lattice_dims_str=args.lattice, 49 | sample_size=args.sample_size, 50 | batchnorm=batch_norm, 51 | save_path=os.path.join(exp_dir, category + '_net_deploy.prototxt')) 52 | else: 53 | assert network.endswith('.prototxt'), 'Please provide a valid prototxt file' 54 | print('Using network defined at {}'.format(network)) 55 | 56 | random_seed = 0 57 | debug_info = False 58 | 59 | solver = create_solver.standard_solver(network, 60 | network, 61 | os.path.join(exp_dir, category)+'_' +args.prefix, 62 | base_lr=args.base_lr, 63 | gamma=args.lr_decay, 64 | stepsize=args.stepsize, 65 | test_iter=args.test_iter, 66 | test_interval=args.test_interval, 67 | max_iter=args.num_iter, 68 | snapshot=args.snapshot_interval, 69 | solver_type=args.solver_type, 70 | weight_decay=args.weight_decay, 71 | iter_size=args.iter_size, 72 | debug_info=debug_info, 73 | random_seed=random_seed, 74 | save_path=os.path.join(exp_dir, category+'_solver.prototxt')) 75 | solver = caffe.get_solver(solver) 76 | 77 | if args.init_model: 78 | if args.init_model.endswith('.caffemodel'): 79 | solver.net.copy_from(args.init_model) 80 | else: 81 | solver.net.copy_from(os.path.join(exp_dir, '{}_iter_{}.caffemodel'.format(category, args.init_model))) 82 | 83 | if args.init_state: 84 | if args.init_state.endswith('.solverstate'): 85 | solver.restore(args.init_state) 86 | else: 87 | solver.restore(os.path.join(exp_dir, '{}_iter_{}.solverstate'.format(category, args.init_state))) 88 | 89 | solver.solve() 90 | caffe.mpi_fin() 91 | 92 | def partseg_train_single_model(network, exp_dir, args): 93 | 94 | if args.cpu: 95 | caffe.set_mode_cpu() 96 | else: 97 | caffe.set_mode_gpu() 98 | caffe.set_device(0) 99 | 100 | if network == 'seq': 101 | batch_norm = True 102 | conv_weight_filler = 'xavier' 103 | network = models.partseg_seq_combined_categories(arch_str=args.arch, 104 | skip_str=args.skips, 105 | renorm_class=args.renorm_class, 106 | dataset=args.dataset, 107 | dataset_params=args.dataset_params, 108 | feat_dims_str=args.feat, 109 | lattice_dims_str=args.lattice, 110 | sample_size=args.sample_size, 111 | batch_size=args.batch_size, 112 | batchnorm=batch_norm, 113 | conv_weight_filler=conv_weight_filler, 114 | save_path=os.path.join(exp_dir, 'net.prototxt')) 115 | 116 | models.partseg_seq_combined_categories(deploy=True, 117 | arch_str=args.arch, 118 | skip_str=args.skips, 119 | renorm_class=args.renorm_class, 120 | dataset=args.dataset, 121 | dataset_params=args.dataset_params, 122 | feat_dims_str=args.feat, 123 | lattice_dims_str=args.lattice, 124 | sample_size=args.sample_size, 125 | batchnorm=batch_norm, 126 | save_path=os.path.join(exp_dir, 'net_deploy.prototxt')) 127 | else: 128 | assert network.endswith('.prototxt'), 'Please provide a valid prototxt file' 129 | print('Using network defined at {}'.format(network)) 130 | 131 | random_seed = 0 132 | debug_info = False 133 | solver = create_solver.standard_solver(network, 134 | network, 135 | os.path.join(exp_dir, 'snapshot'), 136 | base_lr=args.base_lr, 137 | gamma=args.lr_decay, 138 | stepsize=args.stepsize, 139 | test_iter=args.test_iter, 140 | test_interval=args.test_interval, 141 | max_iter=args.num_iter, 142 | snapshot=args.snapshot_interval, 143 | solver_type=args.solver_type, 144 | weight_decay=args.weight_decay, 145 | iter_size=args.iter_size, 146 | debug_info=debug_info, 147 | random_seed=random_seed, 148 | save_path=os.path.join(exp_dir, 'solver.prototxt')) 149 | solver = caffe.get_solver(solver) 150 | 151 | if args.init_model: 152 | if args.init_model.endswith('.caffemodel'): 153 | solver.net.copy_from(args.init_model) 154 | else: 155 | solver.net.copy_from(os.path.join(exp_dir, 'snapshot_iter_{}.caffemodel'.format(args.init_model))) 156 | 157 | if args.init_state: 158 | if args.init_state.endswith('.solverstate'): 159 | solver.restore(args.init_state) 160 | else: 161 | solver.restore(os.path.join(exp_dir, 'snapshot_iter_{}.solverstate'.format(args.init_state))) 162 | solver.solve() 163 | 164 | 165 | if __name__ == '__main__': 166 | parser = argparse.ArgumentParser(add_help=False) 167 | 168 | group = parser.add_argument_group('network options') 169 | group.add_argument('--network', default='seq', help='network type (\'seq\') or path to a .prototxt file') 170 | group.add_argument('--arch', default='64_128_256_256', help='network architecture') 171 | group.add_argument('--dataset', default='shapenet', choices=('shapenet',), help='dataset') 172 | group.add_argument('--dataset_params', nargs='+', help='dataset-specific parameters (key value pairs)') 173 | group.add_argument('--categories', nargs='+', help='pick some categories, otherwise evaluate on all') 174 | group.add_argument('--single_model', action='store_true', help='if True, use a single model across all categories') 175 | group.add_argument('--renorm_class', action='store_true', help='if True, renormalize prediction in true class') 176 | group.add_argument('--skips', nargs='+', help='skip connections') 177 | group.add_argument('--sample_size', type=int, default=3000, help='number of points in a sample') 178 | group.add_argument('--batch_size', type=int, default=32, help='number of samples in a batch') 179 | group.add_argument('--feat', default='x_y_z', help='features to use as input') 180 | group.add_argument('--lattice', nargs='+', help='bnn lattice features and scales') 181 | 182 | group = parser.add_argument_group('solver options') 183 | group.add_argument('--base_lr', type=float, default=0.01, help='starting learning rate') 184 | group.add_argument('--lr_decay', type=float, default=0.1, help='learning rate decay rate') 185 | group.add_argument('--stepsize', type=int, default=1000, help='learning rate decay interval') 186 | group.add_argument('--num_iter', type=int, default=2000, help='number of iterations to train') 187 | group.add_argument('--iter_size', type=int, default=1, help='number of mini-batches per iteration') 188 | group.add_argument('--test_iter', type=int, default=10, help='number of iterations to use at each testing phase') 189 | group.add_argument('--test_interval', type=int, default=20, help='test every such iterations') 190 | group.add_argument('--snapshot_interval', type=int, default=2000, help='snapshot every such iterations') 191 | group.add_argument('--solver_type', default='ADAM', help='optimizer type') 192 | group.add_argument('--prefix', default=None, help='caffemodel prefix') 193 | group.add_argument('--weight_decay', type=float, default=0.001, help='weight decay') 194 | 195 | parser = argparse.ArgumentParser(description='Object part segmentation training', 196 | parents=[parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter) 197 | parser.add_argument('exp_dir') 198 | parser.add_argument('--cpu', action='store_true', help='use cpu') 199 | parser.add_argument('--init_model', default=None, type=str, help='a .caffemodel file, or a snapshot iter') 200 | parser.add_argument('--init_state', default=None, type=str, help='a .solverstate file, or a snapshot iter') 201 | parser.add_argument("--gpus", type=int, nargs='+', default=[0], 202 | help="List of device ids.") 203 | 204 | args = parser.parse_args() 205 | args.gpus = range(args.gpus[0]) 206 | print(args.gpus) 207 | if not args.dataset_params: 208 | args.dataset_params = {} 209 | else: 210 | args.dataset_params = dict(zip(args.dataset_params[::2], args.dataset_params[1::2])) 211 | 212 | assert not (args.init_model and args.init_state) 213 | 214 | if args.exp_dir != '': 215 | if not os.path.exists(args.exp_dir): 216 | os.makedirs(args.exp_dir) 217 | # os.makedirs(args.exp_dir, exist_ok=True) 218 | 219 | if args.single_model: 220 | partseg_train_single_model(args.network, args.exp_dir, args) 221 | else: 222 | if not args.categories: 223 | if args.dataset == 'shapenet': 224 | categories = splatnet.configs.SN_CATEGORY_NAMES 225 | else: 226 | raise ValueError('Unsupported dataset: {}'.format(args.dataset)) 227 | else: 228 | categories = args.categories 229 | del args.categories 230 | 231 | for category in categories: 232 | partseg_train(args.network, args.exp_dir, category, args) 233 | 234 | -------------------------------------------------------------------------------- /sampling-based/splatnet/partseg3d/train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | import os 6 | import argparse 7 | import sys 8 | sys.path.insert(0, '/home/peterwg/repos/deformableSplatnet') 9 | import splatnet.configs 10 | import pdb 11 | from splatnet.partseg3d import models 12 | from splatnet import create_solver 13 | from multiprocessing import Process 14 | os.environ["GLOG_minloglevel"] = "0" 15 | import caffe 16 | 17 | 18 | def partseg_train(network, exp_dir, category, args): 19 | def solve2(solver, args, uid, rank): 20 | if args.cpu: 21 | caffe.set_mode_cpu() 22 | else: 23 | caffe.set_mode_gpu() 24 | caffe.set_device(args.gpus[rank]) 25 | caffe.set_solver_count(len(args.gpus)) 26 | caffe.set_solver_rank(rank) 27 | caffe.set_multiprocess(True) 28 | 29 | solver = caffe.get_solver(solver) 30 | 31 | if args.init_model: 32 | if args.init_model.endswith('.caffemodel'): 33 | solver.net.copy_from(args.init_model) 34 | else: 35 | solver.net.copy_from(os.path.join(exp_dir, '{}_iter_{}.caffemodel'.format(category, args.init_model))) 36 | 37 | if args.init_state: 38 | if args.init_state.endswith('.solverstate'): 39 | solver.restore(args.init_state) 40 | else: 41 | solver.restore(os.path.join(exp_dir, '{}_iter_{}.solverstate'.format(category, args.init_state))) 42 | 43 | nccl = caffe.NCCL(solver, uid) 44 | nccl.bcast() 45 | if solver.param.layer_wise_reduce: 46 | solver.net.after_backward(nccl) 47 | print(rank) 48 | #pdb.set_trace() 49 | solver.step(solver.param.max_iter) 50 | #solver.solve() 51 | 52 | #caffe.set_device(0) 53 | 54 | if network == 'seq': 55 | batch_norm = True 56 | conv_weight_filler = 'xavier' 57 | network = models.partseg_seq(arch_str=args.arch, 58 | skip_str=args.skips, 59 | dataset=args.dataset, 60 | dataset_params=args.dataset_params, 61 | category=category, 62 | feat_dims_str=args.feat, 63 | lattice_dims_str=args.lattice, 64 | sample_size=args.sample_size, 65 | batch_size=args.batch_size, 66 | batchnorm=batch_norm, 67 | conv_weight_filler=conv_weight_filler, 68 | save_path=os.path.join(exp_dir, category + '_net.prototxt')) 69 | 70 | models.partseg_seq(deploy=True, 71 | arch_str=args.arch, 72 | skip_str=args.skips, 73 | dataset=args.dataset, 74 | dataset_params=args.dataset_params, 75 | category=category, 76 | feat_dims_str=args.feat, 77 | lattice_dims_str=args.lattice, 78 | sample_size=args.sample_size, 79 | batchnorm=batch_norm, 80 | save_path=os.path.join(exp_dir, category + '_net_deploy.prototxt')) 81 | else: 82 | assert network.endswith('.prototxt'), 'Please provide a valid prototxt file' 83 | print('Using network defined at {}'.format(network)) 84 | 85 | random_seed = 0 86 | debug_info = False 87 | 88 | solver = create_solver.standard_solver(network, 89 | network, 90 | os.path.join(exp_dir, category)+'_' +args.prefix, 91 | base_lr=args.base_lr, 92 | gamma=args.lr_decay, 93 | stepsize=args.stepsize, 94 | test_iter=args.test_iter, 95 | test_interval=args.test_interval, 96 | max_iter=args.num_iter, 97 | snapshot=args.snapshot_interval, 98 | solver_type=args.solver_type, 99 | weight_decay=args.weight_decay, 100 | iter_size=args.iter_size, 101 | debug_info=debug_info, 102 | random_seed=random_seed, 103 | save_path=os.path.join(exp_dir, category+'_solver.prototxt')) 104 | ## Multiple GPUs 105 | uid = caffe.NCCL.new_uid() 106 | 107 | caffe.init_log(0, True) 108 | caffe.log('Using devices %s' % str(args.gpus)) 109 | procs = [] 110 | 111 | for rank in range(len(args.gpus)): 112 | p = Process(target=solve2, 113 | args=(solver, args, uid, rank)) 114 | p.daemon = True 115 | p.start() 116 | procs.append(p) 117 | for p in procs: 118 | p.join() 119 | 120 | 121 | def partseg_train_single_model(network, exp_dir, args): 122 | 123 | if args.cpu: 124 | caffe.set_mode_cpu() 125 | else: 126 | caffe.set_mode_gpu() 127 | caffe.set_device(0) 128 | 129 | if network == 'seq': 130 | batch_norm = True 131 | conv_weight_filler = 'xavier' 132 | network = models.partseg_seq_combined_categories(arch_str=args.arch, 133 | skip_str=args.skips, 134 | renorm_class=args.renorm_class, 135 | dataset=args.dataset, 136 | dataset_params=args.dataset_params, 137 | feat_dims_str=args.feat, 138 | lattice_dims_str=args.lattice, 139 | sample_size=args.sample_size, 140 | batch_size=args.batch_size, 141 | batchnorm=batch_norm, 142 | conv_weight_filler=conv_weight_filler, 143 | save_path=os.path.join(exp_dir, 'net.prototxt')) 144 | 145 | models.partseg_seq_combined_categories(deploy=True, 146 | arch_str=args.arch, 147 | skip_str=args.skips, 148 | renorm_class=args.renorm_class, 149 | dataset=args.dataset, 150 | dataset_params=args.dataset_params, 151 | feat_dims_str=args.feat, 152 | lattice_dims_str=args.lattice, 153 | sample_size=args.sample_size, 154 | batchnorm=batch_norm, 155 | save_path=os.path.join(exp_dir, 'net_deploy.prototxt')) 156 | else: 157 | assert network.endswith('.prototxt'), 'Please provide a valid prototxt file' 158 | print('Using network defined at {}'.format(network)) 159 | 160 | random_seed = 0 161 | debug_info = False 162 | solver = create_solver.standard_solver(network, 163 | network, 164 | os.path.join(exp_dir, 'snapshot'), 165 | base_lr=args.base_lr, 166 | gamma=args.lr_decay, 167 | stepsize=args.stepsize, 168 | test_iter=args.test_iter, 169 | test_interval=args.test_interval, 170 | max_iter=args.num_iter, 171 | snapshot=args.snapshot_interval, 172 | solver_type=args.solver_type, 173 | weight_decay=args.weight_decay, 174 | iter_size=args.iter_size, 175 | debug_info=debug_info, 176 | random_seed=random_seed, 177 | save_path=os.path.join(exp_dir, 'solver.prototxt')) 178 | solver = caffe.get_solver(solver) 179 | 180 | if args.init_model: 181 | if args.init_model.endswith('.caffemodel'): 182 | solver.net.copy_from(args.init_model) 183 | else: 184 | solver.net.copy_from(os.path.join(exp_dir, 'snapshot_iter_{}.caffemodel'.format(args.init_model))) 185 | 186 | if args.init_state: 187 | if args.init_state.endswith('.solverstate'): 188 | solver.restore(args.init_state) 189 | else: 190 | solver.restore(os.path.join(exp_dir, 'snapshot_iter_{}.solverstate'.format(args.init_state))) 191 | solver.solve() 192 | 193 | 194 | if __name__ == '__main__': 195 | parser = argparse.ArgumentParser(add_help=False) 196 | 197 | group = parser.add_argument_group('network options') 198 | group.add_argument('--network', default='seq', help='network type (\'seq\') or path to a .prototxt file') 199 | group.add_argument('--arch', default='64_128_256_256', help='network architecture') 200 | group.add_argument('--dataset', default='shapenet', choices=('shapenet',), help='dataset') 201 | group.add_argument('--dataset_params', nargs='+', help='dataset-specific parameters (key value pairs)') 202 | group.add_argument('--categories', nargs='+', help='pick some categories, otherwise evaluate on all') 203 | group.add_argument('--single_model', action='store_true', help='if True, use a single model across all categories') 204 | group.add_argument('--renorm_class', action='store_true', help='if True, renormalize prediction in true class') 205 | group.add_argument('--skips', nargs='+', help='skip connections') 206 | group.add_argument('--sample_size', type=int, default=3000, help='number of points in a sample') 207 | group.add_argument('--batch_size', type=int, default=32, help='number of samples in a batch') 208 | group.add_argument('--feat', default='x_y_z', help='features to use as input') 209 | group.add_argument('--lattice', nargs='+', help='bnn lattice features and scales') 210 | 211 | group = parser.add_argument_group('solver options') 212 | group.add_argument('--base_lr', type=float, default=0.01, help='starting learning rate') 213 | group.add_argument('--lr_decay', type=float, default=0.1, help='learning rate decay rate') 214 | group.add_argument('--stepsize', type=int, default=1000, help='learning rate decay interval') 215 | group.add_argument('--num_iter', type=int, default=2000, help='number of iterations to train') 216 | group.add_argument('--iter_size', type=int, default=1, help='number of mini-batches per iteration') 217 | group.add_argument('--test_iter', type=int, default=10, help='number of iterations to use at each testing phase') 218 | group.add_argument('--test_interval', type=int, default=20, help='test every such iterations') 219 | group.add_argument('--snapshot_interval', type=int, default=2000, help='snapshot every such iterations') 220 | group.add_argument('--solver_type', default='ADAM', help='optimizer type') 221 | group.add_argument('--prefix', default=None, help='caffemodel prefix') 222 | group.add_argument('--weight_decay', type=float, default=0.001, help='weight decay') 223 | 224 | parser = argparse.ArgumentParser(description='Object part segmentation training', 225 | parents=[parser], formatter_class=argparse.ArgumentDefaultsHelpFormatter) 226 | parser.add_argument('exp_dir') 227 | parser.add_argument('--cpu', action='store_true', help='use cpu') 228 | parser.add_argument('--init_model', default=None, type=str, help='a .caffemodel file, or a snapshot iter') 229 | parser.add_argument('--init_state', default=None, type=str, help='a .solverstate file, or a snapshot iter') 230 | parser.add_argument("--gpus", type=int, nargs='+', default=[0], 231 | help="List of device ids.") 232 | 233 | args = parser.parse_args() 234 | args.gpus = range(args.gpus[0]) 235 | print(args.gpus) 236 | if not args.dataset_params: 237 | args.dataset_params = {} 238 | else: 239 | args.dataset_params = dict(zip(args.dataset_params[::2], args.dataset_params[1::2])) 240 | 241 | assert not (args.init_model and args.init_state) 242 | 243 | if args.exp_dir != '': 244 | if not os.path.exists(args.exp_dir): 245 | os.makedirs(args.exp_dir) 246 | # os.makedirs(args.exp_dir, exist_ok=True) 247 | 248 | if args.single_model: 249 | partseg_train_single_model(args.network, args.exp_dir, args) 250 | else: 251 | if not args.categories: 252 | if args.dataset == 'shapenet': 253 | categories = splatnet.configs.SN_CATEGORY_NAMES 254 | else: 255 | raise ValueError('Unsupported dataset: {}'.format(args.dataset)) 256 | else: 257 | categories = args.categories 258 | del args.categories 259 | 260 | for category in categories: 261 | partseg_train(args.network, args.exp_dir, category, args) 262 | 263 | -------------------------------------------------------------------------------- /sampling-based/splatnet/semseg3d/models.py: -------------------------------------------------------------------------------- 1 | """ 2 | Copyright (C) 2018 NVIDIA Corporation. All rights reserved. 3 | Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). 4 | """ 5 | from functools import reduce 6 | import caffe 7 | from caffe import layers as L, params as P 8 | import splatnet.configs 9 | from splatnet.utils import get_prototxt, parse_channel_scale, map_channel_scale 10 | 11 | 12 | def semseg_seq(arch_str='64_128_256_256', batchnorm=True, 13 | skip_str=(), # tuple of strings like '4_1_ga' - relu4 <- relu1 w/ options 'ga' 14 | bilateral_nbr=1, 15 | conv_weight_filler='xavier', bltr_weight_filler='gauss_0.001', 16 | dataset='facade', dataset_params=None, 17 | sample_size=10000, batch_size=1, 18 | feat_dims_str='nx_ny_nz_r_g_b_h', lattice_dims_str=None, 19 | deploy=False, create_prototxt=True, save_path=None): 20 | 21 | n = caffe.NetSpec() 22 | 23 | arch_str = [(v[0], int(v[1:])) if v[0] in {'b', 'c'} else ('c', int(v)) for v in arch_str.split('_')] 24 | num_bltr_layers = sum(v[0] == 'b' for v in arch_str) 25 | 26 | if num_bltr_layers > 0: 27 | if type(lattice_dims_str) == str: 28 | lattice_dims_str = (lattice_dims_str,) * num_bltr_layers 29 | elif len(lattice_dims_str) == 1: 30 | lattice_dims_str = lattice_dims_str * num_bltr_layers 31 | else: 32 | assert len(lattice_dims_str) == num_bltr_layers, '{} lattices should be provided'.format(num_bltr_layers) 33 | feat_dims = parse_channel_scale(feat_dims_str, channel_str=True)[0] 34 | lattice_dims = [parse_channel_scale(s, channel_str=True)[0] for s in lattice_dims_str] 35 | input_dims_w_dup = feat_dims + reduce(lambda x, y: x + y, lattice_dims) 36 | input_dims = reduce(lambda x, y: x if y in x else x + [y], input_dims_w_dup, []) 37 | feat_dims_str = map_channel_scale(feat_dims_str, input_dims) 38 | lattice_dims_str = [map_channel_scale(s, input_dims) for s in lattice_dims_str] 39 | input_dims_str = '_'.join(input_dims) 40 | else: 41 | feat_dims = parse_channel_scale(feat_dims_str, channel_str=True)[0] 42 | input_dims = feat_dims 43 | feat_dims_str = map_channel_scale(feat_dims_str, input_dims) 44 | input_dims_str = '_'.join(input_dims) 45 | 46 | # dataset specific settings: nclass, datalayer_train, datalayer_test 47 | if dataset == 'facade': 48 | # default dataset params 49 | dataset_params_new = {} if not dataset_params else dataset_params 50 | dataset_params = dict(subset_train='train', subset_test='test', val_ratio=0.0) 51 | dataset_params.update(dataset_params_new) 52 | 53 | dataset_params['feat_dims'] = input_dims_str 54 | dataset_params['sample_size'] = sample_size 55 | dataset_params['batch_size'] = batch_size 56 | 57 | # dataset params type casting 58 | for v in {'jitter_color', 'jitter_h', 'val_ratio'}: 59 | if v in dataset_params: 60 | dataset_params[v] = float(dataset_params[v]) 61 | for v in {'sample_size', 'batch_size'}: 62 | if v in dataset_params: 63 | dataset_params[v] = int(dataset_params[v]) 64 | for v in {'jitter_rotation'}: 65 | if v in dataset_params: 66 | dataset_params[v] = False if dataset_params[v] == '0' else True 67 | 68 | # training time dataset params 69 | dataset_params_train = dataset_params.copy() 70 | dataset_params_train['subset'] = dataset_params['subset_train'] 71 | del dataset_params_train['subset_train'], dataset_params_train['subset_test'] 72 | 73 | # testing time dataset params: turn off all data augmentations 74 | dataset_params_test = dataset_params.copy() 75 | dataset_params_test['subset'] = dataset_params['subset_test'] 76 | dataset_params_test['jitter_color'] = 0.0 77 | dataset_params_test['jitter_h'] = 0.0 78 | dataset_params_test['jitter_rotation'] = False 79 | del dataset_params_test['subset_train'], dataset_params_test['subset_test'] 80 | 81 | nclass = len(splatnet.configs.FACADE_CATEGORIES) 82 | 83 | # data layers 84 | datalayer_train = L.Python(name='data', include=dict(phase=caffe.TRAIN), ntop=2, 85 | python_param=dict(module='dataset_facade', layer='InputFacade', 86 | param_str=repr(dataset_params_train))) 87 | datalayer_test = L.Python(name='data', include=dict(phase=caffe.TEST), ntop=0, top=['data', 'label'], 88 | python_param=dict(module='dataset_facade', layer='InputFacade', 89 | param_str=repr(dataset_params_test))) 90 | elif dataset == 'stanford3d': 91 | raise NotImplementedError() 92 | else: 93 | raise ValueError('Dataset {} unknown'.format(dataset)) 94 | 95 | # Input/Data layer 96 | if deploy: 97 | n.data = L.Input(shape=dict(dim=[1, len(input_dims), 1, sample_size])) 98 | else: 99 | n.data, n.label = datalayer_train 100 | n.test_data = datalayer_test 101 | n.data_feat = L.Python(n.data, python_param=dict(module='custom_layers', layer='PickAndScale', 102 | param_str=feat_dims_str)) 103 | top_prev = n.data_feat 104 | 105 | if conv_weight_filler in {'xavier', 'msra'}: 106 | conv_weight_filler = dict(type=conv_weight_filler) 107 | elif conv_weight_filler.startswith('gauss_'): 108 | conv_weight_filler = dict(type='gaussian', std=float(conv_weight_filler.split('_')[1])) 109 | else: 110 | conv_weight_filler = eval(conv_weight_filler) 111 | assert bltr_weight_filler.startswith('gauss_') 112 | bltr_weight_filler = dict(type='gaussian', std=float(bltr_weight_filler.split('_')[1])) 113 | 114 | # multiple 1x1 conv-(bn)-relu blocks, optionally with a single global pooling somewhere among them 115 | 116 | idx = 1 117 | bltr_idx = 0 118 | lattices = dict() 119 | last_in_block = dict() 120 | for (layer_type, n_out) in arch_str: 121 | if layer_type == 'c': 122 | n['conv' + str(idx)] = L.Convolution(top_prev, 123 | convolution_param=dict(num_output=n_out, 124 | kernel_size=1, stride=1, pad=0, 125 | weight_filler=conv_weight_filler, 126 | bias_filler=dict(type='constant', value=0)), 127 | param=[dict(lr_mult=1), dict(lr_mult=0.1)]) 128 | elif layer_type == 'b': 129 | lattice_dims_str_curr = lattice_dims_str[bltr_idx] 130 | if lattice_dims_str_curr in lattices: 131 | top_data_lattice, top_lattice = lattices[lattice_dims_str_curr] 132 | n['conv' + str(idx)] = L.Permutohedral(top_prev, top_data_lattice, top_data_lattice, top_lattice, 133 | permutohedral_param=dict(num_output=n_out, 134 | group=1, 135 | neighborhood_size=bilateral_nbr, 136 | bias_term=True, 137 | norm_type=P.Permutohedral.AFTER, 138 | offset_type=P.Permutohedral.NONE, 139 | filter_filler=bltr_weight_filler, 140 | bias_filler=dict(type='constant', 141 | value=0)), 142 | param=[{'lr_mult': 1, 'decay_mult': 1}, 143 | {'lr_mult': 2, 'decay_mult': 0}]) 144 | else: 145 | top_data_lattice = L.Python(n.data, python_param=dict(module='custom_layers', layer='PickAndScale', 146 | param_str=lattice_dims_str_curr)) 147 | n['data_lattice' + str(len(lattices))] = top_data_lattice 148 | if lattice_dims_str.count(lattice_dims_str_curr) > 1: 149 | n['conv' + str(idx)], top_lattice = L.Permutohedral(top_prev, top_data_lattice, top_data_lattice, 150 | ntop=2, 151 | permutohedral_param=dict( 152 | num_output=n_out, 153 | group=1, 154 | neighborhood_size=bilateral_nbr, 155 | bias_term=True, 156 | norm_type=P.Permutohedral.AFTER, 157 | offset_type=P.Permutohedral.NONE, 158 | filter_filler=bltr_weight_filler, 159 | bias_filler=dict(type='constant', 160 | value=0)), 161 | param=[{'lr_mult': 1, 'decay_mult': 1}, 162 | {'lr_mult': 2, 'decay_mult': 0}]) 163 | n['lattice' + str(len(lattices))] = top_lattice 164 | else: 165 | n['conv' + str(idx)] = L.Permutohedral(top_prev, top_data_lattice, top_data_lattice, 166 | permutohedral_param=dict( 167 | num_output=n_out, 168 | group=1, 169 | neighborhood_size=bilateral_nbr, 170 | bias_term=True, 171 | norm_type=P.Permutohedral.AFTER, 172 | offset_type=P.Permutohedral.NONE, 173 | filter_filler=bltr_weight_filler, 174 | bias_filler=dict(type='constant', value=0)), 175 | param=[{'lr_mult': 1, 'decay_mult': 1}, 176 | {'lr_mult': 2, 'decay_mult': 0}]) 177 | top_lattice = None 178 | 179 | lattices[lattice_dims_str_curr] = (top_data_lattice, top_lattice) 180 | 181 | bltr_idx += 1 182 | 183 | top_prev = n['conv' + str(idx)] 184 | if batchnorm: 185 | n['bn'+str(idx)] = L.BatchNorm(top_prev) 186 | top_prev = n['bn'+str(idx)] 187 | n['relu'+str(idx)] = L.ReLU(top_prev, in_place=True) 188 | top_prev = n['relu'+str(idx)] 189 | 190 | # skip connection & global pooling 191 | if skip_str is None: 192 | skip_str = () 193 | skip_tos = [v.split('_')[0] for v in skip_str] 194 | if str(idx) in skip_tos: 195 | skip_idxs = list(filter(lambda i: skip_tos[i] == str(idx), range(len(skip_tos)))) 196 | skip_params = [skip_str[i].split('_') for i in skip_idxs] 197 | if len(skip_params[0]) == 2: 198 | assert all(len(v) == 2 for v in skip_params) 199 | else: 200 | assert all(v[2] == skip_params[0][2] for v in skip_params) 201 | 202 | if len(skip_params[0]) > 2 and 'g' in skip_params[0][2]: # global pooling on current layer 203 | n['gpool' + str(idx)] = L.Python(top_prev, 204 | python_param=dict(module='custom_layers', layer='GlobalPooling')) 205 | top_prev = n['gpool' + str(idx)] 206 | 207 | if len(skip_params[0]) > 2 and 'a' in skip_params[0][2]: # addition instead of concatenation 208 | n['add' + str(idx)] = L.Eltwise(top_prev, *[last_in_block[int(v[1])] for v in skip_params], 209 | eltwise_param=dict(operation=P.Eltwise.SUM)) 210 | top_prev = n['add' + str(idx)] 211 | else: 212 | n['concat' + str(idx)] = L.Concat(top_prev, *[last_in_block[int(v[1])] for v in skip_params]) 213 | top_prev = n['concat' + str(idx)] 214 | 215 | last_in_block[idx] = top_prev 216 | idx += 1 217 | 218 | # classification & loss 219 | n['conv'+str(idx)] = L.Convolution(top_prev, 220 | convolution_param=dict(num_output=nclass, kernel_size=1, stride=1, pad=0, 221 | weight_filler=conv_weight_filler, 222 | bias_filler=dict(type='constant', value=0)), 223 | param=[dict(lr_mult=1), dict(lr_mult=0.1)]) 224 | top_prev = n['conv'+str(idx)] 225 | if deploy: 226 | n.prob = L.Softmax(top_prev) 227 | else: 228 | n.loss = L.SoftmaxWithLoss(top_prev, n.label) 229 | n.accuracy = L.Accuracy(top_prev, n.label) 230 | 231 | net = n.to_proto() 232 | 233 | if create_prototxt: 234 | net = get_prototxt(net, save_path) 235 | 236 | return net 237 | 238 | -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/car_net_deform.prototxt: -------------------------------------------------------------------------------- 1 | layer { 2 | name: "data" 3 | type: "Python" 4 | top: "data" 5 | top: "label" 6 | include { 7 | phase: TRAIN 8 | } 9 | python_param { 10 | module: "dataset_shapenet" 11 | layer: "InputShapenet" 12 | param_str: "{\'sample_size\': 3000, \'jitter_rotation\': 10.0, \'subset\': \'train\', \'category\': \'02958343\', \'root\': \'../../data/shapenet_ericyi_ply\', \'jitter_xyz\': 0.01, \'feat_dims\': \'x_y_z\', \'batch_size\': 28}" 13 | } 14 | } 15 | layer { 16 | name: "data" 17 | type: "Python" 18 | top: "data" 19 | top: "label" 20 | include { 21 | phase: TEST 22 | } 23 | python_param { 24 | module: "dataset_shapenet" 25 | layer: "InputShapenet" 26 | param_str: "{\'sample_size\': 3000, \'jitter_rotation\': 0.0, \'jitter_stretch\': 0.0, \'subset\': \'val\', \'category\': \'02958343\', \'root\': \'../../data/shapenet_ericyi_ply\', \'jitter_xyz\': 0.0, \'feat_dims\': \'x_y_z\', \'batch_size\': 28}" 27 | } 28 | } 29 | layer { 30 | name: "data_feat" 31 | type: "Python" 32 | bottom: "data" 33 | top: "data_feat" 34 | python_param { 35 | module: "custom_layers" 36 | layer: "PickAndScale" 37 | param_str: "0_1_2" 38 | } 39 | } 40 | layer { 41 | name: "conv1" 42 | type: "Convolution" 43 | bottom: "data_feat" 44 | top: "conv1" 45 | param { 46 | lr_mult: 1.0 47 | } 48 | param { 49 | lr_mult: 0.10000000149 50 | } 51 | convolution_param { 52 | num_output: 32 53 | pad: 0 54 | kernel_size: 1 55 | stride: 1 56 | weight_filler { 57 | type: "xavier" 58 | } 59 | bias_filler { 60 | type: "constant" 61 | value: 0.0 62 | } 63 | } 64 | } 65 | layer { 66 | name: "bn1" 67 | type: "BatchNorm" 68 | bottom: "conv1" 69 | top: "bn1" 70 | } 71 | layer { 72 | name: "relu1" 73 | type: "ReLU" 74 | bottom: "bn1" 75 | top: "bn1" 76 | } 77 | layer { 78 | name: "bnD1" 79 | type: "Concat" 80 | bottom: "data_feat" 81 | bottom: "bn1" 82 | top: "bnD1" 83 | } 84 | layer { 85 | name: "fconv1" 86 | type: "Convolution" 87 | bottom: "bnD1" 88 | top: "fconv1" 89 | param { 90 | name: "fconv1_w" 91 | lr_mult: 1.0 92 | } 93 | param { 94 | lr_mult: 0.10000000149 95 | } 96 | convolution_param { 97 | num_output: 3 98 | pad: 0 99 | kernel_size: 1 100 | stride: 1 101 | weight_filler { 102 | type: "xavier" 103 | } 104 | bias_filler { 105 | type: "constant" 106 | value: 0.0 107 | } 108 | } 109 | } 110 | 111 | layer { 112 | name: "param_fconv1" 113 | type: "Parameter" 114 | param { 115 | name: "fconv1_w" 116 | lr_mult: 1.0 117 | } 118 | parameter_param { 119 | shape { 120 | dim: 3 # copy the dimension from below 121 | dim: 35 122 | dim: 1 123 | dim: 1 124 | } 125 | } 126 | top: "param_fconv1" 127 | } 128 | layer { 129 | name: "data_lattice0" 130 | type: "Python" 131 | bottom: "fconv1" 132 | bottom: "param_fconv1" 133 | top: "data_lattice0" 134 | python_param { 135 | module: "custom_layers" 136 | layer: "NormAndScale" 137 | param_str: "0*64_1*64_2*64" 138 | } 139 | } 140 | layer { 141 | name: "conv2" 142 | type: "Permutohedral" 143 | bottom: "bn1" 144 | bottom: "data_lattice0" 145 | bottom: "data_lattice0" 146 | top: "conv2" 147 | param { 148 | lr_mult: 1.0 149 | decay_mult: 1.0 150 | } 151 | param { 152 | lr_mult: 2.0 153 | decay_mult: 0.0 154 | } 155 | permutohedral_param { 156 | num_output: 64 157 | neighborhood_size: 2 158 | filter_filler { 159 | type: "gaussian" 160 | std: 0.0010000000475 161 | } 162 | bias_term: true 163 | bias_filler { 164 | type: "constant" 165 | value: 0.0 166 | } 167 | group: 1 168 | norm_type: AFTER 169 | offset_type: NONE 170 | } 171 | } 172 | layer { 173 | name: "bn2" 174 | type: "BatchNorm" 175 | bottom: "conv2" 176 | top: "bn2" 177 | } 178 | layer { 179 | name: "relu2" 180 | type: "ReLU" 181 | bottom: "bn2" 182 | top: "bn2" 183 | } 184 | layer { 185 | name: "bnD2" 186 | type: "Concat" 187 | bottom: "data_feat" 188 | bottom: "bn2" 189 | top: "bnD2" 190 | } 191 | layer { 192 | name: "fconv2" 193 | type: "Convolution" 194 | bottom: "bnD2" 195 | top: "fconv2" 196 | param { 197 | name: "fconv2_w" 198 | lr_mult: 1.0 199 | } 200 | param { 201 | lr_mult: 0.10000000149 202 | } 203 | convolution_param { 204 | num_output: 3 205 | pad: 0 206 | kernel_size: 1 207 | stride: 1 208 | weight_filler { 209 | type: "xavier" 210 | } 211 | bias_filler { 212 | type: "constant" 213 | value: 0.0 214 | } 215 | } 216 | } 217 | layer { 218 | name: "param_fconv2" 219 | type: "Parameter" 220 | param { 221 | name: "fconv2_w" 222 | lr_mult: 1.0 223 | } 224 | parameter_param { 225 | shape { 226 | dim: 3 # copy the dimension from below 227 | dim: 67 228 | dim: 1 229 | dim: 1 230 | } 231 | } 232 | top: "param_fconv2" 233 | } 234 | layer { 235 | name: "data_lattice1" 236 | type: "Python" 237 | bottom: "fconv2" 238 | bottom: "param_fconv2" 239 | top: "data_lattice1" 240 | python_param { 241 | module: "custom_layers" 242 | layer: "NormAndScale" 243 | param_str: "0*32_1*32_2*32" 244 | } 245 | } 246 | layer { 247 | name: "conv3" 248 | type: "Permutohedral" 249 | bottom: "bn2" 250 | bottom: "data_lattice1" 251 | bottom: "data_lattice1" 252 | top: "conv3" 253 | param { 254 | lr_mult: 1.0 255 | decay_mult: 1.0 256 | } 257 | param { 258 | lr_mult: 2.0 259 | decay_mult: 0.0 260 | } 261 | permutohedral_param { 262 | num_output: 128 263 | neighborhood_size: 2 264 | filter_filler { 265 | type: "gaussian" 266 | std: 0.0010000000475 267 | } 268 | bias_term: true 269 | bias_filler { 270 | type: "constant" 271 | value: 0.0 272 | } 273 | group: 1 274 | norm_type: AFTER 275 | offset_type: NONE 276 | } 277 | } 278 | layer { 279 | name: "bn3" 280 | type: "BatchNorm" 281 | bottom: "conv3" 282 | top: "bn3" 283 | } 284 | layer { 285 | name: "relu3" 286 | type: "ReLU" 287 | bottom: "bn3" 288 | top: "bn3" 289 | } 290 | layer { 291 | name: "bnD3" 292 | type: "Concat" 293 | bottom: "data_feat" 294 | bottom: "bn3" 295 | top: "bnD3" 296 | } 297 | layer { 298 | name: "fconv3" 299 | type: "Convolution" 300 | bottom: "bnD3" 301 | top: "fconv3" 302 | param { 303 | name: "fconv3_w" 304 | lr_mult: 1.0 305 | } 306 | param { 307 | lr_mult: 0.10000000149 308 | } 309 | convolution_param { 310 | num_output: 3 311 | pad: 0 312 | kernel_size: 1 313 | stride: 1 314 | weight_filler { 315 | type: "xavier" 316 | } 317 | bias_filler { 318 | type: "constant" 319 | value: 0.0 320 | } 321 | } 322 | } 323 | 324 | layer { 325 | name: "param_fconv3" 326 | type: "Parameter" 327 | param { 328 | name: "fconv3_w" 329 | lr_mult: 1.0 330 | } 331 | parameter_param { 332 | shape { 333 | dim: 3 # copy the dimension from below 334 | dim: 131 335 | dim: 1 336 | dim: 1 337 | } 338 | } 339 | top: "param_fconv3" 340 | } 341 | layer { 342 | name: "data_lattice2" 343 | type: "Python" 344 | bottom: "fconv3" 345 | bottom: "param_fconv3" 346 | top: "data_lattice2" 347 | python_param { 348 | module: "custom_layers" 349 | layer: "NormAndScale" 350 | param_str: "0*16_1*16_2*16" 351 | } 352 | } 353 | layer { 354 | name: "conv4" 355 | type: "Permutohedral" 356 | bottom: "bn3" 357 | bottom: "data_lattice2" 358 | bottom: "data_lattice2" 359 | top: "conv4" 360 | param { 361 | lr_mult: 1.0 362 | decay_mult: 1.0 363 | } 364 | param { 365 | lr_mult: 2.0 366 | decay_mult: 0.0 367 | } 368 | permutohedral_param { 369 | num_output: 256 370 | neighborhood_size: 2 371 | filter_filler { 372 | type: "gaussian" 373 | std: 0.0010000000475 374 | } 375 | bias_term: true 376 | bias_filler { 377 | type: "constant" 378 | value: 0.0 379 | } 380 | group: 1 381 | norm_type: AFTER 382 | offset_type: NONE 383 | } 384 | } 385 | layer { 386 | name: "bn4" 387 | type: "BatchNorm" 388 | bottom: "conv4" 389 | top: "bn4" 390 | } 391 | layer { 392 | name: "relu4" 393 | type: "ReLU" 394 | bottom: "bn4" 395 | top: "bn4" 396 | } 397 | layer { 398 | name: "bnD4" 399 | type: "Concat" 400 | bottom: "data_feat" 401 | bottom: "bn4" 402 | top: "bnD4" 403 | } 404 | layer { 405 | name: "fconv4" 406 | type: "Convolution" 407 | bottom: "bnD4" 408 | top: "fconv4" 409 | param { 410 | name: "fconv4_w" 411 | lr_mult: 1.0 412 | } 413 | param { 414 | lr_mult: 0.10000000149 415 | } 416 | convolution_param { 417 | num_output: 3 418 | pad: 0 419 | kernel_size: 1 420 | stride: 1 421 | weight_filler { 422 | type: "xavier" 423 | } 424 | bias_filler { 425 | type: "constant" 426 | value: 0.0 427 | } 428 | } 429 | } 430 | 431 | layer { 432 | name: "param_fconv4" 433 | type: "Parameter" 434 | param { 435 | name: "fconv4_w" 436 | lr_mult: 1.0 437 | } 438 | parameter_param { 439 | shape { 440 | dim: 3 # copy the dimension from below 441 | dim: 259 442 | dim: 1 443 | dim: 1 444 | } 445 | } 446 | top: "param_fconv4" 447 | } 448 | layer { 449 | name: "data_lattice3" 450 | type: "Python" 451 | bottom: "fconv4" 452 | bottom: "param_fconv4" 453 | top: "data_lattice3" 454 | python_param { 455 | module: "custom_layers" 456 | layer: "NormAndScale" 457 | param_str: "0*8_1*8_2*8" 458 | } 459 | } 460 | layer { 461 | name: "conv5" 462 | type: "Permutohedral" 463 | bottom: "bn4" 464 | bottom: "data_lattice3" 465 | bottom: "data_lattice3" 466 | top: "conv5" 467 | param { 468 | lr_mult: 1.0 469 | decay_mult: 1.0 470 | } 471 | param { 472 | lr_mult: 2.0 473 | decay_mult: 0.0 474 | } 475 | permutohedral_param { 476 | num_output: 256 477 | neighborhood_size: 2 478 | filter_filler { 479 | type: "gaussian" 480 | std: 0.0010000000475 481 | } 482 | bias_term: true 483 | bias_filler { 484 | type: "constant" 485 | value: 0.0 486 | } 487 | group: 1 488 | norm_type: AFTER 489 | offset_type: NONE 490 | } 491 | } 492 | layer { 493 | name: "bn5" 494 | type: "BatchNorm" 495 | bottom: "conv5" 496 | top: "bn5" 497 | } 498 | layer { 499 | name: "relu5" 500 | type: "ReLU" 501 | bottom: "bn5" 502 | top: "bn5" 503 | } 504 | layer { 505 | name: "bnD5" 506 | type: "Concat" 507 | bottom: "data_feat" 508 | bottom: "bn5" 509 | top: "bnD5" 510 | } 511 | layer { 512 | name: "fconv5" 513 | type: "Convolution" 514 | bottom: "bnD5" 515 | top: "fconv5" 516 | param { 517 | name: "fconv5_w" 518 | lr_mult: 1.0 519 | } 520 | param { 521 | lr_mult: 0.10000000149 522 | } 523 | convolution_param { 524 | num_output: 3 525 | pad: 0 526 | kernel_size: 1 527 | stride: 1 528 | weight_filler { 529 | type: "xavier" 530 | } 531 | bias_filler { 532 | type: "constant" 533 | value: 0.0 534 | } 535 | } 536 | } 537 | 538 | layer { 539 | name: "param_fconv5" 540 | type: "Parameter" 541 | param { 542 | name: "fconv5_w" 543 | lr_mult: 1.0 544 | } 545 | parameter_param { 546 | shape { 547 | dim: 3 # copy the dimension from below 548 | dim: 259 549 | dim: 1 550 | dim: 1 551 | } 552 | } 553 | top: "param_fconv5" 554 | } 555 | layer { 556 | name: "data_lattice4" 557 | type: "Python" 558 | bottom: "fconv5" 559 | bottom: "param_fconv5" 560 | top: "data_lattice4" 561 | python_param { 562 | module: "custom_layers" 563 | layer: "NormAndScale" 564 | param_str: "0*4_1*4_2*4" 565 | } 566 | } 567 | layer { 568 | name: "conv6" 569 | type: "Permutohedral" 570 | bottom: "bn5" 571 | bottom: "data_lattice4" 572 | bottom: "data_lattice4" 573 | top: "conv6" 574 | param { 575 | lr_mult: 1.0 576 | decay_mult: 1.0 577 | } 578 | param { 579 | lr_mult: 2.0 580 | decay_mult: 0.0 581 | } 582 | permutohedral_param { 583 | num_output: 256 584 | neighborhood_size: 2 585 | filter_filler { 586 | type: "gaussian" 587 | std: 0.0010000000475 588 | } 589 | bias_term: true 590 | bias_filler { 591 | type: "constant" 592 | value: 0.0 593 | } 594 | group: 1 595 | norm_type: AFTER 596 | offset_type: NONE 597 | } 598 | } 599 | layer { 600 | name: "bn6" 601 | type: "BatchNorm" 602 | bottom: "conv6" 603 | top: "bn6" 604 | } 605 | layer { 606 | name: "relu6" 607 | type: "ReLU" 608 | bottom: "bn6" 609 | top: "bn6" 610 | } 611 | layer { 612 | name: "concat6" 613 | type: "Concat" 614 | bottom: "bn6" 615 | bottom: "bn2" 616 | bottom: "bn3" 617 | bottom: "bn4" 618 | bottom: "bn5" 619 | top: "concat6" 620 | } 621 | layer { 622 | name: "conv7" 623 | type: "Convolution" 624 | bottom: "concat6" 625 | top: "conv7" 626 | param { 627 | lr_mult: 1.0 628 | } 629 | param { 630 | lr_mult: 0.10000000149 631 | } 632 | convolution_param { 633 | num_output: 128 634 | pad: 0 635 | kernel_size: 1 636 | stride: 1 637 | weight_filler { 638 | type: "xavier" 639 | } 640 | bias_filler { 641 | type: "constant" 642 | value: 0.0 643 | } 644 | } 645 | } 646 | layer { 647 | name: "bn7" 648 | type: "BatchNorm" 649 | bottom: "conv7" 650 | top: "bn7" 651 | } 652 | layer { 653 | name: "relu7" 654 | type: "ReLU" 655 | bottom: "bn7" 656 | top: "bn7" 657 | } 658 | layer { 659 | name: "conv8" 660 | type: "Convolution" 661 | bottom: "bn7" 662 | top: "conv8" 663 | param { 664 | lr_mult: 1.0 665 | } 666 | param { 667 | lr_mult: 0.10000000149 668 | } 669 | convolution_param { 670 | num_output: 4 671 | pad: 0 672 | kernel_size: 1 673 | stride: 1 674 | weight_filler { 675 | type: "xavier" 676 | } 677 | bias_filler { 678 | type: "constant" 679 | value: 0.0 680 | } 681 | } 682 | } 683 | 684 | 685 | layer { 686 | name: "loss" 687 | type: "SoftmaxWithLoss" 688 | bottom: "conv8" 689 | bottom: "label" 690 | top: "loss" 691 | } 692 | layer { 693 | name: "accuracy" 694 | type: "Accuracy" 695 | bottom: "conv8" 696 | bottom: "label" 697 | top: "accuracy" 698 | } 699 | -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/lamp_net_deform.prototxt: -------------------------------------------------------------------------------- 1 | layer { 2 | name: "data" 3 | type: "Python" 4 | top: "data" 5 | top: "label" 6 | include { 7 | phase: TRAIN 8 | } 9 | python_param { 10 | module: "dataset_shapenet" 11 | layer: "InputShapenet" 12 | param_str: "{\'root\': \'../../data/shapenet_ericyi_ply\', \'batch_size\': 28, \'jitter_xyz\': 0.01, \'feat_dims\': \'x_y_z\', \'sample_size\': 3000, \'jitter_rotation\': 10.0, \'category\': \'03636649\', \'subset\': \'train\'}" 13 | } 14 | } 15 | layer { 16 | name: "data" 17 | type: "Python" 18 | top: "data" 19 | top: "label" 20 | include { 21 | phase: TEST 22 | } 23 | python_param { 24 | module: "dataset_shapenet" 25 | layer: "InputShapenet" 26 | param_str: "{\'jitter_stretch\': 0.0, \'root\': \'../../data/shapenet_ericyi_ply\', \'batch_size\': 28, \'jitter_xyz\': 0.0, \'feat_dims\': \'x_y_z\', \'sample_size\': 3000, \'jitter_rotation\': 0.0, \'category\': \'03636649\', \'subset\': \'val\'}" 27 | } 28 | } 29 | layer { 30 | name: "data_feat" 31 | type: "Python" 32 | bottom: "data" 33 | top: "data_feat" 34 | python_param { 35 | module: "custom_layers" 36 | layer: "PickAndScale" 37 | param_str: "0_1_2" 38 | } 39 | } 40 | layer { 41 | name: "conv1" 42 | type: "Convolution" 43 | bottom: "data_feat" 44 | top: "conv1" 45 | param { 46 | lr_mult: 1.0 47 | } 48 | param { 49 | lr_mult: 0.10000000149 50 | } 51 | convolution_param { 52 | num_output: 32 53 | pad: 0 54 | kernel_size: 1 55 | stride: 1 56 | weight_filler { 57 | type: "xavier" 58 | } 59 | bias_filler { 60 | type: "constant" 61 | value: 0.0 62 | } 63 | } 64 | } 65 | layer { 66 | name: "bn1" 67 | type: "BatchNorm" 68 | bottom: "conv1" 69 | top: "bn1" 70 | } 71 | layer { 72 | name: "relu1" 73 | type: "ReLU" 74 | bottom: "bn1" 75 | top: "bn1" 76 | } 77 | layer { 78 | name: "bnD1" 79 | type: "Concat" 80 | bottom: "data_feat" 81 | bottom: "bn1" 82 | top: "bnD1" 83 | } 84 | layer { 85 | name: "fconv1" 86 | type: "Convolution" 87 | bottom: "bnD1" 88 | top: "fconv1" 89 | param { 90 | name: "fconv1_w" 91 | lr_mult: 1.0 92 | } 93 | param { 94 | lr_mult: 0.10000000149 95 | } 96 | convolution_param { 97 | num_output: 3 98 | pad: 0 99 | kernel_size: 1 100 | stride: 1 101 | weight_filler { 102 | type: "xavier" 103 | } 104 | bias_filler { 105 | type: "constant" 106 | value: 0.0 107 | } 108 | } 109 | } 110 | 111 | layer { 112 | name: "param_fconv1" 113 | type: "Parameter" 114 | param { 115 | name: "fconv1_w" 116 | lr_mult: 1.0 117 | } 118 | parameter_param { 119 | shape { 120 | dim: 3 # copy the dimension from below 121 | dim: 35 122 | dim: 1 123 | dim: 1 124 | } 125 | } 126 | top: "param_fconv1" 127 | } 128 | layer { 129 | name: "data_lattice0" 130 | type: "Python" 131 | bottom: "fconv1" 132 | bottom: "param_fconv1" 133 | top: "data_lattice0" 134 | python_param { 135 | module: "custom_layers" 136 | layer: "NormAndScale" 137 | param_str: "0*64_1*64_2*64" 138 | } 139 | } 140 | layer { 141 | name: "conv2" 142 | type: "Permutohedral" 143 | bottom: "bn1" 144 | bottom: "data_lattice0" 145 | bottom: "data_lattice0" 146 | top: "conv2" 147 | param { 148 | lr_mult: 1.0 149 | decay_mult: 1.0 150 | } 151 | param { 152 | lr_mult: 2.0 153 | decay_mult: 0.0 154 | } 155 | permutohedral_param { 156 | num_output: 64 157 | neighborhood_size: 2 158 | filter_filler { 159 | type: "gaussian" 160 | std: 0.0010000000475 161 | } 162 | bias_term: true 163 | bias_filler { 164 | type: "constant" 165 | value: 0.0 166 | } 167 | group: 1 168 | norm_type: AFTER 169 | offset_type: NONE 170 | } 171 | } 172 | layer { 173 | name: "bn2" 174 | type: "BatchNorm" 175 | bottom: "conv2" 176 | top: "bn2" 177 | } 178 | layer { 179 | name: "relu2" 180 | type: "ReLU" 181 | bottom: "bn2" 182 | top: "bn2" 183 | } 184 | layer { 185 | name: "bnD2" 186 | type: "Concat" 187 | bottom: "data_feat" 188 | bottom: "bn2" 189 | top: "bnD2" 190 | } 191 | layer { 192 | name: "fconv2" 193 | type: "Convolution" 194 | bottom: "bnD2" 195 | top: "fconv2" 196 | param { 197 | name: "fconv2_w" 198 | lr_mult: 1.0 199 | } 200 | param { 201 | lr_mult: 0.10000000149 202 | } 203 | convolution_param { 204 | num_output: 3 205 | pad: 0 206 | kernel_size: 1 207 | stride: 1 208 | weight_filler { 209 | type: "xavier" 210 | } 211 | bias_filler { 212 | type: "constant" 213 | value: 0.0 214 | } 215 | } 216 | } 217 | layer { 218 | name: "param_fconv2" 219 | type: "Parameter" 220 | param { 221 | name: "fconv2_w" 222 | lr_mult: 1.0 223 | } 224 | parameter_param { 225 | shape { 226 | dim: 3 # copy the dimension from below 227 | dim: 67 228 | dim: 1 229 | dim: 1 230 | } 231 | } 232 | top: "param_fconv2" 233 | } 234 | layer { 235 | name: "data_lattice1" 236 | type: "Python" 237 | bottom: "fconv2" 238 | bottom: "param_fconv2" 239 | top: "data_lattice1" 240 | python_param { 241 | module: "custom_layers" 242 | layer: "NormAndScale" 243 | param_str: "0*32_1*32_2*32" 244 | } 245 | } 246 | layer { 247 | name: "conv3" 248 | type: "Permutohedral" 249 | bottom: "bn2" 250 | bottom: "data_lattice1" 251 | bottom: "data_lattice1" 252 | top: "conv3" 253 | param { 254 | lr_mult: 1.0 255 | decay_mult: 1.0 256 | } 257 | param { 258 | lr_mult: 2.0 259 | decay_mult: 0.0 260 | } 261 | permutohedral_param { 262 | num_output: 128 263 | neighborhood_size: 2 264 | filter_filler { 265 | type: "gaussian" 266 | std: 0.0010000000475 267 | } 268 | bias_term: true 269 | bias_filler { 270 | type: "constant" 271 | value: 0.0 272 | } 273 | group: 1 274 | norm_type: AFTER 275 | offset_type: NONE 276 | } 277 | } 278 | layer { 279 | name: "bn3" 280 | type: "BatchNorm" 281 | bottom: "conv3" 282 | top: "bn3" 283 | } 284 | layer { 285 | name: "relu3" 286 | type: "ReLU" 287 | bottom: "bn3" 288 | top: "bn3" 289 | } 290 | layer { 291 | name: "bnD3" 292 | type: "Concat" 293 | bottom: "data_feat" 294 | bottom: "bn3" 295 | top: "bnD3" 296 | } 297 | layer { 298 | name: "fconv3" 299 | type: "Convolution" 300 | bottom: "bnD3" 301 | top: "fconv3" 302 | param { 303 | name: "fconv3_w" 304 | lr_mult: 1.0 305 | } 306 | param { 307 | lr_mult: 0.10000000149 308 | } 309 | convolution_param { 310 | num_output: 3 311 | pad: 0 312 | kernel_size: 1 313 | stride: 1 314 | weight_filler { 315 | type: "xavier" 316 | } 317 | bias_filler { 318 | type: "constant" 319 | value: 0.0 320 | } 321 | } 322 | } 323 | 324 | layer { 325 | name: "param_fconv3" 326 | type: "Parameter" 327 | param { 328 | name: "fconv3_w" 329 | lr_mult: 1.0 330 | } 331 | parameter_param { 332 | shape { 333 | dim: 3 # copy the dimension from below 334 | dim: 131 335 | dim: 1 336 | dim: 1 337 | } 338 | } 339 | top: "param_fconv3" 340 | } 341 | layer { 342 | name: "data_lattice2" 343 | type: "Python" 344 | bottom: "fconv3" 345 | bottom: "param_fconv3" 346 | top: "data_lattice2" 347 | python_param { 348 | module: "custom_layers" 349 | layer: "NormAndScale" 350 | param_str: "0*16_1*16_2*16" 351 | } 352 | } 353 | layer { 354 | name: "conv4" 355 | type: "Permutohedral" 356 | bottom: "bn3" 357 | bottom: "data_lattice2" 358 | bottom: "data_lattice2" 359 | top: "conv4" 360 | param { 361 | lr_mult: 1.0 362 | decay_mult: 1.0 363 | } 364 | param { 365 | lr_mult: 2.0 366 | decay_mult: 0.0 367 | } 368 | permutohedral_param { 369 | num_output: 256 370 | neighborhood_size: 2 371 | filter_filler { 372 | type: "gaussian" 373 | std: 0.0010000000475 374 | } 375 | bias_term: true 376 | bias_filler { 377 | type: "constant" 378 | value: 0.0 379 | } 380 | group: 1 381 | norm_type: AFTER 382 | offset_type: NONE 383 | } 384 | } 385 | layer { 386 | name: "bn4" 387 | type: "BatchNorm" 388 | bottom: "conv4" 389 | top: "bn4" 390 | } 391 | layer { 392 | name: "relu4" 393 | type: "ReLU" 394 | bottom: "bn4" 395 | top: "bn4" 396 | } 397 | layer { 398 | name: "bnD4" 399 | type: "Concat" 400 | bottom: "data_feat" 401 | bottom: "bn4" 402 | top: "bnD4" 403 | } 404 | layer { 405 | name: "fconv4" 406 | type: "Convolution" 407 | bottom: "bnD4" 408 | top: "fconv4" 409 | param { 410 | name: "fconv4_w" 411 | lr_mult: 1.0 412 | } 413 | param { 414 | lr_mult: 0.10000000149 415 | } 416 | convolution_param { 417 | num_output: 3 418 | pad: 0 419 | kernel_size: 1 420 | stride: 1 421 | weight_filler { 422 | type: "xavier" 423 | } 424 | bias_filler { 425 | type: "constant" 426 | value: 0.0 427 | } 428 | } 429 | } 430 | 431 | layer { 432 | name: "param_fconv4" 433 | type: "Parameter" 434 | param { 435 | name: "fconv4_w" 436 | lr_mult: 1.0 437 | } 438 | parameter_param { 439 | shape { 440 | dim: 3 # copy the dimension from below 441 | dim: 259 442 | dim: 1 443 | dim: 1 444 | } 445 | } 446 | top: "param_fconv4" 447 | } 448 | layer { 449 | name: "data_lattice3" 450 | type: "Python" 451 | bottom: "fconv4" 452 | bottom: "param_fconv4" 453 | top: "data_lattice3" 454 | python_param { 455 | module: "custom_layers" 456 | layer: "NormAndScale" 457 | param_str: "0*8_1*8_2*8" 458 | } 459 | } 460 | layer { 461 | name: "conv5" 462 | type: "Permutohedral" 463 | bottom: "bn4" 464 | bottom: "data_lattice3" 465 | bottom: "data_lattice3" 466 | top: "conv5" 467 | param { 468 | lr_mult: 1.0 469 | decay_mult: 1.0 470 | } 471 | param { 472 | lr_mult: 2.0 473 | decay_mult: 0.0 474 | } 475 | permutohedral_param { 476 | num_output: 256 477 | neighborhood_size: 2 478 | filter_filler { 479 | type: "gaussian" 480 | std: 0.0010000000475 481 | } 482 | bias_term: true 483 | bias_filler { 484 | type: "constant" 485 | value: 0.0 486 | } 487 | group: 1 488 | norm_type: AFTER 489 | offset_type: NONE 490 | } 491 | } 492 | layer { 493 | name: "bn5" 494 | type: "BatchNorm" 495 | bottom: "conv5" 496 | top: "bn5" 497 | } 498 | layer { 499 | name: "relu5" 500 | type: "ReLU" 501 | bottom: "bn5" 502 | top: "bn5" 503 | } 504 | layer { 505 | name: "bnD5" 506 | type: "Concat" 507 | bottom: "data_feat" 508 | bottom: "bn5" 509 | top: "bnD5" 510 | } 511 | layer { 512 | name: "fconv5" 513 | type: "Convolution" 514 | bottom: "bnD5" 515 | top: "fconv5" 516 | param { 517 | name: "fconv5_w" 518 | lr_mult: 1.0 519 | } 520 | param { 521 | lr_mult: 0.10000000149 522 | } 523 | convolution_param { 524 | num_output: 3 525 | pad: 0 526 | kernel_size: 1 527 | stride: 1 528 | weight_filler { 529 | type: "xavier" 530 | } 531 | bias_filler { 532 | type: "constant" 533 | value: 0.0 534 | } 535 | } 536 | } 537 | 538 | layer { 539 | name: "param_fconv5" 540 | type: "Parameter" 541 | param { 542 | name: "fconv5_w" 543 | lr_mult: 1.0 544 | } 545 | parameter_param { 546 | shape { 547 | dim: 3 # copy the dimension from below 548 | dim: 259 549 | dim: 1 550 | dim: 1 551 | } 552 | } 553 | top: "param_fconv5" 554 | } 555 | layer { 556 | name: "data_lattice4" 557 | type: "Python" 558 | bottom: "fconv5" 559 | bottom: "param_fconv5" 560 | top: "data_lattice4" 561 | python_param { 562 | module: "custom_layers" 563 | layer: "NormAndScale" 564 | param_str: "0*4_1*4_2*4" 565 | } 566 | } 567 | layer { 568 | name: "conv6" 569 | type: "Permutohedral" 570 | bottom: "bn5" 571 | bottom: "data_lattice4" 572 | bottom: "data_lattice4" 573 | top: "conv6" 574 | param { 575 | lr_mult: 1.0 576 | decay_mult: 1.0 577 | } 578 | param { 579 | lr_mult: 2.0 580 | decay_mult: 0.0 581 | } 582 | permutohedral_param { 583 | num_output: 256 584 | neighborhood_size: 2 585 | filter_filler { 586 | type: "gaussian" 587 | std: 0.0010000000475 588 | } 589 | bias_term: true 590 | bias_filler { 591 | type: "constant" 592 | value: 0.0 593 | } 594 | group: 1 595 | norm_type: AFTER 596 | offset_type: NONE 597 | } 598 | } 599 | layer { 600 | name: "bn6" 601 | type: "BatchNorm" 602 | bottom: "conv6" 603 | top: "bn6" 604 | } 605 | layer { 606 | name: "relu6" 607 | type: "ReLU" 608 | bottom: "bn6" 609 | top: "bn6" 610 | } 611 | layer { 612 | name: "concat6" 613 | type: "Concat" 614 | bottom: "bn6" 615 | bottom: "bn2" 616 | bottom: "bn3" 617 | bottom: "bn4" 618 | bottom: "bn5" 619 | top: "concat6" 620 | } 621 | layer { 622 | name: "conv7" 623 | type: "Convolution" 624 | bottom: "concat6" 625 | top: "conv7" 626 | param { 627 | lr_mult: 1.0 628 | } 629 | param { 630 | lr_mult: 0.10000000149 631 | } 632 | convolution_param { 633 | num_output: 128 634 | pad: 0 635 | kernel_size: 1 636 | stride: 1 637 | weight_filler { 638 | type: "xavier" 639 | } 640 | bias_filler { 641 | type: "constant" 642 | value: 0.0 643 | } 644 | } 645 | } 646 | layer { 647 | name: "bn7" 648 | type: "BatchNorm" 649 | bottom: "conv7" 650 | top: "bn7" 651 | } 652 | layer { 653 | name: "relu7" 654 | type: "ReLU" 655 | bottom: "bn7" 656 | top: "bn7" 657 | } 658 | layer { 659 | name: "conv8" 660 | type: "Convolution" 661 | bottom: "bn7" 662 | top: "conv8" 663 | param { 664 | lr_mult: 1.0 665 | } 666 | param { 667 | lr_mult: 0.10000000149 668 | } 669 | convolution_param { 670 | num_output: 4 671 | pad: 0 672 | kernel_size: 1 673 | stride: 1 674 | weight_filler { 675 | type: "xavier" 676 | } 677 | bias_filler { 678 | type: "constant" 679 | value: 0.0 680 | } 681 | } 682 | } 683 | 684 | 685 | layer { 686 | name: "loss" 687 | type: "SoftmaxWithLoss" 688 | bottom: "conv8" 689 | bottom: "label" 690 | top: "loss" 691 | } 692 | layer { 693 | name: "accuracy" 694 | type: "Accuracy" 695 | bottom: "conv8" 696 | bottom: "label" 697 | top: "accuracy" 698 | } 699 | -------------------------------------------------------------------------------- /sampling-based/exp/shapenet3d/mug_net_deform.prototxt: -------------------------------------------------------------------------------- 1 | layer { 2 | name: "data" 3 | type: "Python" 4 | top: "data" 5 | top: "label" 6 | include { 7 | phase: TRAIN 8 | } 9 | python_param { 10 | module: "dataset_shapenet" 11 | layer: "InputShapenet" 12 | param_str: "{\'root\': \'../../data/shapenet_ericyi_ply\', \'jitter_xyz\': 0.01, \'category\': \'03797390\', \'jitter_rotation\': 10.0, \'subset\': \'train\', \'sample_size\': 3000, \'feat_dims\': \'x_y_z\', \'batch_size\': 28}" 13 | } 14 | } 15 | layer { 16 | name: "data" 17 | type: "Python" 18 | top: "data" 19 | top: "label" 20 | include { 21 | phase: TEST 22 | } 23 | python_param { 24 | module: "dataset_shapenet" 25 | layer: "InputShapenet" 26 | param_str: "{\'root\': \'../../data/shapenet_ericyi_ply\', \'jitter_xyz\': 0.0, \'category\': \'03797390\', \'jitter_stretch\': 0.0, \'jitter_rotation\': 0.0, \'subset\': \'val\', \'sample_size\': 3000, \'feat_dims\': \'x_y_z\', \'batch_size\': 28}" 27 | } 28 | } 29 | layer { 30 | name: "data_feat" 31 | type: "Python" 32 | bottom: "data" 33 | top: "data_feat" 34 | python_param { 35 | module: "custom_layers" 36 | layer: "PickAndScale" 37 | param_str: "0_1_2" 38 | } 39 | } 40 | layer { 41 | name: "conv1" 42 | type: "Convolution" 43 | bottom: "data_feat" 44 | top: "conv1" 45 | param { 46 | lr_mult: 1.0 47 | } 48 | param { 49 | lr_mult: 0.10000000149 50 | } 51 | convolution_param { 52 | num_output: 32 53 | pad: 0 54 | kernel_size: 1 55 | stride: 1 56 | weight_filler { 57 | type: "xavier" 58 | } 59 | bias_filler { 60 | type: "constant" 61 | value: 0.0 62 | } 63 | } 64 | } 65 | layer { 66 | name: "bn1" 67 | type: "BatchNorm" 68 | bottom: "conv1" 69 | top: "bn1" 70 | } 71 | layer { 72 | name: "relu1" 73 | type: "ReLU" 74 | bottom: "bn1" 75 | top: "bn1" 76 | } 77 | layer { 78 | name: "bnD1" 79 | type: "Concat" 80 | bottom: "data_feat" 81 | bottom: "bn1" 82 | top: "bnD1" 83 | } 84 | layer { 85 | name: "fconv1" 86 | type: "Convolution" 87 | bottom: "bnD1" 88 | top: "fconv1" 89 | param { 90 | name: "fconv1_w" 91 | lr_mult: 1.0 92 | } 93 | param { 94 | lr_mult: 0.10000000149 95 | } 96 | convolution_param { 97 | num_output: 3 98 | pad: 0 99 | kernel_size: 1 100 | stride: 1 101 | weight_filler { 102 | type: "xavier" 103 | } 104 | bias_filler { 105 | type: "constant" 106 | value: 0.0 107 | } 108 | } 109 | } 110 | 111 | layer { 112 | name: "param_fconv1" 113 | type: "Parameter" 114 | param { 115 | name: "fconv1_w" 116 | lr_mult: 1.0 117 | } 118 | parameter_param { 119 | shape { 120 | dim: 3 # copy the dimension from below 121 | dim: 35 122 | dim: 1 123 | dim: 1 124 | } 125 | } 126 | top: "param_fconv1" 127 | } 128 | layer { 129 | name: "data_lattice0" 130 | type: "Python" 131 | bottom: "fconv1" 132 | bottom: "param_fconv1" 133 | top: "data_lattice0" 134 | python_param { 135 | module: "custom_layers" 136 | layer: "NormAndScale" 137 | param_str: "0*64_1*64_2*64" 138 | } 139 | } 140 | layer { 141 | name: "conv2" 142 | type: "Permutohedral" 143 | bottom: "bn1" 144 | bottom: "data_lattice0" 145 | bottom: "data_lattice0" 146 | top: "conv2" 147 | param { 148 | lr_mult: 1.0 149 | decay_mult: 1.0 150 | } 151 | param { 152 | lr_mult: 2.0 153 | decay_mult: 0.0 154 | } 155 | permutohedral_param { 156 | num_output: 64 157 | neighborhood_size: 2 158 | filter_filler { 159 | type: "gaussian" 160 | std: 0.0010000000475 161 | } 162 | bias_term: true 163 | bias_filler { 164 | type: "constant" 165 | value: 0.0 166 | } 167 | group: 1 168 | norm_type: AFTER 169 | offset_type: NONE 170 | } 171 | } 172 | layer { 173 | name: "bn2" 174 | type: "BatchNorm" 175 | bottom: "conv2" 176 | top: "bn2" 177 | } 178 | layer { 179 | name: "relu2" 180 | type: "ReLU" 181 | bottom: "bn2" 182 | top: "bn2" 183 | } 184 | layer { 185 | name: "bnD2" 186 | type: "Concat" 187 | bottom: "data_feat" 188 | bottom: "bn2" 189 | top: "bnD2" 190 | } 191 | layer { 192 | name: "fconv2" 193 | type: "Convolution" 194 | bottom: "bnD2" 195 | top: "fconv2" 196 | param { 197 | name: "fconv2_w" 198 | lr_mult: 1.0 199 | } 200 | param { 201 | lr_mult: 0.10000000149 202 | } 203 | convolution_param { 204 | num_output: 3 205 | pad: 0 206 | kernel_size: 1 207 | stride: 1 208 | weight_filler { 209 | type: "xavier" 210 | } 211 | bias_filler { 212 | type: "constant" 213 | value: 0.0 214 | } 215 | } 216 | } 217 | layer { 218 | name: "param_fconv2" 219 | type: "Parameter" 220 | param { 221 | name: "fconv2_w" 222 | lr_mult: 1.0 223 | } 224 | parameter_param { 225 | shape { 226 | dim: 3 # copy the dimension from below 227 | dim: 67 228 | dim: 1 229 | dim: 1 230 | } 231 | } 232 | top: "param_fconv2" 233 | } 234 | layer { 235 | name: "data_lattice1" 236 | type: "Python" 237 | bottom: "fconv2" 238 | bottom: "param_fconv2" 239 | top: "data_lattice1" 240 | python_param { 241 | module: "custom_layers" 242 | layer: "NormAndScale" 243 | param_str: "0*32_1*32_2*32" 244 | } 245 | } 246 | layer { 247 | name: "conv3" 248 | type: "Permutohedral" 249 | bottom: "bn2" 250 | bottom: "data_lattice1" 251 | bottom: "data_lattice1" 252 | top: "conv3" 253 | param { 254 | lr_mult: 1.0 255 | decay_mult: 1.0 256 | } 257 | param { 258 | lr_mult: 2.0 259 | decay_mult: 0.0 260 | } 261 | permutohedral_param { 262 | num_output: 128 263 | neighborhood_size: 2 264 | filter_filler { 265 | type: "gaussian" 266 | std: 0.0010000000475 267 | } 268 | bias_term: true 269 | bias_filler { 270 | type: "constant" 271 | value: 0.0 272 | } 273 | group: 1 274 | norm_type: AFTER 275 | offset_type: NONE 276 | } 277 | } 278 | layer { 279 | name: "bn3" 280 | type: "BatchNorm" 281 | bottom: "conv3" 282 | top: "bn3" 283 | } 284 | layer { 285 | name: "relu3" 286 | type: "ReLU" 287 | bottom: "bn3" 288 | top: "bn3" 289 | } 290 | layer { 291 | name: "bnD3" 292 | type: "Concat" 293 | bottom: "data_feat" 294 | bottom: "bn3" 295 | top: "bnD3" 296 | } 297 | layer { 298 | name: "fconv3" 299 | type: "Convolution" 300 | bottom: "bnD3" 301 | top: "fconv3" 302 | param { 303 | name: "fconv3_w" 304 | lr_mult: 1.0 305 | } 306 | param { 307 | lr_mult: 0.10000000149 308 | } 309 | convolution_param { 310 | num_output: 3 311 | pad: 0 312 | kernel_size: 1 313 | stride: 1 314 | weight_filler { 315 | type: "xavier" 316 | } 317 | bias_filler { 318 | type: "constant" 319 | value: 0.0 320 | } 321 | } 322 | } 323 | 324 | layer { 325 | name: "param_fconv3" 326 | type: "Parameter" 327 | param { 328 | name: "fconv3_w" 329 | lr_mult: 1.0 330 | } 331 | parameter_param { 332 | shape { 333 | dim: 3 # copy the dimension from below 334 | dim: 131 335 | dim: 1 336 | dim: 1 337 | } 338 | } 339 | top: "param_fconv3" 340 | } 341 | layer { 342 | name: "data_lattice2" 343 | type: "Python" 344 | bottom: "fconv3" 345 | bottom: "param_fconv3" 346 | top: "data_lattice2" 347 | python_param { 348 | module: "custom_layers" 349 | layer: "NormAndScale" 350 | param_str: "0*16_1*16_2*16" 351 | } 352 | } 353 | layer { 354 | name: "conv4" 355 | type: "Permutohedral" 356 | bottom: "bn3" 357 | bottom: "data_lattice2" 358 | bottom: "data_lattice2" 359 | top: "conv4" 360 | param { 361 | lr_mult: 1.0 362 | decay_mult: 1.0 363 | } 364 | param { 365 | lr_mult: 2.0 366 | decay_mult: 0.0 367 | } 368 | permutohedral_param { 369 | num_output: 256 370 | neighborhood_size: 2 371 | filter_filler { 372 | type: "gaussian" 373 | std: 0.0010000000475 374 | } 375 | bias_term: true 376 | bias_filler { 377 | type: "constant" 378 | value: 0.0 379 | } 380 | group: 1 381 | norm_type: AFTER 382 | offset_type: NONE 383 | } 384 | } 385 | layer { 386 | name: "bn4" 387 | type: "BatchNorm" 388 | bottom: "conv4" 389 | top: "bn4" 390 | } 391 | layer { 392 | name: "relu4" 393 | type: "ReLU" 394 | bottom: "bn4" 395 | top: "bn4" 396 | } 397 | layer { 398 | name: "bnD4" 399 | type: "Concat" 400 | bottom: "data_feat" 401 | bottom: "bn4" 402 | top: "bnD4" 403 | } 404 | layer { 405 | name: "fconv4" 406 | type: "Convolution" 407 | bottom: "bnD4" 408 | top: "fconv4" 409 | param { 410 | name: "fconv4_w" 411 | lr_mult: 1.0 412 | } 413 | param { 414 | lr_mult: 0.10000000149 415 | } 416 | convolution_param { 417 | num_output: 3 418 | pad: 0 419 | kernel_size: 1 420 | stride: 1 421 | weight_filler { 422 | type: "xavier" 423 | } 424 | bias_filler { 425 | type: "constant" 426 | value: 0.0 427 | } 428 | } 429 | } 430 | 431 | layer { 432 | name: "param_fconv4" 433 | type: "Parameter" 434 | param { 435 | name: "fconv4_w" 436 | lr_mult: 1.0 437 | } 438 | parameter_param { 439 | shape { 440 | dim: 3 # copy the dimension from below 441 | dim: 259 442 | dim: 1 443 | dim: 1 444 | } 445 | } 446 | top: "param_fconv4" 447 | } 448 | layer { 449 | name: "data_lattice3" 450 | type: "Python" 451 | bottom: "fconv4" 452 | bottom: "param_fconv4" 453 | top: "data_lattice3" 454 | python_param { 455 | module: "custom_layers" 456 | layer: "NormAndScale" 457 | param_str: "0*8_1*8_2*8" 458 | } 459 | } 460 | layer { 461 | name: "conv5" 462 | type: "Permutohedral" 463 | bottom: "bn4" 464 | bottom: "data_lattice3" 465 | bottom: "data_lattice3" 466 | top: "conv5" 467 | param { 468 | lr_mult: 1.0 469 | decay_mult: 1.0 470 | } 471 | param { 472 | lr_mult: 2.0 473 | decay_mult: 0.0 474 | } 475 | permutohedral_param { 476 | num_output: 256 477 | neighborhood_size: 2 478 | filter_filler { 479 | type: "gaussian" 480 | std: 0.0010000000475 481 | } 482 | bias_term: true 483 | bias_filler { 484 | type: "constant" 485 | value: 0.0 486 | } 487 | group: 1 488 | norm_type: AFTER 489 | offset_type: NONE 490 | } 491 | } 492 | layer { 493 | name: "bn5" 494 | type: "BatchNorm" 495 | bottom: "conv5" 496 | top: "bn5" 497 | } 498 | layer { 499 | name: "relu5" 500 | type: "ReLU" 501 | bottom: "bn5" 502 | top: "bn5" 503 | } 504 | layer { 505 | name: "bnD5" 506 | type: "Concat" 507 | bottom: "data_feat" 508 | bottom: "bn5" 509 | top: "bnD5" 510 | } 511 | layer { 512 | name: "fconv5" 513 | type: "Convolution" 514 | bottom: "bnD5" 515 | top: "fconv5" 516 | param { 517 | name: "fconv5_w" 518 | lr_mult: 1.0 519 | } 520 | param { 521 | lr_mult: 0.10000000149 522 | } 523 | convolution_param { 524 | num_output: 3 525 | pad: 0 526 | kernel_size: 1 527 | stride: 1 528 | weight_filler { 529 | type: "xavier" 530 | } 531 | bias_filler { 532 | type: "constant" 533 | value: 0.0 534 | } 535 | } 536 | } 537 | 538 | layer { 539 | name: "param_fconv5" 540 | type: "Parameter" 541 | param { 542 | name: "fconv5_w" 543 | lr_mult: 1.0 544 | } 545 | parameter_param { 546 | shape { 547 | dim: 3 # copy the dimension from below 548 | dim: 259 549 | dim: 1 550 | dim: 1 551 | } 552 | } 553 | top: "param_fconv5" 554 | } 555 | layer { 556 | name: "data_lattice4" 557 | type: "Python" 558 | bottom: "fconv5" 559 | bottom: "param_fconv5" 560 | top: "data_lattice4" 561 | python_param { 562 | module: "custom_layers" 563 | layer: "NormAndScale" 564 | param_str: "0*4_1*4_2*4" 565 | } 566 | } 567 | layer { 568 | name: "conv6" 569 | type: "Permutohedral" 570 | bottom: "bn5" 571 | bottom: "data_lattice4" 572 | bottom: "data_lattice4" 573 | top: "conv6" 574 | param { 575 | lr_mult: 1.0 576 | decay_mult: 1.0 577 | } 578 | param { 579 | lr_mult: 2.0 580 | decay_mult: 0.0 581 | } 582 | permutohedral_param { 583 | num_output: 256 584 | neighborhood_size: 2 585 | filter_filler { 586 | type: "gaussian" 587 | std: 0.0010000000475 588 | } 589 | bias_term: true 590 | bias_filler { 591 | type: "constant" 592 | value: 0.0 593 | } 594 | group: 1 595 | norm_type: AFTER 596 | offset_type: NONE 597 | } 598 | } 599 | layer { 600 | name: "bn6" 601 | type: "BatchNorm" 602 | bottom: "conv6" 603 | top: "bn6" 604 | } 605 | layer { 606 | name: "relu6" 607 | type: "ReLU" 608 | bottom: "bn6" 609 | top: "bn6" 610 | } 611 | layer { 612 | name: "concat6" 613 | type: "Concat" 614 | bottom: "bn6" 615 | bottom: "bn2" 616 | bottom: "bn3" 617 | bottom: "bn4" 618 | bottom: "bn5" 619 | top: "concat6" 620 | } 621 | layer { 622 | name: "conv7" 623 | type: "Convolution" 624 | bottom: "concat6" 625 | top: "conv7" 626 | param { 627 | lr_mult: 1.0 628 | } 629 | param { 630 | lr_mult: 0.10000000149 631 | } 632 | convolution_param { 633 | num_output: 128 634 | pad: 0 635 | kernel_size: 1 636 | stride: 1 637 | weight_filler { 638 | type: "xavier" 639 | } 640 | bias_filler { 641 | type: "constant" 642 | value: 0.0 643 | } 644 | } 645 | } 646 | layer { 647 | name: "bn7" 648 | type: "BatchNorm" 649 | bottom: "conv7" 650 | top: "bn7" 651 | } 652 | layer { 653 | name: "relu7" 654 | type: "ReLU" 655 | bottom: "bn7" 656 | top: "bn7" 657 | } 658 | layer { 659 | name: "conv8" 660 | type: "Convolution" 661 | bottom: "bn7" 662 | top: "conv8" 663 | param { 664 | lr_mult: 1.0 665 | } 666 | param { 667 | lr_mult: 0.10000000149 668 | } 669 | convolution_param { 670 | num_output: 2 671 | pad: 0 672 | kernel_size: 1 673 | stride: 1 674 | weight_filler { 675 | type: "xavier" 676 | } 677 | bias_filler { 678 | type: "constant" 679 | value: 0.0 680 | } 681 | } 682 | } 683 | 684 | 685 | layer { 686 | name: "loss" 687 | type: "SoftmaxWithLoss" 688 | bottom: "conv8" 689 | bottom: "label" 690 | top: "loss" 691 | } 692 | layer { 693 | name: "accuracy" 694 | type: "Accuracy" 695 | bottom: "conv8" 696 | bottom: "label" 697 | top: "accuracy" 698 | } 699 | --------------------------------------------------------------------------------