박해연

upload modified pcn code

# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import numpy as np
import tensorflow as tf
from tensorpack import dataflow
def resample_pcd(pcd, n):
"""Drop or duplicate points so that pcd has exactly n points"""
idx = np.random.permutation(pcd.shape[0])
if idx.shape[0] < n:
idx = np.concatenate([idx, np.random.randint(pcd.shape[0], size=n-pcd.shape[0])])
return pcd[idx[:n]]
class PreprocessData(dataflow.ProxyDataFlow):
def __init__(self, ds, input_size, output_size):
super(PreprocessData, self).__init__(ds)
self.input_size = input_size
self.output_size = output_size
def get_data(self):
for id, input, gt in self.ds.get_data():
input = resample_pcd(input, self.input_size)
gt = resample_pcd(gt, self.output_size)
yield id, input, gt
class BatchData(dataflow.ProxyDataFlow):
def __init__(self, ds, batch_size, input_size, gt_size, remainder=False, use_list=False):
super(BatchData, self).__init__(ds)
self.batch_size = batch_size
self.input_size = input_size
self.gt_size = gt_size
self.remainder = remainder
self.use_list = use_list
def __len__(self):
ds_size = len(self.ds)
div = ds_size // self.batch_size
rem = ds_size % self.batch_size
if rem == 0:
return div
return div + int(self.remainder)
def __iter__(self):
holder = []
for data in self.ds:
holder.append(data)
if len(holder) == self.batch_size:
yield self._aggregate_batch(holder, self.use_list)
del holder[:]
if self.remainder and len(holder) > 0:
yield self._aggregate_batch(holder, self.use_list)
def _aggregate_batch(self, data_holder, use_list=False):
''' Concatenate input points along the 0-th dimension
Stack all other data along the 0-th dimension
'''
ids = np.stack([x[0] for x in data_holder])
inputs = [resample_pcd(x[1], self.input_size) if x[1].shape[0] > self.input_size else x[1]
for x in data_holder]
inputs = np.expand_dims(np.concatenate([x for x in inputs]), 0).astype(np.float32)
npts = np.stack([x[1].shape[0] if x[1].shape[0] < self.input_size else self.input_size
for x in data_holder]).astype(np.int32)
gts = np.stack([resample_pcd(x[2], self.gt_size) for x in data_holder]).astype(np.float32)
return ids, inputs, npts, gts
def lmdb_dataflow(lmdb_path, batch_size, input_size, output_size, is_training, test_speed=False):
df = dataflow.LMDBSerializer.load(lmdb_path, shuffle=False)
size = df.size()
if is_training:
df = dataflow.LocallyShuffleData(df, buffer_size=2000)
df = dataflow.PrefetchData(df, nr_prefetch=500, nr_proc=1)
df = BatchData(df, batch_size, input_size, output_size)
if is_training:
df = dataflow.PrefetchDataZMQ(df, nr_proc=8)
df = dataflow.RepeatedData(df, -1)
if test_speed:
dataflow.TestDataSpeed(df, size=1000).start()
df.reset_state()
return df, size
def get_queued_data(generator, dtypes, shapes, queue_capacity=10):
assert len(dtypes) == len(shapes), 'dtypes and shapes must have the same length'
queue = tf.FIFOQueue(queue_capacity, dtypes, shapes)
placeholders = [tf.placeholder(dtype, shape) for dtype, shape in zip(dtypes, shapes)]
enqueue_op = queue.enqueue(placeholders)
close_op = queue.close(cancel_pending_enqueues=True)
feed_fn = lambda: {placeholder: value for placeholder, value in zip(placeholders, next(generator))}
queue_runner = tf.contrib.training.FeedingQueueRunner(queue, [enqueue_op], close_op, feed_fns=[feed_fn])
tf.train.add_queue_runner(queue_runner)
return queue.dequeue()
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import numpy as np
#from open3d import *
import open3d as o3d
def read_pcd(filename):
pcd = o3d.io.read_point_cloud(filename)
return np.array(pcd.points)
def save_pcd(filename, points):
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
o3d.io.write_point_cloud(filename, pcd)
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import argparse
import os
from io_util import read_pcd
from tensorpack import DataFlow, dataflow
class pcd_df(DataFlow):
def __init__(self, model_list, num_scans, partial_dir, complete_dir):
self.model_list = model_list
self.num_scans = num_scans
self.partial_dir = partial_dir
self.complete_dir = complete_dir
def size(self):
return len(self.model_list) * self.num_scans
def get_data(self):
for model_id in model_list:
complete = read_pcd(os.path.join(self.complete_dir, '%s.pcd' % model_id))
for i in range(self.num_scans):
partial = read_pcd(os.path.join(self.partial_dir, model_id, '%d.pcd' % i))
yield model_id.replace('/', '_'), partial, complete
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--list_path')
parser.add_argument('--num_scans', type=int)
parser.add_argument('--partial_dir')
parser.add_argument('--complete_dir')
parser.add_argument('--output_path')
args = parser.parse_args()
with open(args.list_path) as file:
model_list = file.read().splitlines()
df = pcd_df(model_list, args.num_scans, args.partial_dir, args.complete_dir)
if os.path.exists(args.output_path):
os.system('rm %s' % args.output_path)
dataflow.LMDBSerializer.save(df, args.output_path)
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from tf_util import mlp, mlp_conv, chamfer, add_train_summary, add_valid_summary
class Model:
def __init__(self, inputs, gt, alpha):
self.num_output_points = 16384
self.features = self.create_encoder(inputs)
self.outputs = self.create_decoder(self.features)
self.loss, self.update = self.create_loss(self.outputs, gt)
self.visualize_ops = [inputs[0], self.outputs[0], gt[0]]
self.visualize_titles = ['input', 'output', 'ground truth']
def create_encoder(self, inputs):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = tf.reduce_max(features, axis=1, keep_dims=True, name='maxpool_0')
features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = tf.reduce_max(features, axis=1, name='maxpool_1')
return features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
outputs = mlp(features, [1024, 1024, self.num_output_points * 3])
outputs = tf.reshape(outputs, [-1, self.num_output_points, 3])
return outputs
def create_loss(self, outputs, gt):
loss = chamfer(outputs, gt)
add_train_summary('train/loss', loss)
update_loss = add_valid_summary('valid/loss', loss)
return loss, update_loss
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from tf_util import mlp_conv, chamfer, add_train_summary, add_valid_summary
class Model:
def __init__(self, inputs, gt, alpha):
self.grid_size = 128
self.grid_scale = 0.5
self.num_output_points = 16384
self.features = self.create_encoder(inputs)
fold1, fold2 = self.create_decoder(self.features)
self.outputs = fold2
self.loss, self.update = self.create_loss(self.outputs, gt)
self.visualize_ops = [inputs[0], fold1[0], fold2[0], gt[0]]
self.visualize_titles = ['input', '1st folding', '2nd folding', 'ground truth']
def create_encoder(self, inputs):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = tf.reduce_max(features, axis=1, keep_dims=True, name='maxpool_0')
features = tf.concat([features, tf.tile(features_global, [1, tf.shape(inputs)[1], 1])], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = tf.reduce_max(features, axis=1, name='maxpool_1')
return features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
x = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
y = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
grid = tf.meshgrid(x, y)
grid = tf.reshape(tf.stack(grid, axis=2), [-1, 2])
grid = tf.tile(tf.expand_dims(grid, 0), [features.shape[0], 1, 1])
features = tf.tile(tf.expand_dims(features, 1), [1, self.num_output_points, 1])
with tf.variable_scope('folding_1'):
fold1 = mlp_conv(tf.concat([features, grid], axis=2), [512, 512, 3])
with tf.variable_scope('folding_2'):
fold2 = mlp_conv(tf.concat([features, fold1], axis=2), [512, 512, 3])
return fold1, fold2
def create_loss(self, outputs, gt):
loss = chamfer(outputs, gt)
add_train_summary('train/loss', loss)
update_loss = add_valid_summary('valid/loss', loss)
return loss, update_loss
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from tf_util import *
class Model:
def __init__(self, inputs, npts, gt, alpha):
self.num_coarse = 1024
self.grid_size = 4
self.grid_scale = 0.05
self.num_fine = self.grid_size ** 2 * self.num_coarse
self.features = self.create_encoder(inputs, npts)
self.coarse, self.fine = self.create_decoder(self.features)
self.loss, self.update = self.create_loss(self.coarse, self.fine, gt, alpha)
self.outputs = self.fine
self.visualize_ops = [tf.split(inputs[0], npts, axis=0), self.coarse, self.fine, gt]
self.visualize_titles = ['input', 'coarse output', 'fine output', 'ground truth']
def create_encoder(self, inputs, npts):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = point_unpool(point_maxpool(features, npts, keepdims=True), npts)
features = tf.concat([features, features_global], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = point_maxpool(features, npts)
return features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
coarse = mlp(features, [1024, 1024, self.num_coarse * 3])
coarse = tf.reshape(coarse, [-1, self.num_coarse, 3])
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
grid = tf.meshgrid(tf.linspace(-0.05, 0.05, self.grid_size), tf.linspace(-0.05, 0.05, self.grid_size))
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]), 0)
grid_feat = tf.tile(grid, [features.shape[0], self.num_coarse, 1])
point_feat = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
point_feat = tf.reshape(point_feat, [-1, self.num_fine, 3])
global_feat = tf.tile(tf.expand_dims(features, 1), [1, self.num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
center = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
center = tf.reshape(center, [-1, self.num_fine, 3])
fine = mlp_conv(feat, [512, 512, 3]) + center
return coarse, fine
def create_loss(self, coarse, fine, gt, alpha):
loss_coarse = chamfer(coarse, gt)
add_train_summary('train/coarse_loss', loss_coarse)
update_coarse = add_valid_summary('valid/coarse_loss', loss_coarse)
loss_fine = chamfer(fine, gt)
add_train_summary('train/fine_loss', loss_fine)
update_fine = add_valid_summary('valid/fine_loss', loss_fine)
loss = loss_coarse + alpha * loss_fine
add_train_summary('train/loss', loss)
update_loss = add_valid_summary('valid/loss', loss)
return loss, [update_coarse, update_fine, update_loss]
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from tf_util import *
class Model:
def __init__(self, inputs,my_inputs, npts, gt, alpha,beta):
self.num_coarse = 1024
self.grid_size = 4
self.grid_scale = 0.05
self.num_fine = self.grid_size ** 2 * self.num_coarse
self.features = self.create_encoder(inputs, npts)
self.coarse, self.fine = self.create_decoder(self.features)
self.loss, self.update = self.create_loss(my_inputs,self.coarse, self.fine, gt, alpha,beta)
self.outputs = self.fine
self.visualize_ops = [tf.split(inputs[0], npts, axis=0), self.coarse, self.fine, gt]
self.visualize_titles = ['input', 'coarse output', 'fine output', 'ground truth']
def create_encoder(self, inputs, npts):
with tf.variable_scope('encoder_0', reuse=tf.AUTO_REUSE):
features = mlp_conv(inputs, [128, 256])
features_global = point_unpool(point_maxpool(features, npts, keepdims=True), npts)
features = tf.concat([features, features_global], axis=2)
with tf.variable_scope('encoder_1', reuse=tf.AUTO_REUSE):
features = mlp_conv(features, [512, 1024])
features = point_maxpool(features, npts)
return features
def create_decoder(self, features):
with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
coarse = mlp(features, [1024, 1024, self.num_coarse * 3])
coarse = tf.reshape(coarse, [-1, self.num_coarse, 3])
with tf.variable_scope('folding', reuse=tf.AUTO_REUSE):
x = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
y = tf.linspace(-self.grid_scale, self.grid_scale, self.grid_size)
grid = tf.meshgrid(x, y)
grid = tf.expand_dims(tf.reshape(tf.stack(grid, axis=2), [-1, 2]), 0)
grid_feat = tf.tile(grid, [features.shape[0], self.num_coarse, 1])
point_feat = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
point_feat = tf.reshape(point_feat, [-1, self.num_fine, 3])
global_feat = tf.tile(tf.expand_dims(features, 1), [1, self.num_fine, 1])
feat = tf.concat([grid_feat, point_feat, global_feat], axis=2)
center = tf.tile(tf.expand_dims(coarse, 2), [1, 1, self.grid_size ** 2, 1])
center = tf.reshape(center, [-1, self.num_fine, 3])
fine = mlp_conv(feat, [512, 512, 3]) + center
return coarse, fine
def create_loss(self,my_inputs,coarse, fine, gt, alpha,beta):
gt_ds = gt[:, :coarse.shape[1], :]
loss_coarse = earth_mover(coarse, gt_ds)
add_train_summary('train/coarse_loss', loss_coarse)
update_coarse = add_valid_summary('valid/coarse_loss', loss_coarse)
loss_fine = chamfer(fine, gt)
add_train_summary('train/fine_loss', loss_fine)
update_fine = add_valid_summary('valid/fine_loss', loss_fine)
loss_input = my_chamfer(my_inputs,fine)
add_train_summary('train/input_loss',loss_input)
update_input = add_valid_summary('valid/input_loss',loss_input)
loss = loss_coarse + alpha * loss_fine + beta*loss_input
add_train_summary('train/loss', loss)
update_loss = add_valid_summary('valid/loss', loss)
return loss, [update_coarse, update_fine, update_input,update_loss]
cuda_inc = /usr/local/cuda-10.0/include/
cuda_lib = /usr/local/cuda-10.0/lib64/
nvcc = /usr/local/cuda-10.0/bin/nvcc
tf_inc = /usr/local/lib/python3.6/dist-packages/tensorflow_core/include
tf_lib = /usr/local/lib/python3.6/dist-packages/tensorflow_core
all: tf_nndistance_so.so tf_approxmatch_so.so
tf_nndistance.cu.o: tf_nndistance.cu
$(nvcc) tf_nndistance.cu -o tf_nndistance.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
tf_nndistance_so.so: tf_nndistance.cpp tf_nndistance.cu.o
g++ tf_nndistance.cpp tf_nndistance.cu.o -o tf_nndistance_so.so \
-I $(cuda_inc) -I $(tf_inc) -L $(cuda_lib) -lcudart -L $(tf_lib) -ltensorflow_framework \
-shared -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -fPIC -O2
tf_approxmatch.cu.o: tf_approxmatch.cu
$(nvcc) tf_approxmatch.cu -o tf_approxmatch.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC
tf_approxmatch_so.so: tf_approxmatch.cpp tf_approxmatch.cu.o
g++ -shared $(CPPFLAGS) tf_approxmatch.cpp tf_approxmatch.cu.o -o tf_approxmatch_so.so \
-I $(cuda_inc) -I $(tf_inc) -L $(cuda_lib) -lcudart -L $(tf_lib) -ltensorflow_framework \
-shared -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -fPIC -O2
clean:
rm -rf *.o *.so
This diff is collapsed. Click to expand it.
__global__ void approxmatch(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,float * __restrict__ match,float * temp){
float * remainL=temp+blockIdx.x*(n+m)*2, * remainR=temp+blockIdx.x*(n+m)*2+n,*ratioL=temp+blockIdx.x*(n+m)*2+n+m,*ratioR=temp+blockIdx.x*(n+m)*2+n+m+n;
float multiL,multiR;
if (n>=m){
multiL=1;
multiR=n/m;
}else{
multiL=m/n;
multiR=1;
}
const int Block=1024;
__shared__ float buf[Block*4];
for (int i=blockIdx.x;i<b;i+=gridDim.x){
for (int j=threadIdx.x;j<n*m;j+=blockDim.x)
match[i*n*m+j]=0;
for (int j=threadIdx.x;j<n;j+=blockDim.x)
remainL[j]=multiL;
for (int j=threadIdx.x;j<m;j+=blockDim.x)
remainR[j]=multiR;
__syncthreads();
for (int j=7;j>=-2;j--){
float level=-powf(4.0f,j);
if (j==-2){
level=0;
}
for (int k0=0;k0<n;k0+=blockDim.x){
int k=k0+threadIdx.x;
float x1=0,y1=0,z1=0;
if (k<n){
x1=xyz1[i*n*3+k*3+0];
y1=xyz1[i*n*3+k*3+1];
z1=xyz1[i*n*3+k*3+2];
}
float suml=1e-9f;
for (int l0=0;l0<m;l0+=Block){
int lend=min(m,l0+Block)-l0;
for (int l=threadIdx.x;l<lend;l+=blockDim.x){
float x2=xyz2[i*m*3+l0*3+l*3+0];
float y2=xyz2[i*m*3+l0*3+l*3+1];
float z2=xyz2[i*m*3+l0*3+l*3+2];
buf[l*4+0]=x2;
buf[l*4+1]=y2;
buf[l*4+2]=z2;
buf[l*4+3]=remainR[l0+l];
}
__syncthreads();
for (int l=0;l<lend;l++){
float x2=buf[l*4+0];
float y2=buf[l*4+1];
float z2=buf[l*4+2];
float d=level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1));
float w=__expf(d)*buf[l*4+3];
suml+=w;
}
__syncthreads();
}
if (k<n)
ratioL[k]=remainL[k]/suml;
}
/*for (int k=threadIdx.x;k<n;k+=gridDim.x){
float x1=xyz1[i*n*3+k*3+0];
float y1=xyz1[i*n*3+k*3+1];
float z1=xyz1[i*n*3+k*3+2];
float suml=1e-9f;
for (int l=0;l<m;l++){
float x2=xyz2[i*m*3+l*3+0];
float y2=xyz2[i*m*3+l*3+1];
float z2=xyz2[i*m*3+l*3+2];
float w=expf(level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1)))*remainR[l];
suml+=w;
}
ratioL[k]=remainL[k]/suml;
}*/
__syncthreads();
for (int l0=0;l0<m;l0+=blockDim.x){
int l=l0+threadIdx.x;
float x2=0,y2=0,z2=0;
if (l<m){
x2=xyz2[i*m*3+l*3+0];
y2=xyz2[i*m*3+l*3+1];
z2=xyz2[i*m*3+l*3+2];
}
float sumr=0;
for (int k0=0;k0<n;k0+=Block){
int kend=min(n,k0+Block)-k0;
for (int k=threadIdx.x;k<kend;k+=blockDim.x){
buf[k*4+0]=xyz1[i*n*3+k0*3+k*3+0];
buf[k*4+1]=xyz1[i*n*3+k0*3+k*3+1];
buf[k*4+2]=xyz1[i*n*3+k0*3+k*3+2];
buf[k*4+3]=ratioL[k0+k];
}
__syncthreads();
for (int k=0;k<kend;k++){
float x1=buf[k*4+0];
float y1=buf[k*4+1];
float z1=buf[k*4+2];
float w=__expf(level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1)))*buf[k*4+3];
sumr+=w;
}
__syncthreads();
}
if (l<m){
sumr*=remainR[l];
float consumption=fminf(remainR[l]/(sumr+1e-9f),1.0f);
ratioR[l]=consumption*remainR[l];
remainR[l]=fmaxf(0.0f,remainR[l]-sumr);
}
}
/*for (int l=threadIdx.x;l<m;l+=blockDim.x){
float x2=xyz2[i*m*3+l*3+0];
float y2=xyz2[i*m*3+l*3+1];
float z2=xyz2[i*m*3+l*3+2];
float sumr=0;
for (int k=0;k<n;k++){
float x1=xyz1[i*n*3+k*3+0];
float y1=xyz1[i*n*3+k*3+1];
float z1=xyz1[i*n*3+k*3+2];
float w=expf(level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1)))*ratioL[k];
sumr+=w;
}
sumr*=remainR[l];
float consumption=fminf(remainR[l]/(sumr+1e-9f),1.0f);
ratioR[l]=consumption*remainR[l];
remainR[l]=fmaxf(0.0f,remainR[l]-sumr);
}*/
__syncthreads();
for (int k0=0;k0<n;k0+=blockDim.x){
int k=k0+threadIdx.x;
float x1=0,y1=0,z1=0;
if (k<n){
x1=xyz1[i*n*3+k*3+0];
y1=xyz1[i*n*3+k*3+1];
z1=xyz1[i*n*3+k*3+2];
}
float suml=0;
for (int l0=0;l0<m;l0+=Block){
int lend=min(m,l0+Block)-l0;
for (int l=threadIdx.x;l<lend;l+=blockDim.x){
buf[l*4+0]=xyz2[i*m*3+l0*3+l*3+0];
buf[l*4+1]=xyz2[i*m*3+l0*3+l*3+1];
buf[l*4+2]=xyz2[i*m*3+l0*3+l*3+2];
buf[l*4+3]=ratioR[l0+l];
}
__syncthreads();
float rl=ratioL[k];
if (k<n){
for (int l=0;l<lend;l++){
float x2=buf[l*4+0];
float y2=buf[l*4+1];
float z2=buf[l*4+2];
float w=__expf(level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1)))*rl*buf[l*4+3];
match[i*n*m+(l0+l)*n+k]+=w;
suml+=w;
}
}
__syncthreads();
}
if (k<n)
remainL[k]=fmaxf(0.0f,remainL[k]-suml);
}
/*for (int k=threadIdx.x;k<n;k+=blockDim.x){
float x1=xyz1[i*n*3+k*3+0];
float y1=xyz1[i*n*3+k*3+1];
float z1=xyz1[i*n*3+k*3+2];
float suml=0;
for (int l=0;l<m;l++){
float x2=xyz2[i*m*3+l*3+0];
float y2=xyz2[i*m*3+l*3+1];
float z2=xyz2[i*m*3+l*3+2];
float w=expf(level*((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1)))*ratioL[k]*ratioR[l];
match[i*n*m+l*n+k]+=w;
suml+=w;
}
remainL[k]=fmaxf(0.0f,remainL[k]-suml);
}*/
__syncthreads();
}
}
}
void approxmatchLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,float * match,float * temp){
approxmatch<<<32,512>>>(b,n,m,xyz1,xyz2,match,temp);
}
__global__ void matchcost(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ out){
__shared__ float allsum[512];
const int Block=1024;
__shared__ float buf[Block*3];
for (int i=blockIdx.x;i<b;i+=gridDim.x){
float subsum=0;
for (int k0=0;k0<n;k0+=blockDim.x){
int k=k0+threadIdx.x;
float x1=0,y1=0,z1=0;
if (k<n){
x1=xyz1[i*n*3+k*3+0];
y1=xyz1[i*n*3+k*3+1];
z1=xyz1[i*n*3+k*3+2];
}
for (int l0=0;l0<m;l0+=Block){
int lend=min(m,l0+Block)-l0;
for (int l=threadIdx.x;l<lend*3;l+=blockDim.x)
buf[l]=xyz2[i*m*3+l0*3+l];
__syncthreads();
if (k<n){
for (int l=0;l<lend;l++){
float x2=buf[l*3+0];
float y2=buf[l*3+1];
float z2=buf[l*3+2];
float d=sqrtf((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1));
subsum+=d*match[i*n*m+(l0+l)*n+k];
}
}
__syncthreads();
}
}
allsum[threadIdx.x]=subsum;
for (int j=1;j<blockDim.x;j<<=1){
__syncthreads();
if ((threadIdx.x&j)==0 && threadIdx.x+j<blockDim.x){
allsum[threadIdx.x]+=allsum[threadIdx.x+j];
}
}
if (threadIdx.x==0)
out[i]=allsum[0];
__syncthreads();
}
}
void matchcostLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * out){
matchcost<<<32,512>>>(b,n,m,xyz1,xyz2,match,out);
}
__global__ void matchcostgrad2(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ grad2){
__shared__ float sum_grad[256*3];
for (int i=blockIdx.x;i<b;i+=gridDim.x){
int kbeg=m*blockIdx.y/gridDim.y;
int kend=m*(blockIdx.y+1)/gridDim.y;
for (int k=kbeg;k<kend;k++){
float x2=xyz2[(i*m+k)*3+0];
float y2=xyz2[(i*m+k)*3+1];
float z2=xyz2[(i*m+k)*3+2];
float subsumx=0,subsumy=0,subsumz=0;
for (int j=threadIdx.x;j<n;j+=blockDim.x){
float x1=x2-xyz1[(i*n+j)*3+0];
float y1=y2-xyz1[(i*n+j)*3+1];
float z1=z2-xyz1[(i*n+j)*3+2];
float d=match[i*n*m+k*n+j]*rsqrtf(fmaxf(x1*x1+y1*y1+z1*z1,1e-20f));
subsumx+=x1*d;
subsumy+=y1*d;
subsumz+=z1*d;
}
sum_grad[threadIdx.x*3+0]=subsumx;
sum_grad[threadIdx.x*3+1]=subsumy;
sum_grad[threadIdx.x*3+2]=subsumz;
for (int j=1;j<blockDim.x;j<<=1){
__syncthreads();
int j1=threadIdx.x;
int j2=threadIdx.x+j;
if ((j1&j)==0 && j2<blockDim.x){
sum_grad[j1*3+0]+=sum_grad[j2*3+0];
sum_grad[j1*3+1]+=sum_grad[j2*3+1];
sum_grad[j1*3+2]+=sum_grad[j2*3+2];
}
}
if (threadIdx.x==0){
grad2[(i*m+k)*3+0]=sum_grad[0];
grad2[(i*m+k)*3+1]=sum_grad[1];
grad2[(i*m+k)*3+2]=sum_grad[2];
}
__syncthreads();
}
}
}
__global__ void matchcostgrad1(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ grad1){
for (int i=blockIdx.x;i<b;i+=gridDim.x){
for (int l=threadIdx.x;l<n;l+=blockDim.x){
float x1=xyz1[i*n*3+l*3+0];
float y1=xyz1[i*n*3+l*3+1];
float z1=xyz1[i*n*3+l*3+2];
float dx=0,dy=0,dz=0;
for (int k=0;k<m;k++){
float x2=xyz2[i*m*3+k*3+0];
float y2=xyz2[i*m*3+k*3+1];
float z2=xyz2[i*m*3+k*3+2];
float d=match[i*n*m+k*n+l]*rsqrtf(fmaxf((x1-x2)*(x1-x2)+(y1-y2)*(y1-y2)+(z1-z2)*(z1-z2),1e-20f));
dx+=(x1-x2)*d;
dy+=(y1-y2)*d;
dz+=(z1-z2)*d;
}
grad1[i*n*3+l*3+0]=dx;
grad1[i*n*3+l*3+1]=dy;
grad1[i*n*3+l*3+2]=dz;
}
}
}
void matchcostgradLauncher(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * grad1,float * grad2){
matchcostgrad1<<<32,512>>>(b,n,m,xyz1,xyz2,match,grad1);
matchcostgrad2<<<dim3(32,32),256>>>(b,n,m,xyz1,xyz2,match,grad2);
}
import tensorflow as tf
from tensorflow.python.framework import ops
import os.path as osp
base_dir = osp.dirname(osp.abspath(__file__))
approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so'))
def approx_match(xyz1,xyz2):
'''
input:
xyz1 : batch_size * #dataset_points * 3
xyz2 : batch_size * #query_points * 3
returns:
match : batch_size * #query_points * #dataset_points
'''
return approxmatch_module.approx_match(xyz1,xyz2)
ops.NoGradient('ApproxMatch')
#@tf.RegisterShape('ApproxMatch')
@ops.RegisterShape('ApproxMatch')
def _approx_match_shape(op):
shape1=op.inputs[0].get_shape().with_rank(3)
shape2=op.inputs[1].get_shape().with_rank(3)
return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])]
def match_cost(xyz1,xyz2,match):
'''
input:
xyz1 : batch_size * #dataset_points * 3
xyz2 : batch_size * #query_points * 3
match : batch_size * #query_points * #dataset_points
returns:
cost : batch_size
'''
return approxmatch_module.match_cost(xyz1,xyz2,match)
#@tf.RegisterShape('MatchCost')
@ops.RegisterShape('MatchCost')
def _match_cost_shape(op):
shape1=op.inputs[0].get_shape().with_rank(3)
shape2=op.inputs[1].get_shape().with_rank(3)
shape3=op.inputs[2].get_shape().with_rank(3)
return [tf.TensorShape([shape1.dims[0]])]
@tf.RegisterGradient('MatchCost')
def _match_cost_grad(op,grad_cost):
xyz1=op.inputs[0]
xyz2=op.inputs[1]
match=op.inputs[2]
grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match)
return [grad_1*tf.expand_dims(tf.expand_dims(grad_cost,1),2),grad_2*tf.expand_dims(tf.expand_dims(grad_cost,1),2),None]
if __name__=='__main__':
alpha=0.5
beta=2.0
import bestmatch
import numpy as np
import math
import random
import cv2
import tf_nndistance
npoint=100
with tf.device('/gpu:2'):
pt_in=tf.placeholder(tf.float32,shape=(1,npoint*4,3))
mypoints=tf.Variable(np.random.randn(1,npoint,3).astype('float32'))
match=approx_match(pt_in,mypoints)
loss=tf.reduce_sum(match_cost(pt_in,mypoints,match))
#match=approx_match(mypoints,pt_in)
#loss=tf.reduce_sum(match_cost(mypoints,pt_in,match))
#distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints)
#loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5
#loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint
optimizer=tf.train.GradientDescentOptimizer(1e-4).minimize(loss)
with tf.Session('') as sess:
sess.run(tf.initialize_all_variables())
while True:
meanloss=0
meantrueloss=0
for i in xrange(1001):
#phi=np.random.rand(4*npoint)*math.pi*2
#tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:]
#tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:]
tpoints=np.hstack([np.linspace(-1,1,400)[:,None],(random.random()*2*np.linspace(1,0,400)**2)[:,None],np.zeros((400,1))])[None,:,:]
trainloss,_=sess.run([loss,optimizer],feed_dict={pt_in:tpoints.astype('float32')})
trainloss,trainmatch=sess.run([loss,match],feed_dict={pt_in:tpoints.astype('float32')})
#trainmatch=trainmatch.transpose((0,2,1))
show=np.zeros((400,400,3),dtype='uint8')^255
trainmypoints=sess.run(mypoints)
for i in xrange(len(tpoints[0])):
u=np.random.choice(range(len(trainmypoints[0])),p=trainmatch[0].T[i])
cv2.line(show,
(int(tpoints[0][i,1]*100+200),int(tpoints[0][i,0]*100+200)),
(int(trainmypoints[0][u,1]*100+200),int(trainmypoints[0][u,0]*100+200)),
cv2.cv.CV_RGB(0,255,0))
for x,y,z in tpoints[0]:
cv2.circle(show,(int(y*100+200),int(x*100+200)),2,cv2.cv.CV_RGB(255,0,0))
for x,y,z in trainmypoints[0]:
cv2.circle(show,(int(y*100+200),int(x*100+200)),3,cv2.cv.CV_RGB(0,0,255))
cost=((tpoints[0][:,None,:]-np.repeat(trainmypoints[0][None,:,:],4,axis=1))**2).sum(axis=2)**0.5
#trueloss=bestmatch.bestmatch(cost)[0]
print(trainloss) #,trueloss
cv2.imshow('show',show)
cmd=cv2.waitKey(10)%256
if cmd==ord('q'):
break
This diff is collapsed. Click to expand it.
#if GOOGLE_CUDA
#define EIGEN_USE_GPU
// #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
const int batch=512;
__shared__ float buf[batch*3];
for (int i=blockIdx.x;i<b;i+=gridDim.x){
for (int k2=0;k2<m;k2+=batch){
int end_k=min(m,k2+batch)-k2;
for (int j=threadIdx.x;j<end_k*3;j+=blockDim.x){
buf[j]=xyz2[(i*m+k2)*3+j];
}
__syncthreads();
for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
float x1=xyz[(i*n+j)*3+0];
float y1=xyz[(i*n+j)*3+1];
float z1=xyz[(i*n+j)*3+2];
int best_i=0;
float best=0;
int end_ka=end_k-(end_k&3);
if (end_ka==batch){
for (int k=0;k<batch;k+=4){
{
float x2=buf[k*3+0]-x1;
float y2=buf[k*3+1]-y1;
float z2=buf[k*3+2]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (k==0 || d<best){
best=d;
best_i=k+k2;
}
}
{
float x2=buf[k*3+3]-x1;
float y2=buf[k*3+4]-y1;
float z2=buf[k*3+5]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+1;
}
}
{
float x2=buf[k*3+6]-x1;
float y2=buf[k*3+7]-y1;
float z2=buf[k*3+8]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+2;
}
}
{
float x2=buf[k*3+9]-x1;
float y2=buf[k*3+10]-y1;
float z2=buf[k*3+11]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+3;
}
}
}
}else{
for (int k=0;k<end_ka;k+=4){
{
float x2=buf[k*3+0]-x1;
float y2=buf[k*3+1]-y1;
float z2=buf[k*3+2]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (k==0 || d<best){
best=d;
best_i=k+k2;
}
}
{
float x2=buf[k*3+3]-x1;
float y2=buf[k*3+4]-y1;
float z2=buf[k*3+5]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+1;
}
}
{
float x2=buf[k*3+6]-x1;
float y2=buf[k*3+7]-y1;
float z2=buf[k*3+8]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+2;
}
}
{
float x2=buf[k*3+9]-x1;
float y2=buf[k*3+10]-y1;
float z2=buf[k*3+11]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (d<best){
best=d;
best_i=k+k2+3;
}
}
}
}
for (int k=end_ka;k<end_k;k++){
float x2=buf[k*3+0]-x1;
float y2=buf[k*3+1]-y1;
float z2=buf[k*3+2]-z1;
float d=x2*x2+y2*y2+z2*z2;
if (k==0 || d<best){
best=d;
best_i=k+k2;
}
}
if (k2==0 || result[(i*n+j)]>best){
result[(i*n+j)]=best;
result_i[(i*n+j)]=best_i;
}
}
__syncthreads();
}
}
}
void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){
NmDistanceKernel<<<dim3(32,16,1),512>>>(b,n,xyz,m,xyz2,result,result_i);
NmDistanceKernel<<<dim3(32,16,1),512>>>(b,m,xyz2,n,xyz,result2,result2_i);
}
__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
for (int i=blockIdx.x;i<b;i+=gridDim.x){
for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
float x1=xyz1[(i*n+j)*3+0];
float y1=xyz1[(i*n+j)*3+1];
float z1=xyz1[(i*n+j)*3+2];
int j2=idx1[i*n+j];
float x2=xyz2[(i*m+j2)*3+0];
float y2=xyz2[(i*m+j2)*3+1];
float z2=xyz2[(i*m+j2)*3+2];
float g=grad_dist1[i*n+j]*2;
atomicAdd(&(grad_xyz1[(i*n+j)*3+0]),g*(x1-x2));
atomicAdd(&(grad_xyz1[(i*n+j)*3+1]),g*(y1-y2));
atomicAdd(&(grad_xyz1[(i*n+j)*3+2]),g*(z1-z2));
atomicAdd(&(grad_xyz2[(i*m+j2)*3+0]),-(g*(x1-x2)));
atomicAdd(&(grad_xyz2[(i*m+j2)*3+1]),-(g*(y1-y2)));
atomicAdd(&(grad_xyz2[(i*m+j2)*3+2]),-(g*(z1-z2)));
}
}
}
void NmDistanceGradKernelLauncher(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2){
cudaMemset(grad_xyz1,0,b*n*3*4);
cudaMemset(grad_xyz2,0,b*m*3*4);
NmDistanceGradKernel<<<dim3(1,16,1),256>>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2);
NmDistanceGradKernel<<<dim3(1,16,1),256>>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1);
}
#endif
import os, sys
import tensorflow as tf
from tensorflow.python.framework import ops
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
nn_distance_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_nndistance_so.so'))
def nn_distance(xyz1,xyz2):
'''
Computes the distance of nearest neighbors for a pair of point clouds
input: xyz1: (batch_size,#points_1,3) the first point cloud
input: xyz2: (batch_size,#points_2,3) the second point cloud
output: dist1: (batch_size,#point_1) distance from first to second
output: idx1: (batch_size,#point_1) nearest neighbor from first to second
output: dist2: (batch_size,#point_2) distance from second to first
output: idx2: (batch_size,#point_2) nearest neighbor from second to first
'''
return nn_distance_module.nn_distance(xyz1,xyz2)
#@tf.RegisterShape('NnDistance')
#def _nn_distance_shape(op):
#shape1=op.inputs[0].get_shape().with_rank(3)
#shape2=op.inputs[1].get_shape().with_rank(3)
#return [tf.TensorShape([shape1.dims[0],shape1.dims[1]]),tf.TensorShape([shape1.dims[0],shape1.dims[1]]),
#tf.TensorShape([shape2.dims[0],shape2.dims[1]]),tf.TensorShape([shape2.dims[0],shape2.dims[1]])]
@ops.RegisterGradient('NnDistance')
def _nn_distance_grad(op,grad_dist1,grad_idx1,grad_dist2,grad_idx2):
xyz1=op.inputs[0]
xyz2=op.inputs[1]
idx1=op.outputs[1]
idx2=op.outputs[3]
return nn_distance_module.nn_distance_grad(xyz1,xyz2,grad_dist1,idx1,grad_dist2,idx2)
if __name__=='__main__':
import numpy as np
import random
import time
from tensorflow.python.ops.gradient_checker import compute_gradient
random.seed(100)
np.random.seed(100)
with tf.Session('') as sess:
xyz1=np.random.randn(32,16384,3).astype('float32')
xyz2=np.random.randn(32,1024,3).astype('float32')
#with tf.device('/gpu:0'):
if True:
inp1=tf.Variable(xyz1)
inp2=tf.constant(xyz2)
reta,retb,retc,retd=nn_distance(inp1,inp2)
loss=tf.reduce_sum(reta)+tf.reduce_sum(retc)
train=tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss)
sess.run(tf.initialize_all_variables())
t0=time.time()
t1=t0
best=1e100
for i in xrange(100):
trainloss,_=sess.run([loss,train])
newt=time.time()
best=min(best,newt-t1)
print(i,trainloss,(newt-t0)/(i+1),best)
t1=newt
#print sess.run([inp1,retb,inp2,retd])
#grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2])
#for i,j in grads:
#print i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j))
#for i in xrange(10):
#t0=time.time()
#a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2})
#print 'time',time.time()-t0
#print a.shape,b.shape,c.shape,d.shape
#print a.dtype,b.dtype,c.dtype,d.dtype
#samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32')
#dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1)
#idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1)
#print np.abs(dist1-a[:,samples]).max()
#print np.abs(idx1-b[:,samples]).max()
#dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1)
#idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1)
#print np.abs(dist2-c[:,samples]).max()
#print np.abs(idx2-d[:,samples]).max()
*.png
*.log
\ No newline at end of file
This directory contains code that generates partial point clouds from ShapeNet models. To use it:
1. Install [Blender](https://blender.org/download/).
2. Create a model list. Each line of the model list should be in the format `[synset_id]/[model_id]`.
3. Run `blender -b -P render_depth.py [ShapeNet directory] [model list] [output directory] [num scans per model]` to render the depth images. The images will be stored in OpenEXR format.
4. Run `python3 process_exr.py [model list] [intrinsics file] [output directory] [num scans per model]` to convert the `.exr` files into 16 bit PNG depth images and point clouds in the model's coordinate frame.
\ No newline at end of file
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import Imath
import OpenEXR
import argparse
import array
import numpy as np
import os
from open3d import *
def read_exr(exr_path, height, width):
file = OpenEXR.InputFile(exr_path)
depth_arr = array.array('f', file.channel('R', Imath.PixelType(Imath.PixelType.FLOAT)))
depth = np.array(depth_arr).reshape((height, width))
depth[depth < 0] = 0
depth[np.isinf(depth)] = 0
return depth
def depth2pcd(depth, intrinsics, pose):
inv_K = np.linalg.inv(intrinsics)
inv_K[2, 2] = -1
depth = np.flipud(depth)
y, x = np.where(depth > 0)
# image coordinates -> camera coordinates
points = np.dot(inv_K, np.stack([x, y, np.ones_like(x)] * depth[y, x], 0))
# camera coordinates -> world coordinates
points = np.dot(pose, np.concatenate([points, np.ones((1, points.shape[1]))], 0)).T[:, :3]
return points
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('list_file')
parser.add_argument('intrinsics_file')
parser.add_argument('output_dir')
parser.add_argument('num_scans', type=int)
args = parser.parse_args()
with open(args.list_file) as file:
model_list = file.read().splitlines()
intrinsics = np.loadtxt(args.intrinsics_file)
width = int(intrinsics[0, 2] * 2)
height = int(intrinsics[1, 2] * 2)
for model_id in model_list:
depth_dir = os.path.join(args.output_dir, 'depth', model_id)
pcd_dir = os.path.join(args.output_dir, 'pcd', model_id)
os.makedirs(depth_dir, exist_ok=True)
os.makedirs(pcd_dir, exist_ok=True)
for i in range(args.num_scans):
exr_path = os.path.join(args.output_dir, 'exr', model_id, '%d.exr' % i)
pose_path = os.path.join(args.output_dir, 'pose', model_id, '%d.txt' % i)
depth = read_exr(exr_path, height, width)
depth_img = Image(np.uint16(depth * 1000))
write_image(os.path.join(depth_dir, '%d.png' % i), depth_img)
pose = np.loadtxt(pose_path)
points = depth2pcd(depth, intrinsics, pose)
pcd = PointCloud()
pcd.points = Vector3dVector(points)
write_point_cloud(os.path.join(pcd_dir, '%d.pcd' % i), pcd)
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import bpy
import mathutils
import numpy as np
import os
import sys
import time
# Usage: blender -b -P render_depth.py [ShapeNet directory] [model list] [output directory] [num scans per model]
def random_pose():
angle_x = np.random.uniform() * 2 * np.pi
angle_y = np.random.uniform() * 2 * np.pi
angle_z = np.random.uniform() * 2 * np.pi
Rx = np.array([[1, 0, 0],
[0, np.cos(angle_x), -np.sin(angle_x)],
[0, np.sin(angle_x), np.cos(angle_x)]])
Ry = np.array([[np.cos(angle_y), 0, np.sin(angle_y)],
[0, 1, 0],
[-np.sin(angle_y), 0, np.cos(angle_y)]])
Rz = np.array([[np.cos(angle_z), -np.sin(angle_z), 0],
[np.sin(angle_z), np.cos(angle_z), 0],
[0, 0, 1]])
R = np.dot(Rz, np.dot(Ry, Rx))
# Set camera pointing to the origin and 1 unit away from the origin
t = np.expand_dims(R[:, 2], 1)
pose = np.concatenate([np.concatenate([R, t], 1), [[0, 0, 0, 1]]], 0)
return pose
def setup_blender(width, height, focal_length):
# camera
camera = bpy.data.objects['Camera']
camera.data.angle = np.arctan(width / 2 / focal_length) * 2
# render layer
scene = bpy.context.scene
scene.render.filepath = 'buffer'
scene.render.image_settings.color_depth = '16'
scene.render.resolution_percentage = 100
scene.render.resolution_x = width
scene.render.resolution_y = height
# compositor nodes
scene.use_nodes = True
tree = scene.node_tree
rl = tree.nodes.new('CompositorNodeRLayers')
output = tree.nodes.new('CompositorNodeOutputFile')
output.base_path = ''
output.format.file_format = 'OPEN_EXR'
tree.links.new(rl.outputs['Depth'], output.inputs[0])
# remove default cube
bpy.data.objects['Cube'].select = True
bpy.ops.object.delete()
return scene, camera, output
if __name__ == '__main__':
model_dir = sys.argv[-4]
list_path = sys.argv[-3]
output_dir = sys.argv[-2]
num_scans = int(sys.argv[-1])
width = 160
height = 120
focal = 100
scene, camera, output = setup_blender(width, height, focal)
intrinsics = np.array([[focal, 0, width / 2], [0, focal, height / 2], [0, 0, 1]])
with open(os.path.join(list_path)) as file:
model_list = [line.strip() for line in file]
open('blender.log', 'w+').close()
os.system('rm -rf %s' % output_dir)
os.makedirs(output_dir)
np.savetxt(os.path.join(output_dir, 'intrinsics.txt'), intrinsics, '%f')
for model_id in model_list:
start = time.time()
exr_dir = os.path.join(output_dir, 'exr', model_id)
pose_dir = os.path.join(output_dir, 'pose', model_id)
os.makedirs(exr_dir)
os.makedirs(pose_dir)
# Redirect output to log file
old_os_out = os.dup(1)
os.close(1)
os.open('blender.log', os.O_WRONLY)
# Import mesh model
model_path = os.path.join(model_dir, model_id, 'model.obj')
bpy.ops.import_scene.obj(filepath=model_path)
# Rotate model by 90 degrees around x-axis (z-up => y-up) to match ShapeNet's coordinates
bpy.ops.transform.rotate(value=-np.pi / 2, axis=(1, 0, 0))
# Render
for i in range(num_scans):
scene.frame_set(i)
pose = random_pose()
camera.matrix_world = mathutils.Matrix(pose)
output.file_slots[0].path = os.path.join(exr_dir, '#.exr')
bpy.ops.render.render(write_still=True)
np.savetxt(os.path.join(pose_dir, '%d.txt' % i), pose, '%f')
# Clean up
bpy.ops.object.delete()
for m in bpy.data.meshes:
bpy.data.meshes.remove(m)
for m in bpy.data.materials:
m.user_clear()
bpy.data.materials.remove(m)
# Show time
os.close(1)
os.dup(old_os_out)
os.close(old_os_out)
print('%s done, time=%.4f sec' % (model_id, time.time() - start))
lmdb>=0.9
msgpack==0.5.6
pyarrow==0.10.0
tensorpack==0.8.9
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import argparse
import csv
import importlib
import models
import numpy as np
import os
import tensorflow as tf
import time
from io_util import read_pcd, save_pcd
from tf_util import chamfer, earth_mover
from visu_util import plot_pcd_three_views
def test(args):
inputs = tf.placeholder(tf.float32, (1, None, 3))
npts = tf.placeholder(tf.int32, (1,))
gt = tf.placeholder(tf.float32, (1, args.num_gt_points, 3))
model_module = importlib.import_module('.%s' % args.model_type, 'models')
model = model_module.Model(inputs, npts, gt, tf.constant(1.0))
output = tf.placeholder(tf.float32, (1, args.num_gt_points, 3))
cd_op = chamfer(output, gt)
emd_op = earth_mover(output, gt)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
sess = tf.Session(config=config)
saver = tf.train.Saver()
saver.restore(sess, args.checkpoint)
os.makedirs(args.results_dir, exist_ok=True)
csv_file = open(os.path.join(args.results_dir, 'results.csv'), 'w')
writer = csv.writer(csv_file)
writer.writerow(['id', 'cd', 'emd'])
with open(args.list_path) as file:
model_list = file.read().splitlines()
total_time = 0
total_cd = 0
total_emd = 0
cd_per_cat = {}
emd_per_cat = {}
for i, model_id in enumerate(model_list):
partial = read_pcd(os.path.join(args.data_dir, 'partial', '%s.pcd' % model_id))
complete = read_pcd(os.path.join(args.data_dir, 'complete', '%s.pcd' % model_id))
start = time.time()
completion = sess.run(model.outputs, feed_dict={inputs: [partial], npts: [partial.shape[0]]})
total_time += time.time() - start
cd, emd = sess.run([cd_op, emd_op], feed_dict={output: completion, gt: [complete]})
total_cd += cd
total_emd += emd
writer.writerow([model_id, cd, emd])
synset_id, model_id = model_id.split('/')
if not cd_per_cat.get(synset_id):
cd_per_cat[synset_id] = []
if not emd_per_cat.get(synset_id):
emd_per_cat[synset_id] = []
cd_per_cat[synset_id].append(cd)
emd_per_cat[synset_id].append(emd)
if i % args.plot_freq == 0:
os.makedirs(os.path.join(args.results_dir, 'plots', synset_id), exist_ok=True)
plot_path = os.path.join(args.results_dir, 'plots', synset_id, '%s.png' % model_id)
plot_pcd_three_views(plot_path, [partial, completion[0], complete],
['input', 'output', 'ground truth'],
'CD %.4f EMD %.4f' % (cd, emd),
[5, 0.5, 0.5])
if args.save_pcd:
os.makedirs(os.path.join(args.results_dir, 'pcds', synset_id), exist_ok=True)
save_pcd(os.path.join(args.results_dir, 'pcds', '%s.pcd' % model_id), completion[0])
csv_file.close()
sess.close()
print('Average time: %f' % (total_time / len(model_list)))
print('Average Chamfer distance: %f' % (total_cd / len(model_list)))
print('Average Earth mover distance: %f' % (total_emd / len(model_list)))
print('Chamfer distance per category')
for synset_id in cd_per_cat.keys():
print(synset_id, '%f' % np.mean(cd_per_cat[synset_id]))
print('Earth mover distance per category')
for synset_id in emd_per_cat.keys():
print(synset_id, '%f' % np.mean(emd_per_cat[synset_id]))
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--list_path', default='data/shapenet/test.list')
parser.add_argument('--data_dir', default='data/shapenet/test')
parser.add_argument('--model_type', default='pcn_emd')
parser.add_argument('--checkpoint', default='data/trained_models/pcn_emd')
parser.add_argument('--results_dir', default='results/shapenet_pcn_emd')
parser.add_argument('--num_gt_points', type=int, default=16384)
parser.add_argument('--plot_freq', type=int, default=100)
parser.add_argument('--save_pcd', action='store_true')
args = parser.parse_args()
test(args)
This diff is collapsed. Click to expand it.
{
"cells": [
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[32m[0603 12:20:31 @format.py:92]\u001b[0m Found 100 entries in ../../../data/shapenet-car/valid.lmdb\n",
"INFO:tensorflow:Restoring parameters from ./log/pcn_addbeta_lr/model-81500\n",
"Average Chamfer distance: 0.009029\n",
"Average Earth mover distance: 0.053403\n"
]
}
],
"source": [
"# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018\n",
"\n",
"import argparse\n",
"import csv\n",
"import importlib\n",
"import models\n",
"import numpy as np\n",
"import os\n",
"import tensorflow as tf\n",
"import time\n",
"\n",
"from tf_util import chamfer, earth_mover\n",
"from visu_util import plot_pcd_three_views\n",
"from data_util import lmdb_dataflow, get_queued_data\n",
"\n",
"##################\n",
"#from io_util import read_pcd, save_pcd 이부분 그냥 포함시켜버렸다...ㅎㅎ\n",
"import numpy as np\n",
"#from open3d import *\n",
"import open3d as o3d\n",
"\n",
"\n",
"def read_pcd(filename):\n",
" pcd = o3d.io.read_point_cloud(filename)\n",
" return np.array(pcd.points)\n",
"\n",
"\n",
"def save_pcd(filename, points):\n",
" pcd = o3d.geometry.PointCloud()\n",
" pcd.points = o3d.utility.Vector3dVector(points)\n",
" o3d.io.write_point_cloud(filename, pcd)\n",
"##################\n",
"\n",
"#list_path ='../../data/shapenet/car_test.list' \n",
"#data_dir = '../../data/shapenet/test'\n",
"model_type = 'pcn_emd'\n",
"checkpoint = './log/pcn_addbeta_lr'\n",
"results_dir ='results/pcn_addbeta_lr'\n",
"num_gt_points = 16384\n",
"plot_freq = 1\n",
"_save_pcd = True\n",
"lmdb_valid='../../../data/shapenet-car/valid.lmdb'\n",
"num_input_points=3000\n",
"\n",
"def test():\n",
" inputs = tf.placeholder(tf.float32, (1, None, 3))\n",
" my_inputs = tf.placeholder(tf.float32, (1, None, 3))\n",
" npts = tf.placeholder(tf.int32, (1,))\n",
" gt = tf.placeholder(tf.float32, (1, num_gt_points, 3))\n",
" model_module = importlib.import_module('.%s' % model_type, 'models')\n",
" model = model_module.Model(inputs,my_inputs, npts, gt, tf.constant(1.0),tf.constant(1.0))\n",
"\n",
" output = tf.placeholder(tf.float32, (1, num_gt_points, 3))\n",
" cd_op = chamfer(output, gt)\n",
" emd_op = earth_mover(output, gt)\n",
"\n",
" ###\n",
" df_valid, num_valid = lmdb_dataflow(\n",
" lmdb_valid, 1, num_input_points, num_gt_points, is_training=False)\n",
" valid_gen = df_valid.get_data()\n",
" \n",
" config = tf.ConfigProto()\n",
" config.gpu_options.allow_growth = True\n",
" config.allow_soft_placement = True\n",
" sess = tf.Session(config=config)\n",
"\n",
" saver = tf.train.Saver()\n",
" saver.restore(sess, tf.train.latest_checkpoint(checkpoint))\n",
"\n",
" os.makedirs(results_dir, exist_ok=True)\n",
" csv_file = open(os.path.join(results_dir, 'results.csv'), 'w') # 각 항목별로 cd, emd 구해줌.\n",
" writer = csv.writer(csv_file)\n",
" writer.writerow(['id', 'cd', 'emd'])\n",
"\n",
" ###\n",
" total_time = 0\n",
" total_cd = 0\n",
" total_emd = 0\n",
" for i in range(num_valid):\n",
" ids,iinputs,inpts,igt = next(valid_gen)\n",
"\n",
" completion = sess.run(model.outputs, feed_dict={inputs:iinputs, my_inputs:iinputs, npts:inpts})\n",
" cd,emd = sess.run([cd_op,emd_op],feed_dict={output: completion, gt:igt})\n",
" total_cd +=cd\n",
" total_emd +=emd\n",
" writer.writerow([ids,cd,emd]) #항목별 cd,emd\n",
" \n",
" if i % plot_freq == 0:\n",
" os.makedirs(os.path.join(results_dir, 'plots'), exist_ok=True)\n",
" plot_path = os.path.join(results_dir, 'plots', '%s.png' % ids)\n",
"# print(iinputs.shape,completion[0].shape,igt.shape)###\n",
"\n",
" plot_pcd_three_views(plot_path, [iinputs.reshape(iinputs.shape[1],iinputs.shape[2]), completion[0], igt.reshape(igt.shape[1],igt.shape[2])],\n",
" ['input', 'output', 'ground truth'],\n",
" 'CD %.4f EMD %.4f' % (cd, emd),\n",
" [5, 0.5, 0.5])\n",
" if _save_pcd:\n",
" os.makedirs(os.path.join(results_dir, 'pcds'), exist_ok=True)\n",
" save_pcd(os.path.join(results_dir, 'pcds', '%s.pcd' % ids), completion[0])\n",
" csv_file.close()\n",
" sess.close()\n",
"\n",
" print('Average Chamfer distance: %f' % (total_cd / num_valid))\n",
" print('Average Earth mover distance: %f' % (total_emd / num_valid))\n",
"######################\n",
"''' \n",
" with open(list_path) as file:\n",
" model_list = file.read().splitlines()\n",
" total_time = 0\n",
" total_cd = 0\n",
" total_emd = 0\n",
" cd_per_cat = {}\n",
" emd_per_cat = {}\n",
" for i, model_id in enumerate(model_list):\n",
" partial = read_pcd(os.path.join(data_dir, 'partial', '%s.pcd' % model_id))\n",
" complete = read_pcd(os.path.join(data_dir, 'complete', '%s.pcd' % model_id))\n",
" start = time.time()\n",
" completion = sess.run(model.outputs, feed_dict={inputs: [partial],my_inputs:[partial], npts: [partial.shape[0]]})\n",
" total_time += time.time() - start\n",
" cd, emd = sess.run([cd_op, emd_op], feed_dict={output: completion, gt: [complete]})\n",
" total_cd += cd\n",
" total_emd += emd\n",
" writer.writerow([model_id, cd, emd]) #항목별 cd,emd 써줌\n",
"\n",
" # 카테고리별 cd,emd 얻음\n",
" synset_id, model_id = model_id.split('/')\n",
" if not cd_per_cat.get(synset_id):\n",
" cd_per_cat[synset_id] = []\n",
" if not emd_per_cat.get(synset_id):\n",
" emd_per_cat[synset_id] = []\n",
" cd_per_cat[synset_id].append(cd)\n",
" emd_per_cat[synset_id].append(emd)\n",
" \n",
" # 3가지 view에서 모델 input,gt,output보여줌.\n",
" if i % plot_freq == 0:\n",
" os.makedirs(os.path.join(results_dir, 'plots', synset_id), exist_ok=True)\n",
" plot_path = os.path.join(results_dir, 'plots', synset_id, '%s.png' % model_id)\n",
" plot_pcd_three_views(plot_path, [partial, completion[0], complete],\n",
" ['input', 'output', 'ground truth'],\n",
" 'CD %.4f EMD %.4f' % (cd, emd),\n",
" [5, 0.5, 0.5])\n",
" if _save_pcd:\n",
" os.makedirs(os.path.join(results_dir, 'pcds', synset_id), exist_ok=True)\n",
" save_pcd(os.path.join(results_dir, 'pcds', '%s.pcd' % model_id), completion[0])\n",
" csv_file.close()\n",
" sess.close()\n",
"\n",
" print('Average time: %f' % (total_time / len(model_list)))\n",
" print('Average Chamfer distance: %f' % (total_cd / len(model_list)))\n",
" print('Average Earth mover distance: %f' % (total_emd / len(model_list)))\n",
" print('Chamfer distance per category')\n",
" for synset_id in cd_per_cat.keys():\n",
" print(synset_id, '%f' % np.mean(cd_per_cat[synset_id]))\n",
" print('Earth mover distance per category')\n",
" for synset_id in emd_per_cat.keys():\n",
" print(synset_id, '%f' % np.mean(emd_per_cat[synset_id]))\n",
" '''\n",
"'''\n",
"\n",
"if __name__ == '__main__':\n",
" parser = argparse.ArgumentParser()\n",
" parser.add_argument('--list_path', default='data/shapenet/test.list')\n",
" parser.add_argument('--data_dir', default='data/shapenet/test')\n",
" parser.add_argument('--model_type', default='pcn_emd')\n",
" parser.add_argument('--checkpoint', default='data/trained_models/pcn_emd')\n",
" parser.add_argument('--results_dir', default='results/shapenet_pcn_emd')\n",
" parser.add_argument('--num_gt_points', type=int, default=16384)\n",
" parser.add_argument('--plot_freq', type=int, default=100)\n",
" parser.add_argument('--save_pcd', action='store_true')\n",
" args = parser.parse_args()\n",
"\n",
" test(args)\n",
"'''\n",
"test()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.9"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
import tensorflow as tf
from pc_distance import tf_nndistance, tf_approxmatch
def mlp(features, layer_dims, bn=None, bn_params=None):
for i, num_outputs in enumerate(layer_dims[:-1]):
features = tf.contrib.layers.fully_connected(
features, num_outputs,
normalizer_fn=bn,
normalizer_params=bn_params,
scope='fc_%d' % i)
outputs = tf.contrib.layers.fully_connected(
features, layer_dims[-1],
activation_fn=None,
scope='fc_%d' % (len(layer_dims) - 1))
return outputs
def mlp_conv(inputs, layer_dims, bn=None, bn_params=None):
for i, num_out_channel in enumerate(layer_dims[:-1]):
inputs = tf.contrib.layers.conv1d(
inputs, num_out_channel,
kernel_size=1,
normalizer_fn=bn,
normalizer_params=bn_params,
scope='conv_%d' % i)
outputs = tf.contrib.layers.conv1d(
inputs, layer_dims[-1],
kernel_size=1,
activation_fn=None,
scope='conv_%d' % (len(layer_dims) - 1))
return outputs
def point_maxpool(inputs, npts, keepdims=False):
outputs = [tf.reduce_max(f, axis=1, keepdims=keepdims)
for f in tf.split(inputs, npts, axis=1)]
return tf.concat(outputs, axis=0)
def point_unpool(inputs, npts):
inputs = tf.split(inputs, inputs.shape[0], axis=0)
outputs = [tf.tile(f, [1, npts[i], 1]) for i,f in enumerate(inputs)]
return tf.concat(outputs, axis=1)
def chamfer(pcd1, pcd2):
dist1, _, dist2, _ = tf_nndistance.nn_distance(pcd1, pcd2)
dist1 = tf.reduce_mean(tf.sqrt(dist1))
dist2 = tf.reduce_mean(tf.sqrt(dist2))
return (dist1 + dist2) / 2
def my_chamfer(pcd1, pcd2):
dist1, _, dist2, _ = tf_nndistance.nn_distance(pcd1, pcd2)
dist1 = tf.reduce_mean(tf.sqrt(dist1))
#dist2 = tf.reduce_mean(tf.sqrt(dist2))
return dist1
def earth_mover(pcd1, pcd2):
assert pcd1.shape[1] == pcd2.shape[1]
num_points = tf.cast(pcd1.shape[1], tf.float32)
match = tf_approxmatch.approx_match(pcd1, pcd2)
cost = tf_approxmatch.match_cost(pcd1, pcd2, match)
return tf.reduce_mean(cost / num_points)
def add_train_summary(name, value):
tf.summary.scalar(name, value, collections=['train_summary'])
def add_valid_summary(name, value):
avg, update = tf.metrics.mean(value)
tf.summary.scalar(name, avg, collections=['valid_summary'])
return update
# Author: Wentao Yuan (wyuan1@cs.cmu.edu) 05/31/2018
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def plot_pcd_three_views(filename, pcds, titles, suptitle='', sizes=None, cmap='Reds', zdir='y',
xlim=(-0.3, 0.3), ylim=(-0.3, 0.3), zlim=(-0.3, 0.3)):
if sizes is None:
sizes = [0.5 for i in range(len(pcds))]
fig = plt.figure(figsize=(len(pcds) * 3, 9))
for i in range(3):
elev = 30
azim = -45 + 90 * i
for j, (pcd, size) in enumerate(zip(pcds, sizes)):
color = pcd[:, 0]
ax = fig.add_subplot(3, len(pcds), i * len(pcds) + j + 1, projection='3d')
ax.view_init(elev, azim)
ax.scatter(pcd[:, 0], pcd[:, 1], pcd[:, 2], zdir=zdir, c=color, s=size, cmap=cmap, vmin=-1, vmax=0.5)
ax.set_title(titles[j])
ax.set_axis_off()
ax.set_xlim(xlim)
ax.set_ylim(ylim)
ax.set_zlim(zlim)
plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.9, wspace=0.1, hspace=0.1)
plt.suptitle(suptitle)
fig.savefig(filename)
plt.close(fig)