├── README.md ├── data_generate ├── _data_gen.py ├── _load_point.py ├── data_zip.py ├── gen_pix3d.py ├── obj_io.py ├── obj_io.pyc ├── render_all.py ├── render_single.py ├── rendering_metadata.txt ├── resample.py ├── sample_all.py ├── sample_single.py ├── sample_single.pyc ├── test.py ├── tf_sampling.cpp ├── tf_sampling.py ├── tf_sampling.pyc ├── tf_sampling_compile.sh ├── tf_sampling_g.cu ├── tf_sampling_g.cu.o └── tf_sampling_so.so ├── images ├── overview_network.jpg ├── qualitative_result_pix3d.jpg └── qualitative_result_shapenet.jpg ├── metric ├── __init__.pyc ├── approxmatch.cpp ├── approxmatch.cu ├── makefile ├── tf_approxmatch.cpp ├── tf_approxmatch.py ├── tf_approxmatch.pyc ├── tf_approxmatch_g.cu ├── tf_approxmatch_g.cu.o ├── tf_approxmatch_so.so ├── tf_nndistance.cpp ├── tf_nndistance.py ├── tf_nndistance.pyc ├── tf_nndistance_g.cu ├── tf_nndistance_g.cu.o └── tf_nndistance_so.so ├── model ├── dense.py ├── layer │ ├── attention_module.py │ ├── basic_layers.py │ ├── encoders_decoders.py │ ├── pointnet_util.py │ ├── tf_ops │ │ ├── __init__.py │ │ ├── __init__.pyc │ │ ├── cd │ │ │ ├── __init__.py │ │ │ ├── __init__.pyc │ │ │ ├── tf_nndistance.cpp │ │ │ ├── tf_nndistance.py │ │ │ ├── tf_nndistance.pyc │ │ │ ├── tf_nndistance_g.cu │ │ │ ├── tf_nndistance_g.cu.o │ │ │ └── tf_nndistance_so.so │ │ ├── emd │ │ │ ├── __init__.py │ │ │ ├── __init__.pyc │ │ │ ├── tf_approxmatch.cpp │ │ │ ├── tf_approxmatch.py │ │ │ ├── tf_approxmatch.pyc │ │ │ ├── tf_approxmatch_g.cu │ │ │ ├── tf_approxmatch_g.cu.o │ │ │ └── tf_approxmatch_so.so │ │ ├── grouping │ │ │ ├── __init__.py │ │ │ ├── __init__.pyc │ │ │ ├── query_ball_point.cpp │ │ │ ├── query_ball_point.cu │ │ │ ├── query_ball_point_block.cu │ │ │ ├── query_ball_point_grid.cu │ │ │ ├── selection_sort.cpp │ │ │ ├── selection_sort.cu │ │ │ ├── selection_sort_const.cu │ │ │ ├── test_knn.py │ │ │ ├── tf_grouping.cpp │ │ │ ├── tf_grouping.py │ │ │ ├── tf_grouping.pyc │ │ │ ├── tf_grouping_compile.sh │ │ │ ├── tf_grouping_g.cu │ │ │ ├── tf_grouping_g.cu.o │ │ │ ├── tf_grouping_op_test.py │ │ │ └── tf_grouping_so.so │ │ ├── interpolating │ │ │ ├── __init__.py │ │ │ ├── __init__.pyc │ │ │ ├── interpolate.cpp │ │ │ ├── tf_interpolate.cpp │ │ │ ├── tf_interpolate.py │ │ │ ├── tf_interpolate.pyc │ │ │ ├── tf_interpolate_compile.sh │ │ │ ├── tf_interpolate_op_test.py │ │ │ ├── tf_interpolate_so.so │ │ │ └── visu_interpolation.py │ │ └── sampling │ │ │ ├── __init__.py │ │ │ ├── __init__.pyc │ │ │ ├── tf_sampling.cpp │ │ │ ├── tf_sampling.py │ │ │ ├── tf_sampling.pyc │ │ │ ├── tf_sampling_compile.sh │ │ │ ├── tf_sampling_g.cu │ │ │ ├── tf_sampling_g.cu.o │ │ │ └── tf_sampling_so.so │ ├── tf_util.py │ └── tf_utils.py ├── layers.py ├── model.py ├── pu_util.py └── sparse.py ├── test.py ├── train_dense.py ├── train_finetune.py ├── train_sparse.py └── utils ├── datafetcher.py ├── datafetcher_dense.py ├── datafetcher_finetune.py ├── datafetcher_sparse.py ├── render_balls_so.cpp ├── render_balls_so.so ├── show3d_balls.py └── visualize_utils.py /README.md: -------------------------------------------------------------------------------- 1 | # AttentionDPCR 2 | This repository contains the source codes for the paper "Attention-Based Dense Point Cloud Reconstruction from a Single Image" 3 | 4 | # Overview 5 | 3D Reconstruction has drawn much attention in computer vision. Generating a dense point cloud from a single image is a more challenging task. However, generating dense point clouds directly costs expensively in calculation and memory and may cause the network hard to train. In this work, we propose a two-stage training dense point cloud generation network. We first train our attention-based sparse point cloud generation network to generate a sparse point cloud from a single image. Then we train our dense point cloud generation network to densify the generated sparse point cloud. After combining the two stages and finetuning, we obtain an end-to-end network that generates a dense point cloud from a single image. Through evaluation of both synthetic and real-world datasets, we demonstrate that our approach outperforms state of the art works in dense point cloud generation. 6 | ![overview](https://github.com/VIM-Lab/AttentionDPCR/blob/master/images/overview_network.jpg) 7 | 8 | # Dataset 9 | We train and validate our model on the ShapeNet dataset. The code of data preprocessing is released and the guideline is coming soon. 10 | 11 | # Setup 12 | Before get started, you may need to install numpy, tensorflow-gpu-1.4, python-opencv etc. by pip or sudo apt-get. Then you need to complie some CUDA code. 13 | ``` 14 | $ cd metric 15 | $ make 16 | 17 | $ cd ..\model\layer\tf_ops 18 | $ cd grouping 19 | $ bash tf_grouping_compile.sh 20 | $ cd ..\interpolating 21 | $ bash tf_interpolating_compile.sh 22 | $ cd ..\sampling 23 | $ bash tf_sampling_compile.sh 24 | 25 | ``` 26 | 27 | # Training 28 | Our network is a two-stage training network. 29 | * To train sparse point cloud generation network, run: 30 | ``` 31 | python train_sparse.py sparse 32 | ``` 33 | * To train dense point cloud generation network, run: 34 | ``` 35 | python train_dense.py dense 36 | ``` 37 | * To finetune two stages, run: 38 | ``` 39 | python train_finetune.py finetune 40 | ``` 41 | 42 | # Visualization 43 | **ShapeNet**
44 | Below are a few sample reconstructions from our trained model tested on ShapeNet. 45 | ![shapenet](https://github.com/VIM-Lab/AttentionDPCR/blob/master/images/qualitative_result_shapenet.jpg) 46 | 47 | **Pix3D**
48 | Below are a few sample reconstructions from our trained model tested on real-world Pix3D dataset. Note that we mask out the background using the provided masks before passing the images through the network. 49 | ![pix3d](https://github.com/VIM-Lab/AttentionDPCR/blob/master/images/qualitative_result_pix3d.jpg) 50 | -------------------------------------------------------------------------------- /data_generate/_data_gen.py: -------------------------------------------------------------------------------- 1 | import cPickle as pickle 2 | import numpy as np 3 | from filecmp import dircmp 4 | import os 5 | import cv2 6 | from progress.bar import IncrementalBar 7 | import time 8 | import h5py 9 | 10 | shapenet_category_to_id = { 11 | 'airplane' : '02691156', 12 | 'bench' : '02828884', 13 | 'cabinet' : '02933112', 14 | 'car' : '02958343', 15 | 'chair' : '03001627', 16 | 'lamp' : '03636649', 17 | 'monitor' : '03211117', 18 | 'rifle' : '04090263', 19 | 'sofa' : '04256520', 20 | 'speaker' : '03691459', 21 | 'table' : '04379243', 22 | 'telephone' : '04401088', 23 | 'vessel' : '04530566' 24 | } 25 | 26 | def get_path(): 27 | image_prefix = '/media/tree/data1/projects/AttentionBased/data/image_png' 28 | point_prefix = '/media/tree/data1/projects/AttentionBased/data/pointcloud/16384' 29 | image_output_folder = '/media/tree/backup/projects/AttentionBased/data/image' 30 | point_output_folder = '/media/tree/backup/projects/AttentionBased/data/pointcloud_16384' 31 | npz_output_folder = '/media/tree/backup/projects/AttentionBased/data/data' 32 | 33 | image_path = [] 34 | point_path = [] 35 | # common_path = [] 36 | print('preparing data') 37 | for dir_top, subdir_cmps in dircmp(image_prefix,point_prefix).subdirs.items(): 38 | # print(dir_top, subdir_cmps) 39 | for dir_bottom in subdir_cmps.common_dirs: 40 | # print('preparing') 41 | # common_path.append(os.path.join(dir_top, dir_bottom) 42 | for index in range(24): 43 | image_path.append(os.path.join(image_prefix, dir_top, dir_bottom, '{0:02d}.png'.format(int(index)))) 44 | point_path.append(os.path.join(point_prefix, dir_top, dir_bottom, '{0:02d}.npy'.format(int(index)))) 45 | print('data prepared') 46 | # print(image_path[0]) 47 | # print(image_path[0].split('/')[-3]) 48 | return image_path, point_path, image_output_folder, point_output_folder, npz_output_folder 49 | 50 | def dump_data(): 51 | image_path, point_path, image_output_folder, point_output_folder, npz_output_folder = get_path() 52 | bar = IncrementalBar(max=len(image_path)) 53 | count = 0 54 | batch_image = np.empty((24,256,256,3), dtype = np.uint8) 55 | batch_point = np.empty((24,16384,3), dtype = np.float32) 56 | for i in range(24): 57 | # for i in range(len(image_path)): 58 | image = cv2.imread(image_path[i]) 59 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 60 | image_array = np.array(image) 61 | image_array = image_array.astype(np.uint8) 62 | point_array = np.load(point_path[i]) 63 | point_array = point_array.astype(np.float32) 64 | # np.save(os.path.join(image_output_folder,"{}_{}_{}.npy".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2], image_path[i].split('/')[-1].split('.')[0])), image_array) 65 | # np.save(os.path.join(point_output_folder,"{}_{}_{}.npy".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2], image_path[i].split('/')[-1].split('.')[0])), point_array) 66 | 67 | batch_image[count, ...] = image_array 68 | batch_point[count, ...] = point_array 69 | count += 1 70 | 71 | while count == 24: 72 | count = 0 73 | np.save(os.path.join(image_output_folder,"{}_{}.npy".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2])), batch_image) 74 | np.save(os.path.join(point_output_folder,"{}_{}.npy".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2])), batch_point) 75 | np.savez_compressed(os.path.join(npz_output_folder,"{}_{}.npz".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2])), image = batch_image, point = batch_point) 76 | # print() 77 | # print(batch_image.shape) 78 | # print(batch_image.dtype) 79 | # print(batch_image.size) 80 | # print(batch_point.shape) 81 | # print(batch_point.dtype) 82 | # print(batch_point.size) 83 | # load = np.load(os.path.join(npz_output_folder,"{}_{}.npz".format(image_path[i].split('/')[-3], image_path[i].split('/')[-2]))) 84 | # load_image = load['image'] 85 | # load_point = load['point'] 86 | # print(load_image.shape) 87 | # print(load_image.dtype) 88 | # print(load_image.size) 89 | # print(load_point.shape) 90 | # print(load_point.dtype) 91 | # print(load_point.size) 92 | batch_image = np.empty((24,256,256,3), dtype = np.uint8) 93 | batch_point = np.empty((24,16384,3), dtype = np.float32) 94 | 95 | # with open(os.path.join(output_folder,'{}_{}_{}.pkl'.format(image_path[i].split('/')[-3], image_path[i].split('/')[-2], image_path[i].split('/')[-1].split('.')[0])),'w') as f: 96 | # pickle.dump([image_array,point_array], f) 97 | bar.next() 98 | bar.finish() 99 | 100 | def dump_single(): 101 | image_path = '/media/tree/data1/projects/AttentionBased/data/image_png/04090263/2d203e283c13fd16494585aaf374e961/00.png' 102 | point_path = '/media/tree/data1/projects/AttentionBased/data/pointcloud/16384/04090263/2d203e283c13fd16494585aaf374e961/00.npy' 103 | output_folder = '/media/tree/data1/projects/AttentionBased/PSGN' 104 | pixel2mesh_path = '/media/tree/data1/projects/Pixel2Mesh/pixel2mesh/data/ShapeNetTrain/04090263_2d203e283c13fd16494585aaf374e961_00.dat' 105 | image = cv2.imread(image_path) 106 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 107 | image_array = np.array(image) 108 | print(image_array.dtype) 109 | # image_array = image.astype('float32') / 255.0 110 | point_array = np.load(point_path) 111 | # print(point_array.shape) 112 | # np.savez(os.path.join(output_folder,"{}_{}_{}.npz".format(image_path.split('/')[-3], image_path.split('/')[-2], image_path.split('/')[-1].split('.')[0])), image = image_array, point = point_array) 113 | 114 | # pkl = pickle.load(open(pixel2mesh_path, 'rb')) 115 | # pixel_image = pkl[0] 116 | # pixel_label = pkl[1] 117 | # print(pixel_label.shape) 118 | 119 | # np.savez(os.path.join(output_folder,"pixel.npz"), image = pixel_image, point = pixel_label) 120 | 121 | # with open(os.path.join(output_folder,"{}_{}_{}.pkl".format(image_path.split('/')[-3], image_path.split('/')[-2], image_path.split('/')[-1].split('.')[0])), 'wb') as f: 122 | 123 | # pickle.dump((image_array,point_array), f) 124 | # with open(os.path.join(output_folder,'{}_{}_{}.pkl'.format(image_path.split('/')[-3], image_path.split('/')[-2], image_path.split('/')[-1].split('.')[0])),'w') as f: 125 | # pickle.dump([image_array,point_array], f, 2) 126 | 127 | if __name__ == '__main__': 128 | print(time.strftime("%m-%d %H:%M:%S", time.localtime())) 129 | dump_data() 130 | print(time.strftime("%m-%d %H:%M:%S", time.localtime())) 131 | -------------------------------------------------------------------------------- /data_generate/_load_point.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os 4 | 5 | input_folder = '1a6f615e8b1b5ae4dbbc9440457e303e/temp' 6 | for index in range(24): 7 | point = np.load(os.path.join(input_folder,'{0:02d}.npy'.format(int(index)))) 8 | print(len(point)) 9 | img = cv2.imread(os.path.join(input_folder,'{0:02d}.png'.format(int(index)))) 10 | X,Y,Z = point.T 11 | F = 284 12 | h = (-Y)/(-Z)*F + 256/2.0 13 | w = X/(-Z)*F + 256/2.0 14 | h = np.minimum(np.maximum(h, 0), 255) 15 | w = np.minimum(np.maximum(w, 0), 255) 16 | img[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 17 | img[np.round(h).astype(int), np.round(w).astype(int), 1] = 255 18 | cv2.imwrite(os.path.join(input_folder,'{0:02d}.png'.format(int(index))).replace('.png','_prj.png'), img) 19 | 20 | -------------------------------------------------------------------------------- /data_generate/gen_pix3d.py: -------------------------------------------------------------------------------- 1 | from obj_io import parse_obj_file 2 | from sample_single import sample_faces 3 | import os 4 | import json 5 | import numpy as np 6 | from progress.bar import IncrementalBar 7 | data_dir_imgs = '/media/tree/data1/data/pix3d' 8 | 9 | def get_pix3d_models(): 10 | 11 | with open(os.path.join(data_dir_imgs, 'pix3d.json'), 'r') as f: 12 | models_dict = json.load(f) 13 | models = [] 14 | 15 | cats = ['chair','sofa','table'] 16 | 17 | 18 | # Check for truncation and occlusion before adding a model to the evaluation list 19 | for d in models_dict: 20 | if d['category'] in cats: 21 | if not d['truncated'] and not d['occluded'] and not d['slightly_occluded']: 22 | models.append(d) 23 | 24 | print 'Total models = {}\n'.format(len(models)) 25 | return models 26 | 27 | def sample_single(obj_path): 28 | # 1 sampling 29 | with open(obj_path,'r') as f: 30 | vertices, faces = parse_obj_file(f)[:2] 31 | sample = sample_faces(vertices, faces, 16384) 32 | position = sample * 0.57 33 | npy_path = obj_path.split('.')[0] 34 | np.save(npy_path, position) 35 | 36 | def sample_all(models): 37 | bar = IncrementalBar(max=len(models)) 38 | for i in range(len(models)): 39 | sample_single(os.path.join(data_dir_imgs, models[i]['model'])) 40 | bar.next() 41 | bar.finish() 42 | 43 | if __name__ == '__main__': 44 | pix3d_js = get_pix3d_models() 45 | # for i in pix3d_js: 46 | # print i['model'] 47 | sample_all(pix3d_js) 48 | -------------------------------------------------------------------------------- /data_generate/obj_io.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def parse_obj_file(open_file): 5 | """ 6 | Parse the supplied file. 7 | 8 | Args: 9 | `open_file`: file-like object with `readlines` method. 10 | 11 | Returns: positions, face_positions, texcoords, face_texcoords, \ 12 | normals, face_normals 13 | """ 14 | positions = [] 15 | texcoords = [] 16 | normals = [] 17 | face_positions = [] 18 | face_texcoords = [] 19 | face_normals = [] 20 | 21 | def parse_face(values): 22 | if len(values) != 3: 23 | raise ValueError('not a triangle at line' % lineno) 24 | for v in values: 25 | for j, index in enumerate(v.split('/')): 26 | if len(index): 27 | if j == 0: 28 | face_positions.append(int(index) - 1) 29 | elif j == 1: 30 | face_texcoords.append(int(index) - 1) 31 | elif j == 2: 32 | face_normals.append(int(index) - 1) 33 | 34 | parse_fns = { 35 | 'v': lambda values: positions.append([float(x) for x in values]), 36 | 'vt': lambda values: texcoords.append([float(x) for x in values]), 37 | 'vn': lambda values: normals.append([float(x) for x in values]), 38 | 'f': parse_face, 39 | 'mtllib': lambda values: None, 40 | 'o': lambda values: None, 41 | 'usemtl': lambda values: None, 42 | 's': lambda values: None, 43 | 'newmtl': lambda values: None, 44 | 'Ns': lambda values: None, 45 | 'Ni': lambda values: None, 46 | 'Ka': lambda values: None, 47 | 'Kd': lambda values: None, 48 | 'Ks': lambda values: None, 49 | 'd': lambda values: None, 50 | 'illum': lambda values: None, 51 | 'map_Kd': lambda values: None, 52 | } 53 | 54 | def parse_line(line): 55 | line = line.strip() 56 | if len(line) > 0 and line[0] != '#': 57 | values = line.split(' ') 58 | code = values[0] 59 | values = values[1:] 60 | if code in parse_fns: 61 | parse_fns[code](values) 62 | 63 | for lineno, line in enumerate(open_file.readlines()): 64 | parse_line(line) 65 | 66 | positions = np.array(positions, dtype=np.float32) 67 | texcoords = np.array(texcoords, dtype=np.float32) if len(texcoords) > 0 \ 68 | else None 69 | normals = np.array(normals, dtype=np.float32) \ 70 | if len(normals) > 0 else None 71 | face_positions = np.array(face_positions, dtype=np.uint32).reshape(-1, 3) 72 | 73 | face_texcoords = np.array(face_texcoords, dtype=np.uint32).reshape(-1, 3) \ 74 | if len(face_texcoords) > 0 else None 75 | face_normals = np.array(face_normals, dtype=np.uint32).reshape(-1, 3) \ 76 | if len(face_normals) > 0 else None 77 | 78 | return positions, face_positions, texcoords, face_texcoords, \ 79 | normals, face_normals 80 | 81 | 82 | def parse_obj(file_or_filename): 83 | """ 84 | Parse the given file, or opens the file at filename. 85 | 86 | See `parse_obj_file` for return values. 87 | """ 88 | if hasattr(file_or_filename, 'readlines'): 89 | return parse_obj_file(file_or_filename) 90 | else: 91 | with open(file_or_filename, 'r') as f: 92 | return parse_obj_file(f) 93 | 94 | 95 | def write_obj_file(fp, vertices, faces): 96 | for vertex in vertices: 97 | fp.write('v %s\n' % ' '.join((str(v) for v in vertex))) 98 | for face in faces: 99 | fp.write('f %s\n' % ' '.join((str(f+1) for f in face))) 100 | 101 | 102 | def write_obj(path_or_fp, vertices, faces): 103 | if hasattr(path_or_fp, 'write'): 104 | write_obj_file(path_or_fp, vertices, faces) 105 | else: 106 | with open(path_or_fp, 'w') as fp: 107 | write_obj_file(fp, vertices, faces) 108 | -------------------------------------------------------------------------------- /data_generate/obj_io.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/data_generate/obj_io.pyc -------------------------------------------------------------------------------- /data_generate/render_all.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | from filecmp import dircmp 4 | from progress.bar import IncrementalBar 5 | import numpy as np 6 | 7 | _FNULL = open(os.devnull, 'w') 8 | 9 | shapenet_category_to_id = { 10 | 'airplane' : '02691156', 11 | 'bench' : '02828884', 12 | 'cabinet' : '02933112', 13 | 'car' : '02958343', 14 | 'chair' : '03001627', 15 | 'lamp' : '03636649', 16 | 'monitor' : '03211117', 17 | 'rifle' : '04090263', 18 | 'sofa' : '04256520', 19 | 'speaker' : '03691459', 20 | 'table' : '04379243', 21 | 'telephone' : '04401088', 22 | 'vessel' : '04530566' 23 | } 24 | 25 | def get_path(cat_id): 26 | obj_prefix = os.path.join('/media/tree/data1/data/ShapeNetCore.v1',cat_id) 27 | view_prefix = os.path.join('/media/tree/data1/projects/PointGAN/3d-lmnet/data/ShapeNetRendering',cat_id) 28 | output_prefix = os.path.join('/media/tree/data1/projects/AttentionBased/data/image_256_256_12_with_texture',cat_id) 29 | 30 | view_path = '/media/tree/data1/projects/AttentionBased/PSGN/data_generate/rendering_metadata.txt' 31 | if not os.path.isdir(output_prefix): 32 | os.makedirs(output_prefix) 33 | obj_path = [] 34 | # view_path = [] 35 | output_folder = [] 36 | 37 | for i in dircmp(obj_prefix,view_prefix).common_dirs: 38 | obj_path.append(os.path.join(obj_prefix, i, 'model.obj')) 39 | # view_path.append(os.path.join(view_prefix, i, 'rendering/rendering_metadata.txt')) 40 | output_folder.append(os.path.join(output_prefix, i)) 41 | 42 | for i in output_folder: 43 | if not os.path.isdir(i): 44 | os.makedirs(i) 45 | return obj_path, view_path, output_folder 46 | 47 | def render_all(): 48 | script_path = os.path.join(os.path.realpath(os.path.dirname(__file__)), 'render_single.py') 49 | for cat, cat_id in shapenet_category_to_id.items(): 50 | obj_path, view_path, output_folder = get_path(cat_id) 51 | print('Rendering %d images for cat %s' % (len(obj_path),cat_id)) 52 | bar = IncrementalBar(max=len(obj_path)) 53 | call_kwargs = dict(stdout=_FNULL, stderr=subprocess.STDOUT) 54 | for i in range(len(obj_path)): 55 | bar.next() 56 | subprocess.call([ 57 | 'blender', 58 | '--background', 59 | '--python', script_path, '--', 60 | '--obj_path', str(obj_path[i]), 61 | '--view_path', str(view_path), 62 | '--output_folder', str(output_folder[i]), 63 | ],**call_kwargs) 64 | bar.finish() 65 | print('All cats rendered!') 66 | 67 | 68 | if __name__ == '__main__': 69 | render_all() 70 | 71 | 72 | 73 | -------------------------------------------------------------------------------- /data_generate/render_single.py: -------------------------------------------------------------------------------- 1 | import os 2 | import math 3 | import bpy 4 | import numpy as np 5 | 6 | def main(obj_path, view_path, output_folder, scale, shape): 7 | 8 | def camera_info(param): 9 | theta = np.deg2rad(param[0]) 10 | phi = np.deg2rad(param[1]) 11 | 12 | camY = param[3]*np.sin(phi) 13 | temp = param[3]*np.cos(phi) 14 | camX = temp * np.cos(theta) 15 | camZ = -temp * np.sin(theta) 16 | cam_pos = np.array([camX, camZ, camY]) 17 | 18 | return cam_pos 19 | 20 | cam_params = np.loadtxt(view_path) 21 | 22 | bpy.data.objects['Cube'].select = True #background black -> transparent or delete the original cube 23 | bpy.ops.object.delete() 24 | 25 | def merge_all(): 26 | bpy.ops.object.select_by_type(type="MESH") 27 | bpy.context.scene.objects.active = bpy.context.selected_objects[0] 28 | bpy.ops.object.join() 29 | obj = bpy.context.scene.objects.active 30 | bpy.ops.object.mode_set(mode='OBJECT') 31 | bpy.ops.object.location_clear() # Clear location - set location to (0,0,0) 32 | return obj 33 | 34 | bpy.ops.import_scene.obj(filepath=obj_path) 35 | 36 | #scale 37 | cur_obj = merge_all() 38 | cur_obj.scale = (scale, scale, scale) 39 | 40 | #remove texture 41 | obj = bpy.context.object 42 | mats = obj.data.materials 43 | 44 | for mat in mats: 45 | mat.user_clear() 46 | 47 | matLength = len(mats) 48 | for i in range(matLength): 49 | obj.data.materials.pop(0, update_data=True) 50 | 51 | leftOverMatBlocks = [block for block in bpy.data.materials if block.users == 0] 52 | for block in leftOverMatBlocks: 53 | bpy.data.materials.remove(block) 54 | bpy.context.scene.update() 55 | 56 | scene = bpy.context.scene 57 | #world light 58 | scene.world.light_settings.use_ambient_occlusion = True # turn AO on 59 | scene.world.light_settings.ao_factor = 0.5 # set it to 0.5 60 | 61 | #render setting 62 | scene.render.resolution_x = shape[1] 63 | scene.render.resolution_y = shape[0] 64 | scene.render.resolution_percentage = 100 65 | scene.render.alpha_mode = 'TRANSPARENT' 66 | scene.render.image_settings.file_format = 'PNG' # set output format to png 67 | 68 | def parent_obj_to_camera(b_camera): 69 | origin = (0, 0, 0) 70 | b_empty = bpy.data.objects.new("Empty", None) 71 | b_empty.location = origin 72 | b_camera.parent = b_empty # setup parenting 73 | return b_empty 74 | 75 | cam = scene.objects['Camera'] 76 | cam_constraint = cam.constraints.new(type='TRACK_TO') 77 | cam_constraint.track_axis = 'TRACK_NEGATIVE_Z' 78 | cam_constraint.up_axis = 'UP_Y' 79 | b_empty = parent_obj_to_camera(cam) 80 | cam_constraint.target = b_empty 81 | 82 | cam = bpy.data.objects["Camera"] 83 | 84 | 85 | for index, param in enumerate(cam_params): 86 | cam_pos = camera_info(param) 87 | scene.render.filepath = os.path.join(output_folder,'{0:02d}'.format(int(index))) 88 | cam.location = cam_pos 89 | bpy.ops.render.render(write_still=True) # render still 90 | 91 | 92 | def get_args(): 93 | import argparse 94 | import sys 95 | parser = argparse.ArgumentParser( 96 | description='Renders given obj file by rotation a camera around it.') 97 | parser.add_argument('--obj_path', type=str, default='1a6f615e8b1b5ae4dbbc9440457e303e/model.obj', 98 | help='Path to the obj file to be rendered.') 99 | parser.add_argument('--view_path', type=str, default='1a6f615e8b1b5ae4dbbc9440457e303e/rendering/rendering_metadata.txt', 100 | help='Path to the view file which indicates camera position.') 101 | parser.add_argument('--output_folder', type=str, default='1a6f615e8b1b5ae4dbbc9440457e303e/temp/', 102 | help='The path the output will be dumped to.') 103 | parser.add_argument('--scale', type=float, default=0.57, 104 | help='Scaling factor applied to model. ' 105 | 'Depends on size of mesh.') 106 | parser.add_argument('--shape', type=int, default=[256, 256], nargs=2, 107 | help='2D shape of rendered images.') 108 | argv = sys.argv[sys.argv.index("--") + 1:] 109 | args = parser.parse_args(argv) 110 | return args 111 | 112 | args = get_args() 113 | main(args.obj_path, args.view_path, args.output_folder, args.scale, args.shape) 114 | -------------------------------------------------------------------------------- /data_generate/rendering_metadata.txt: -------------------------------------------------------------------------------- 1 | 30.0 27.0 0 0.8 25 2 | 45.0 27.0 0 0.8 25 3 | 60.0 27.0 0 0.8 25 4 | 120.0 27.0 0 0.8 25 5 | 135.0 27.0 0 0.8 25 6 | 150.0 27.0 0 0.8 25 7 | 210.0 27.0 0 0.8 25 8 | 225.0 27.0 0 0.8 25 9 | 240.0 27.0 0 0.8 25 10 | 300.0 27.0 0 0.8 25 11 | 315.0 27.0 0 0.8 25 12 | 330.0 27.0 0 0.8 25 13 | -------------------------------------------------------------------------------- /data_generate/resample.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from tf_sampling import farthest_point_sample, gather_point 3 | import tensorflow as tf 4 | import os 5 | import h5py 6 | from progress.bar import IncrementalBar 7 | from utils.datafetcher import DataFetcher 8 | import time 9 | 10 | shapenet_id_to_category = { 11 | '02691156': 'airplane', 12 | '02828884': 'bench', 13 | '02933112': 'cabinet', 14 | '02958343': 'car', 15 | '03001627': 'chair', 16 | '03211117': 'monitor', 17 | '03636649': 'lamp', 18 | '03691459': 'speaker', 19 | '04090263': 'rifle', 20 | '04256520': 'sofa', 21 | '04379243': 'table', 22 | '04401088': 'telephone', 23 | '04530566': 'vessel' 24 | } 25 | 26 | def resample(pointcloud, npoint): 27 | new_pointcloud = gather_point(pointcloud, farthest_point_sample(npoint, pointcloud)) 28 | return new_pointcloud 29 | 30 | def resample_all(): 31 | data_prefix = '/media/tree/data1/projects/AttentionBased/data' 32 | train_output_folder = '/media/tree/backup/projects/AttentionBased/data/train' 33 | test_output_folder = '/media/tree/backup/projects/AttentionBased/data/test' 34 | # point_output_folder = 'point_resample_point_num_12' 35 | point_output_folder = 'point_4096_12' 36 | batch_size = 12 37 | resample_point_num = 4096 38 | 39 | #GPU settings 90% memory usage 40 | config = tf.ConfigProto() 41 | config.gpu_options.per_process_gpu_memory_fraction = 0.9 42 | config.gpu_options.allow_growth = True 43 | 44 | data = DataFetcher('train', batch_size = batch_size, epoch = 1) 45 | data.setDaemon(True) 46 | data.start() 47 | 48 | pt = tf.placeholder(tf.float32, shape = (None, 16384, 3), name = 'point') 49 | 50 | with tf.Session(config = config) as sess: 51 | 52 | for cat, batch in data.cats_batches.iteritems(): 53 | print('Resampling train {} starts at {}'.format(shapenet_id_to_category[cat], time.strftime("%m-%d %H:%M:%S", time.localtime()))) 54 | train_point_save = h5py.File(os.path.join(train_output_folder, point_output_folder, '{}.h5'.format(cat)), mode = 'w') 55 | train_pt_shape = (batch*batch_size, resample_point_num, 3) 56 | train_point_save.create_dataset('point', train_pt_shape, np.float32) 57 | bar = IncrementalBar(max = batch) 58 | for i in range(batch): 59 | _image, point = data.fetch() 60 | pointcloud = resample(pt, resample_point_num) 61 | point_array = sess.run(pointcloud, feed_dict = {pt:point}) 62 | # point_array = np.reshape(new_pt,(resample_point_num, 3)) 63 | # np.savetxt(os.path.join(train_output_folder, point_output_folder, '{}_{}.xyz'.format(cat, i)), new_pt) 64 | train_point_save['point'][i*batch_size:i*batch_size+batch_size, ...] = point_array 65 | bar.next() 66 | bar.finish() 67 | train_point_save.close() 68 | data.shutdown() 69 | 70 | data = DataFetcher('test', batch_size = batch_size, epoch = 1) 71 | data.setDaemon(True) 72 | data.start() 73 | 74 | with tf.Session(config = config) as sess: 75 | # pt = tf.placeholder(tf.float32, shape = (None, 16384, 3), name = 'point') 76 | for cat, batch in data.cats_batches.iteritems(): 77 | print('Resampling test {} starts at {}'.format(shapenet_id_to_category[cat], time.strftime("%m-%d %H:%M:%S", time.localtime()))) 78 | test_point_save = h5py.File(os.path.join(test_output_folder, point_output_folder, '{}.h5'.format(cat)), mode = 'w') 79 | test_pt_shape = (batch*batch_size, resample_point_num, 3) 80 | test_point_save.create_dataset('point', test_pt_shape, np.float32) 81 | bar = IncrementalBar(max = batch) 82 | for i in range(batch): 83 | _image, point = data.fetch() 84 | pointcloud = resample(pt, resample_point_num) 85 | point_array = sess.run(pointcloud, feed_dict = {pt:point}) 86 | # point_array = np.reshape(new_pt,(resample_point_num, 3)) 87 | # np.savetxt(os.path.join(train_output_folder, point_output_folder, '{}_{}.xyz'.format(cat, i)), new_pt) 88 | test_point_save['point'][i*batch_size:i*batch_size+batch_size, ...] = point_array 89 | bar.next() 90 | bar.finish() 91 | test_point_save.close() 92 | data.shutdown() 93 | 94 | if __name__ == '__main__': 95 | resample_all() 96 | -------------------------------------------------------------------------------- /data_generate/sample_all.py: -------------------------------------------------------------------------------- 1 | import os 2 | from filecmp import dircmp 3 | from progress.bar import IncrementalBar 4 | import numpy as np 5 | from sample_single import sample_single 6 | 7 | shapenet_category_to_id = { 8 | 'airplane' : '02691156', 9 | 'bench' : '02828884', 10 | 'cabinet' : '02933112', 11 | 'car' : '02958343', 12 | 'chair' : '03001627', 13 | 'lamp' : '03636649', 14 | 'monitor' : '03211117', 15 | 'rifle' : '04090263', 16 | 'sofa' : '04256520', 17 | 'speaker' : '03691459', 18 | 'table' : '04379243', 19 | 'telephone' : '04401088', 20 | 'vessel' : '04530566' 21 | } 22 | 23 | def get_path(cat_id): 24 | obj_prefix = os.path.join('/media/tree/data1/data/ShapeNetCore.v1',cat_id) 25 | view_prefix = os.path.join('/media/tree/data1/projects/PointGAN/3d-lmnet/data/ShapeNetRendering',cat_id) 26 | output_prefix = os.path.join('/media/tree/data1/projects/AttentionBased/data/pointcloud_12/16384',cat_id) 27 | view_path = '/media/tree/data1/projects/AttentionBased/PSGN/data_generate/rendering_metadata.txt' 28 | if not os.path.isdir(output_prefix): 29 | os.makedirs(output_prefix) 30 | obj_path = [] 31 | # view_path = [] 32 | output_folder = [] 33 | 34 | for i in dircmp(obj_prefix,view_prefix).common_dirs: 35 | obj_path.append(os.path.join(obj_prefix, i, 'model.obj')) 36 | # view_path.append(os.path.join(view_prefix, i, 'rendering/rendering_metadata.txt')) 37 | output_folder.append(os.path.join(output_prefix, i)) 38 | 39 | for i in output_folder: 40 | if not os.path.isdir(i): 41 | os.makedirs(i) 42 | return obj_path, view_path, output_folder 43 | 44 | def sample_all(): 45 | for cat, cat_id in shapenet_category_to_id.items(): 46 | obj_path, view_path, output_folder = get_path(cat_id) 47 | print('Sampling %d pointclouds for cat %s' % (len(obj_path),cat_id)) 48 | bar = IncrementalBar(max=len(obj_path)) 49 | for i in range(len(obj_path)): 50 | bar.next() 51 | sample_single(obj_path[i], view_path, output_folder[i]) 52 | bar.finish() 53 | print('All cats sampled!') 54 | 55 | if __name__ == '__main__': 56 | sample_all() 57 | -------------------------------------------------------------------------------- /data_generate/sample_single.py: -------------------------------------------------------------------------------- 1 | import os,sys 2 | import numpy as np 3 | import cv2 4 | from obj_io import parse_obj_file 5 | import sklearn.preprocessing 6 | 7 | def camera_info(param): 8 | theta = np.deg2rad(param[0]) 9 | phi = np.deg2rad(param[1]) 10 | 11 | camY = param[3]*np.sin(phi) 12 | temp = param[3]*np.cos(phi) 13 | camX = temp * np.cos(theta) 14 | camZ = temp * np.sin(theta) 15 | cam_pos = np.array([camX, camY, camZ]) 16 | 17 | axisZ = cam_pos.copy() 18 | axisY = np.array([0,1,0]) 19 | axisX = np.cross(axisY, axisZ) 20 | axisY = np.cross(axisZ, axisX) 21 | 22 | cam_mat = np.array([axisX, axisY, axisZ]) 23 | cam_mat = sklearn.preprocessing.normalize(cam_mat, axis=1) 24 | return cam_mat, cam_pos 25 | 26 | def sample_triangle(v, n=None): 27 | if hasattr(n, 'dtype'): 28 | n = np.asscalar(n) 29 | if n is None: 30 | size = v.shape[:-2] + (2,) 31 | elif isinstance(n, int): 32 | size = (n, 2) 33 | elif isinstance(n, tuple): 34 | size = n + (2,) 35 | elif isinstance(n, list): 36 | size = tuple(n) + (2,) 37 | else: 38 | raise TypeError('n must be int, tuple or list, got %s' % str(n)) 39 | assert(v.shape[-2] == 2) 40 | a = np.random.uniform(size=size) 41 | mask = np.sum(a, axis=-1) > 1 42 | a[mask] *= -1 43 | a[mask] += 1 44 | a = np.expand_dims(a, axis=-1) 45 | return np.sum(a*v, axis=-2) 46 | 47 | def sample_faces(vertices, faces, n_total): 48 | if len(faces) == 0: 49 | raise ValueError('Cannot sample points from zero faces.') 50 | tris = vertices[faces] 51 | n_faces = len(faces) 52 | d0 = tris[..., 0:1, :] 53 | ds = tris[..., 1:, :] - d0 54 | assert(ds.shape[1:] == (2, 3)) 55 | areas = 0.5 * np.sqrt(np.sum(np.cross(ds[:, 0], ds[:, 1])**2, axis=-1)) 56 | cum_area = np.cumsum(areas) 57 | cum_area *= (n_total / cum_area[-1]) 58 | cum_area = np.round(cum_area).astype(np.int32) 59 | 60 | positions = [] 61 | last = 0 62 | for i in range(n_faces): 63 | n = cum_area[i] - last 64 | last = cum_area[i] 65 | if n > 0: 66 | positions.append(d0[i] + sample_triangle(ds[i], n)) 67 | return np.concatenate(positions, axis=0) 68 | 69 | def sample_single(obj_path, view_path, output_folder): 70 | # 1 sampling 71 | with open(obj_path,'r') as f: 72 | vertices, faces = parse_obj_file(f)[:2] 73 | sample = sample_faces(vertices, faces, 16384) 74 | 75 | # 2 tranform to camera view 76 | position = sample * 0.57 77 | cam_params = np.loadtxt(view_path) 78 | 79 | for index, param in enumerate(cam_params): 80 | cam_mat, cam_pos = camera_info(param) 81 | pt_trans = np.dot(position-cam_pos, cam_mat.transpose()) 82 | pt_trans = pt_trans.astype(np.float32) 83 | npy_path = os.path.join(output_folder,'{0:02d}.npy'.format(int(index))) 84 | np.save(npy_path, pt_trans) 85 | # np.savetxt(npy_path.replace('npy','xyz'), pt_trans) 86 | 87 | 88 | if __name__ == '__main__': 89 | obj_path = '1a6f615e8b1b5ae4dbbc9440457e303e/model.obj' 90 | view_path = '1a6f615e8b1b5ae4dbbc9440457e303e/rendering/rendering_metadata.txt' 91 | output_folder = '1a6f615e8b1b5ae4dbbc9440457e303e/temp' 92 | sample_single(obj_path,view_path,output_folder) 93 | 94 | 95 | # # 1 sampling 96 | # obj_path = '1a6f615e8b1b5ae4dbbc9440457e303e/model.obj' 97 | # mesh_list = trimesh.load_mesh(obj_path) 98 | # if not isinstance(mesh_list, list): 99 | # mesh_list = [mesh_list] 100 | 101 | # area_sum = 0 102 | # for mesh in mesh_list: 103 | # area_sum += np.sum(mesh.area_faces) 104 | 105 | # # sample_16k = np.zeros((0,3), dtype=np.float32) 106 | # sample = np.zeros((0,3), dtype=np.float32) 107 | # # normal = np.zeros((0,3), dtype=np.float32) 108 | # total = 0 109 | # for mesh in mesh_list: 110 | # # number_16k = int(round(16384*np.sum(mesh.area_faces)/area_sum)) 111 | # number = int(round(1024*np.sum(mesh.area_faces)/area_sum)) 112 | # if number < 1: 113 | # continue 114 | # # points_16k, _index_16k = trimesh.sample.sample_surface(mesh, number_16k) 115 | # points, _index = trimesh.sample.sample_surface(mesh, number) 116 | 117 | # # sample_16k = np.append(sample_16k, points_16k, axis=0) 118 | # sample = np.append(sample, points, axis=0) 119 | 120 | # # triangles = mesh.triangles[index] 121 | # # pt1 = triangles[:,0,:] 122 | # # pt2 = triangles[:,1,:] 123 | # # pt3 = triangles[:,2,:] 124 | # # norm = np.cross(pt3-pt1, pt2-pt1) 125 | # # norm = sklearn.preprocessing.normalize(norm, axis=1) 126 | # # normal = np.append(normal, norm, axis=0) 127 | 128 | # # 2 tranform to camera view 129 | # # position_16k = sample_16k * 0.57 130 | # position = sample * 0.57 131 | 132 | # view_path = '1a6f615e8b1b5ae4dbbc9440457e303e/rendering/rendering_metadata.txt' 133 | # cam_params = np.loadtxt(view_path) 134 | # for index, param in enumerate(cam_params): 135 | # # camera tranform 136 | # cam_mat, cam_pos = camera_info(param) 137 | 138 | # # pt_trans_16k = np.dot(position_16k-cam_pos, cam_mat.transpose()) 139 | # pt_trans = np.dot(position-cam_pos, cam_mat.transpose()) 140 | # # nom_trans = np.dot(normal, cam_mat.transpose()) 141 | # # train_data = np.hstack((pt_trans, nom_trans)) 142 | # # train_data = pt_trans 143 | 144 | # img_path = os.path.join(os.path.split(view_path)[0], '%02d.png'%index) 145 | # # np.savetxt(img_path.replace('png','xyz'), train_data) 146 | 147 | # # np.savetxt(img_path.replace('png','xyz'), pt_trans_16k) 148 | # np.savetxt(img_path.replace('png','xyz'), pt_trans) 149 | 150 | # # #### project for sure 151 | # # img = cv2.imread(img_path) 152 | 153 | # # X,Y,Z = pt_trans.T 154 | # # F = 284 155 | # # h = (-Y)/(-Z)*F + 256/2.0 156 | # # w = X/(-Z)*F + 256/2.0 157 | # # h = np.minimum(np.maximum(h, 0), 255) 158 | # # w = np.minimum(np.maximum(w, 0), 255) 159 | # # img[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 160 | # # img[np.round(h).astype(int), np.round(w).astype(int), 1] = 255 161 | # # cv2.imwrite(img_path.replace('.png','_prj.png'), img) 162 | -------------------------------------------------------------------------------- /data_generate/sample_single.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/data_generate/sample_single.pyc -------------------------------------------------------------------------------- /data_generate/test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os 3 | 4 | image_prefix = '/media/tree/backup/projects/AttentionBased/data/image' 5 | point_prefix = '/media/tree/backup/projects/AttentionBased/data/pointcloud_16384' 6 | output_folder = '/media/tree/backup/projects/AttentionBased/data/data' 7 | 8 | image_path = [] 9 | point_path = [] 10 | # common_path = [] 11 | 12 | image = os.listdir(image_prefix) 13 | image_path = [os.path.join(image_prefix, i) for i in image] 14 | point = os.listdir(point_prefix) 15 | point_path = [os.path.join(point_prefix, i) for i in point] 16 | 17 | 18 | def work(idx): 19 | # batch_img = [ np.load(self.image_path[idx+i]) for i in range(self.batch_size) ] 20 | # batch_label = [ np.load(self.point_path[idx+i]) for i in range(self.batch_size) ] 21 | # batch_model_id = [] 22 | 23 | batch_img = [] 24 | batch_label = [] 25 | 26 | for i in range(self.batch_size): 27 | 28 | image_path = image_path[idx+i] 29 | point_path = point_path[idx+i] 30 | # single_model_id = image_path.split('/')[-1] 31 | # image = cv2.imread(image_path) 32 | image_array = np.load(image_path) 33 | point_array = np.load(point_path) 34 | 35 | batch_img.append(image_array) 36 | batch_label.append(point_array) 37 | 38 | # batch_model_id.append(single_model_id) 39 | 40 | return np.array(batch_img), np.array(batch_label) 41 | # return np.array(batch_img), np.array(batch_label), batch_model_id 42 | 43 | def fetch(): 44 | idx = 0 45 | while idx <= len(image_path): 46 | work(idx) 47 | idx += 32 48 | 49 | if __name__ == '__main__': 50 | import time 51 | load_start = time.time() 52 | fetch() 53 | load_end = time.time() 54 | print('load_data', load_end - load_start) 55 | -------------------------------------------------------------------------------- /data_generate/tf_sampling.py: -------------------------------------------------------------------------------- 1 | ''' Furthest point sampling 2 | Original author: Haoqiang Fan 3 | Modified by Charles R. Qi 4 | All Rights Reserved. 2017. 5 | ''' 6 | import tensorflow as tf 7 | from tensorflow.python.framework import ops 8 | import sys 9 | import os 10 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 11 | sys.path.append(BASE_DIR) 12 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 13 | def prob_sample(inp,inpr): 14 | ''' 15 | input: 16 | batch_size * ncategory float32 17 | batch_size * npoints float32 18 | returns: 19 | batch_size * npoints int32 20 | ''' 21 | return sampling_module.prob_sample(inp,inpr) 22 | ops.NoGradient('ProbSample') 23 | # TF1.0 API requires set shape in C++ 24 | #@tf.RegisterShape('ProbSample') 25 | #def _prob_sample_shape(op): 26 | # shape1=op.inputs[0].get_shape().with_rank(2) 27 | # shape2=op.inputs[1].get_shape().with_rank(2) 28 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 29 | def gather_point(inp,idx): 30 | ''' 31 | input: 32 | batch_size * ndataset * 3 float32 33 | batch_size * npoints int32 34 | returns: 35 | batch_size * npoints * 3 float32 36 | ''' 37 | return sampling_module.gather_point(inp,idx) 38 | #@tf.RegisterShape('GatherPoint') 39 | #def _gather_point_shape(op): 40 | # shape1=op.inputs[0].get_shape().with_rank(3) 41 | # shape2=op.inputs[1].get_shape().with_rank(2) 42 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 43 | @tf.RegisterGradient('GatherPoint') 44 | def _gather_point_grad(op,out_g): 45 | inp=op.inputs[0] 46 | idx=op.inputs[1] 47 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 48 | def farthest_point_sample(npoint,inp): 49 | ''' 50 | input: 51 | int32 52 | batch_size * ndataset * 3 float32 53 | returns: 54 | batch_size * npoint int32 55 | ''' 56 | return sampling_module.farthest_point_sample(inp, npoint) 57 | ops.NoGradient('FarthestPointSample') 58 | 59 | 60 | if __name__=='__main__': 61 | import numpy as np 62 | np.random.seed(100) 63 | # triangles=np.random.rand(1,5,3,3).astype('float32') 64 | point_path = '/media/tree/data1/projects/AttentionBased/data/pointcloud/16384/02691156/1a04e3eab45ca15dd86060f189eb133/00.npy' 65 | triangles = np.load(point_path) 66 | triangles = triangles[np.newaxis,:,:] 67 | with tf.device('/gpu:1'): 68 | inp=tf.constant(triangles) 69 | tria=inp[:,:,0,:] 70 | trib=inp[:,:,1,:] 71 | tric=inp[:,:,2,:] 72 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 73 | randomnumbers=tf.random_uniform((1,8192)) 74 | triids=prob_sample(areas,randomnumbers) 75 | tria_sample=gather_point(tria,triids) 76 | trib_sample=gather_point(trib,triids) 77 | tric_sample=gather_point(tric,triids) 78 | us=tf.random_uniform((1,8192)) 79 | vs=tf.random_uniform((1,8192)) 80 | uplusv=1-tf.abs(us+vs-1) 81 | uminusv=us-vs 82 | us=(uplusv+uminusv)*0.5 83 | vs=(uplusv-uminusv)*0.5 84 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 85 | print 'pt_sample: ', len(pt_sample) 86 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 87 | print 'reduced_sample', len(reduced_sample) 88 | with tf.Session('') as sess: 89 | ret=sess.run(reduced_sample) 90 | print ret.shape,ret.dtype 91 | # import cPickle as pickle 92 | # pickle.dump(ret,open('1.pkl','wb'),-1) 93 | -------------------------------------------------------------------------------- /data_generate/tf_sampling.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/data_generate/tf_sampling.pyc -------------------------------------------------------------------------------- /data_generate/tf_sampling_compile.sh: -------------------------------------------------------------------------------- 1 | #/bin/bash 2 | /usr/local/cuda-8.0/bin/nvcc tf_sampling_g.cu -o tf_sampling_g.cu.o -c -O2 -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC 3 | 4 | g++ -std=c++11 tf_sampling.cpp tf_sampling_g.cu.o -o tf_sampling_so.so -shared -fPIC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include -I /usr/local/cuda-8.0/include -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -lcudart -L /usr/local/cuda-8.0/lib64/ -L /usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework -O2 -D_GLIBCXX_USE_CXX11_ABI=0 5 | -------------------------------------------------------------------------------- /data_generate/tf_sampling_g.cu: -------------------------------------------------------------------------------- 1 | /* Furthest point sampling GPU implementation 2 | * Original author: Haoqiang Fan 3 | * Modified by Charles R. Qi 4 | * All Rights Reserved. 2017. 5 | */ 6 | 7 | __global__ void cumsumKernel(int b,int n,const float * __restrict__ inp,float * __restrict__ out){ 8 | const int BlockSize=2048; 9 | const int paddingLevel=5; 10 | __shared__ float buffer4[BlockSize*4]; 11 | __shared__ float buffer[BlockSize+(BlockSize>>paddingLevel)]; 12 | for (int i=blockIdx.x;i>2; 18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4; 33 | }else{ 34 | float v=0; 35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v; 43 | } 44 | } 45 | int u=0; 46 | for (;(2<>(u+1));k+=blockDim.x){ 49 | int i1=(((k<<1)+2)<>paddingLevel; 52 | i2+=i2>>paddingLevel; 53 | buffer[i1]+=buffer[i2]; 54 | } 55 | } 56 | u--; 57 | for (;u>=0;u--){ 58 | __syncthreads(); 59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){ 60 | int i1=(((k<<1)+3)<>paddingLevel; 63 | i2+=i2>>paddingLevel; 64 | buffer[i1]+=buffer[i2]; 65 | } 66 | } 67 | __syncthreads(); 68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel); 71 | buffer4[k]+=buffer[k2]; 72 | buffer4[k+1]+=buffer[k2]; 73 | buffer4[k+2]+=buffer[k2]; 74 | buffer4[k+3]+=buffer[k2]; 75 | } 76 | } 77 | __syncthreads(); 78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2; 82 | float r2=runningsum+t; 83 | runningsum2=t-(r2-runningsum); 84 | runningsum=r2; 85 | __syncthreads(); 86 | } 87 | } 88 | } 89 | 90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){ 91 | int base=1; 92 | while (base=1;k>>=1) 99 | if (r>=k && dataset[i*n+r-k]>=q) 100 | r-=k; 101 | result[i*m+j]=r; 102 | } 103 | } 104 | } 105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){ 106 | if (m<=0) 107 | return; 108 | const int BlockSize=512; 109 | __shared__ float dists[BlockSize]; 110 | __shared__ int dists_i[BlockSize]; 111 | const int BufferSize=3072; 112 | __shared__ float buf[BufferSize*3]; 113 | for (int i=blockIdx.x;ibest){ 147 | best=d2; 148 | besti=k; 149 | } 150 | } 151 | dists[threadIdx.x]=best; 152 | dists_i[threadIdx.x]=besti; 153 | for (int u=0;(1<>(u+1))){ 156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out); 196 | } 197 | //require b*n working space 198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){ 199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp); 200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out); 201 | } 202 | //require 32*n working space 203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){ 204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out); 205 | } 206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){ 207 | gatherpointKernel<<>>(b,n,m,inp,idx,out); 208 | } 209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){ 210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g); 211 | } 212 | -------------------------------------------------------------------------------- /data_generate/tf_sampling_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/data_generate/tf_sampling_g.cu.o -------------------------------------------------------------------------------- /data_generate/tf_sampling_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/data_generate/tf_sampling_so.so -------------------------------------------------------------------------------- /images/overview_network.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/images/overview_network.jpg -------------------------------------------------------------------------------- /images/qualitative_result_pix3d.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/images/qualitative_result_pix3d.jpg -------------------------------------------------------------------------------- /images/qualitative_result_shapenet.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/images/qualitative_result_shapenet.jpg -------------------------------------------------------------------------------- /metric/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/__init__.pyc -------------------------------------------------------------------------------- /metric/approxmatch.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | void approxmatch_cpu(int b,int n,int m,float * xyz1,float * xyz2,float * match){ 18 | for (int i=0;i saturatedl(n,double(factorl)),saturatedr(m,double(factorr)); 22 | vector weight(n*m); 23 | for (int j=0;j=-2;j--){ 26 | //printf("i=%d j=%d\n",i,j); 27 | double level=-powf(4.0,j); 28 | if (j==-2) 29 | level=0; 30 | for (int k=0;k ss(m,1e-9); 42 | for (int k=0;k ss2(m,0); 59 | for (int k=0;k1){ 154 | printf("bad i=%d j=%d k=%d u=%f\n",i,j,k,u); 155 | } 156 | s+=u; 157 | } 158 | if (s<0.999 || s>1.001){ 159 | printf("bad i=%d j=%d s=%f\n",i,j,s); 160 | } 161 | } 162 | for (int j=0;j4.001){ 168 | printf("bad i=%d j=%d s=%f\n",i,j,s); 169 | } 170 | } 171 | }*/ 172 | /*for (int j=0;j1e-3) 222 | if (fabs(double(match[i*n*m+k*n+j]-match_cpu[i*n*m+j*m+k]))>1e-2){ 223 | printf("i %d j %d k %d m %f %f\n",i,j,k,match[i*n*m+k*n+j],match_cpu[i*n*m+j*m+k]); 224 | flag=false; 225 | break; 226 | } 227 | //emax=max(emax,fabs(double(match[i*n*m+k*n+j]-match_cpu[i*n*m+j*m+k]))); 228 | emax+=fabs(double(match[i*n*m+k*n+j]-match_cpu[i*n*m+j*m+k])); 229 | } 230 | } 231 | printf("emax_match=%f\n",emax/2/n/m); 232 | emax=0; 233 | for (int i=0;i<2;i++) 234 | emax+=fabs(double(cost[i]-cost_cpu[i])); 235 | printf("emax_cost=%f\n",emax/2); 236 | emax=0; 237 | for (int i=0;i<2*m*3;i++) 238 | emax+=fabs(double(grad[i]-grad_cpu[i])); 239 | //for (int i=0;i<3*m;i++){ 240 | //if (grad[i]!=0) 241 | //printf("i %d %f %f\n",i,grad[i],grad_cpu[i]); 242 | //} 243 | printf("emax_grad=%f\n",emax/(2*m*3)); 244 | 245 | cudaFree(xyz1_g); 246 | cudaFree(xyz2_g); 247 | cudaFree(match_g); 248 | cudaFree(cost_g); 249 | cudaFree(grad_g); 250 | 251 | return 0; 252 | } 253 | 254 | -------------------------------------------------------------------------------- /metric/approxmatch.cu: -------------------------------------------------------------------------------- 1 | //n<=4096, m<=1024 2 | __global__ void approxmatch(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,float * __restrict__ match){ 3 | const int MaxN=4096,MaxM=1024; 4 | __shared__ float remainL[MaxN],remainR[MaxM],ratioR[MaxM],ratioL[MaxN]; 5 | __shared__ int listR[MaxM],lc; 6 | float multiL,multiR; 7 | if (n>=m){ 8 | multiL=1; 9 | multiR=n/m; 10 | }else{ 11 | multiL=m/n; 12 | multiR=1; 13 | } 14 | for (int i=blockIdx.x;i=-2;j--){ 23 | float level=-powf(4.0f,j); 24 | if (j==-2){ 25 | level=0; 26 | } 27 | if (threadIdx.x==0){ 28 | lc=0; 29 | for (int k=0;k0) 31 | listR[lc++]=k; 32 | } 33 | __syncthreads(); 34 | int _lc=lc; 35 | for (int k=threadIdx.x;k>>(b,n,m,xyz1,xyz2,match); 94 | } 95 | __global__ void matchcost(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * __restrict__ out){ 96 | __shared__ float allsum[512]; 97 | const int Block=256; 98 | __shared__ float buf[Block*3]; 99 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,out); 138 | } 139 | __global__ void matchcostgrad(int b,int n,int m,const float * __restrict__ xyz1,const float * __restrict__ xyz2,const float * __restrict__ match,float * grad2){ 140 | __shared__ float sum_grad[256*3]; 141 | for (int i=blockIdx.x;i>>(b,n,m,xyz1,xyz2,match,grad2); 182 | } 183 | 184 | -------------------------------------------------------------------------------- /metric/makefile: -------------------------------------------------------------------------------- 1 | nvcc = /usr/local/cuda-8.0/bin/nvcc 2 | cudalib = /usr/local/cuda-8.0/lib64 3 | tensorflow = /usr/local/lib/python2.7/dist-packages/tensorflow/include 4 | 5 | all: tf_approxmatch_so.so tf_approxmatch_g.cu.o tf_nndistance_so.so tf_nndistance_g.cu.o 6 | 7 | 8 | tf_approxmatch_so.so: tf_approxmatch_g.cu.o tf_approxmatch.cpp 9 | g++ -std=c++11 tf_approxmatch.cpp tf_approxmatch_g.cu.o -o tf_approxmatch_so.so -shared -fPIC -I $(tensorflow) -lcudart -L $(cudalib) -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -I $TF_INC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -L $TF_LIB -L /usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework 10 | 11 | 12 | tf_approxmatch_g.cu.o: tf_approxmatch_g.cu 13 | $(nvcc) -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -c -o tf_approxmatch_g.cu.o tf_approxmatch_g.cu -I $(tensorflow) -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC -O2 -I $TF_INC -I$TF_INC/external/nsync/public -L$TF_LIB -ltensorflow_framework 14 | 15 | 16 | tf_nndistance_so.so: tf_nndistance_g.cu.o tf_nndistance.cpp 17 | g++ -std=c++11 tf_nndistance.cpp tf_nndistance_g.cu.o -o tf_nndistance_so.so -shared -fPIC -I $(tensorflow) -lcudart -L $(cudalib) -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -I $TF_INC -I /usr/local/lib/python2.7/dist-packages/tensorflow/include/external/nsync/public -L $TF_LIB -L /usr/local/lib/python2.7/dist-packages/tensorflow -ltensorflow_framework 18 | 19 | 20 | tf_nndistance_g.cu.o: tf_nndistance_g.cu 21 | $(nvcc) -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 -c -o tf_nndistance_g.cu.o tf_nndistance_g.cu -I $(tensorflow) -DGOOGLE_CUDA=1 -x cu -Xcompiler -fPIC -O2 -I $TF_INC -I$TF_INC/external/nsync/public -L$TF_LIB -ltensorflow_framework 22 | 23 | 24 | clean: 25 | rm tf_approxmatch_so.so 26 | rm tf_nndistance_so.so 27 | rm *.cu.o 28 | -------------------------------------------------------------------------------- /metric/tf_approxmatch.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import os.path as osp 4 | 5 | base_dir = osp.dirname(osp.abspath(__file__)) 6 | 7 | approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so')) 8 | 9 | 10 | def approx_match(xyz1,xyz2): 11 | ''' 12 | input: 13 | xyz1 : batch_size * #dataset_points * 3 14 | xyz2 : batch_size * #query_points * 3 15 | returns: 16 | match : batch_size * #query_points * #dataset_points 17 | ''' 18 | return approxmatch_module.approx_match(xyz1,xyz2) 19 | ops.NoGradient('ApproxMatch') 20 | #@tf.RegisterShape('ApproxMatch') 21 | @ops.RegisterShape('ApproxMatch') 22 | def _approx_match_shape(op): 23 | shape1=op.inputs[0].get_shape().with_rank(3) 24 | shape2=op.inputs[1].get_shape().with_rank(3) 25 | return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])] 26 | 27 | def match_cost(xyz1,xyz2,match): 28 | ''' 29 | input: 30 | xyz1 : batch_size * #dataset_points * 3 31 | xyz2 : batch_size * #query_points * 3 32 | match : batch_size * #query_points * #dataset_points 33 | returns: 34 | cost : batch_size 35 | ''' 36 | return approxmatch_module.match_cost(xyz1,xyz2,match) 37 | #@tf.RegisterShape('MatchCost') 38 | @ops.RegisterShape('MatchCost') 39 | def _match_cost_shape(op): 40 | shape1=op.inputs[0].get_shape().with_rank(3) 41 | shape2=op.inputs[1].get_shape().with_rank(3) 42 | shape3=op.inputs[2].get_shape().with_rank(3) 43 | return [tf.TensorShape([shape1.dims[0]])] 44 | @tf.RegisterGradient('MatchCost') 45 | def _match_cost_grad(op,grad_cost): 46 | xyz1=op.inputs[0] 47 | xyz2=op.inputs[1] 48 | match=op.inputs[2] 49 | grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match) 50 | 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] 51 | 52 | if __name__=='__main__': 53 | alpha=0.5 54 | beta=2.0 55 | import bestmatch 56 | import numpy as np 57 | import math 58 | import random 59 | import cv2 60 | 61 | import tf_nndistance 62 | 63 | npoint=100 64 | 65 | with tf.device('/gpu:2'): 66 | pt_in=tf.placeholder(tf.float32,shape=(1,npoint*4,3)) 67 | mypoints=tf.Variable(np.random.randn(1,npoint,3).astype('float32')) 68 | match=approx_match(pt_in,mypoints) 69 | loss=tf.reduce_sum(match_cost(pt_in,mypoints,match)) 70 | #match=approx_match(mypoints,pt_in) 71 | #loss=tf.reduce_sum(match_cost(mypoints,pt_in,match)) 72 | #distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints) 73 | #loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5 74 | #loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint 75 | 76 | optimizer=tf.train.GradientDescentOptimizer(1e-4).minimize(loss) 77 | with tf.Session('') as sess: 78 | sess.run(tf.initialize_all_variables()) 79 | while True: 80 | meanloss=0 81 | meantrueloss=0 82 | for i in xrange(1001): 83 | #phi=np.random.rand(4*npoint)*math.pi*2 84 | #tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:] 85 | #tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:] 86 | tpoints=np.hstack([np.linspace(-1,1,400)[:,None],(random.random()*2*np.linspace(1,0,400)**2)[:,None],np.zeros((400,1))])[None,:,:] 87 | trainloss,_=sess.run([loss,optimizer],feed_dict={pt_in:tpoints.astype('float32')}) 88 | trainloss,trainmatch=sess.run([loss,match],feed_dict={pt_in:tpoints.astype('float32')}) 89 | #trainmatch=trainmatch.transpose((0,2,1)) 90 | show=np.zeros((400,400,3),dtype='uint8')^255 91 | trainmypoints=sess.run(mypoints) 92 | for i in xrange(len(tpoints[0])): 93 | u=np.random.choice(range(len(trainmypoints[0])),p=trainmatch[0].T[i]) 94 | cv2.line(show, 95 | (int(tpoints[0][i,1]*100+200),int(tpoints[0][i,0]*100+200)), 96 | (int(trainmypoints[0][u,1]*100+200),int(trainmypoints[0][u,0]*100+200)), 97 | cv2.cv.CV_RGB(0,255,0)) 98 | for x,y,z in tpoints[0]: 99 | cv2.circle(show,(int(y*100+200),int(x*100+200)),2,cv2.cv.CV_RGB(255,0,0)) 100 | for x,y,z in trainmypoints[0]: 101 | cv2.circle(show,(int(y*100+200),int(x*100+200)),3,cv2.cv.CV_RGB(0,0,255)) 102 | cost=((tpoints[0][:,None,:]-np.repeat(trainmypoints[0][None,:,:],4,axis=1))**2).sum(axis=2)**0.5 103 | #trueloss=bestmatch.bestmatch(cost)[0] 104 | print trainloss#,trueloss 105 | cv2.imshow('show',show) 106 | cmd=cv2.waitKey(10)%256 107 | if cmd==ord('q'): 108 | break 109 | -------------------------------------------------------------------------------- /metric/tf_approxmatch.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_approxmatch.pyc -------------------------------------------------------------------------------- /metric/tf_approxmatch_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_approxmatch_g.cu.o -------------------------------------------------------------------------------- /metric/tf_approxmatch_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_approxmatch_so.so -------------------------------------------------------------------------------- /metric/tf_nndistance.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import os.path as osp 4 | 5 | base_dir = osp.dirname(osp.abspath(__file__)) 6 | 7 | nn_distance_module = tf.load_op_library(osp.join(base_dir, 'tf_nndistance_so.so')) 8 | 9 | 10 | def nn_distance(xyz1, xyz2): 11 | ''' 12 | Computes the distance of nearest neighbors for a pair of point clouds 13 | input: xyz1: (batch_size,#points_1,3) the first point cloud 14 | input: xyz2: (batch_size,#points_2,3) the second point cloud 15 | output: dist1: (batch_size,#point_1) distance from first to second 16 | output: idx1: (batch_size,#point_1) nearest neighbor from first to second 17 | output: dist2: (batch_size,#point_2) distance from second to first 18 | output: idx2: (batch_size,#point_2) nearest neighbor from second to first 19 | ''' 20 | return nn_distance_module.nn_distance(xyz1,xyz2) 21 | 22 | #@tf.RegisterShape('NnDistance') 23 | @ops.RegisterShape('NnDistance') 24 | def _nn_distance_shape(op): 25 | shape1=op.inputs[0].get_shape().with_rank(3) 26 | shape2=op.inputs[1].get_shape().with_rank(3) 27 | return [tf.TensorShape([shape1.dims[0],shape1.dims[1]]),tf.TensorShape([shape1.dims[0],shape1.dims[1]]), 28 | tf.TensorShape([shape2.dims[0],shape2.dims[1]]),tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 29 | @ops.RegisterGradient('NnDistance') 30 | def _nn_distance_grad(op,grad_dist1,grad_idx1,grad_dist2,grad_idx2): 31 | xyz1=op.inputs[0] 32 | xyz2=op.inputs[1] 33 | idx1=op.outputs[1] 34 | idx2=op.outputs[3] 35 | return nn_distance_module.nn_distance_grad(xyz1,xyz2,grad_dist1,idx1,grad_dist2,idx2) 36 | 37 | 38 | if __name__=='__main__': 39 | import numpy as np 40 | import random 41 | import time 42 | #from tensorflow.python.kernel_tests.gradient_checker import compute_gradient 43 | from tensorflow.python.ops.gradient_checker import compute_gradient 44 | random.seed(100) 45 | np.random.seed(100) 46 | with tf.Session('') as sess: 47 | xyz1=np.random.randn(32,16384,3).astype('float32') 48 | xyz2=np.random.randn(32,1024,3).astype('float32') 49 | with tf.device('/gpu:0'): 50 | inp1=tf.Variable(xyz1) 51 | inp2=tf.constant(xyz2) 52 | reta,retb,retc,retd=nn_distance(inp1,inp2) 53 | loss=tf.reduce_sum(reta)+tf.reduce_sum(retc) 54 | train=tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss) 55 | sess.run(tf.initialize_all_variables()) 56 | t0=time.time() 57 | t1=t0 58 | best=1e100 59 | for i in xrange(100): 60 | trainloss,_=sess.run([loss,train]) 61 | newt=time.time() 62 | best=min(best,newt-t1) 63 | print i,trainloss,(newt-t0)/(i+1),best 64 | t1=newt 65 | #print sess.run([inp1,retb,inp2,retd]) 66 | #grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2]) 67 | #for i,j in grads: 68 | #print i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j)) 69 | #for i in xrange(10): 70 | #t0=time.time() 71 | #a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2}) 72 | #print 'time',time.time()-t0 73 | #print a.shape,b.shape,c.shape,d.shape 74 | #print a.dtype,b.dtype,c.dtype,d.dtype 75 | #samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32') 76 | #dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 77 | #idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 78 | #print np.abs(dist1-a[:,samples]).max() 79 | #print np.abs(idx1-b[:,samples]).max() 80 | #dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 81 | #idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 82 | #print np.abs(dist2-c[:,samples]).max() 83 | #print np.abs(idx2-d[:,samples]).max() 84 | 85 | -------------------------------------------------------------------------------- /metric/tf_nndistance.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_nndistance.pyc -------------------------------------------------------------------------------- /metric/tf_nndistance_g.cu: -------------------------------------------------------------------------------- 1 | #if GOOGLE_CUDA 2 | #define EIGEN_USE_GPU 3 | #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 4 | 5 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 6 | const int batch=512; 7 | __shared__ float buf[batch*3]; 8 | for (int i=blockIdx.x;ibest){ 120 | result[(i*n+j)]=best; 121 | result_i[(i*n+j)]=best_i; 122 | } 123 | } 124 | __syncthreads(); 125 | } 126 | } 127 | } 128 | void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){ 129 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 130 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 131 | } 132 | __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){ 133 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 156 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 157 | } 158 | 159 | #endif 160 | -------------------------------------------------------------------------------- /metric/tf_nndistance_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_nndistance_g.cu.o -------------------------------------------------------------------------------- /metric/tf_nndistance_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/metric/tf_nndistance_so.so -------------------------------------------------------------------------------- /model/layer/attention_module.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | attention module of Residual Attention Network 4 | """ 5 | 6 | import tensorflow as tf 7 | from tensorflow.python.keras.layers import UpSampling2D 8 | 9 | from basic_layers import ResidualBlock 10 | 11 | 12 | class AttentionModule(object): 13 | """AttentionModuleClass""" 14 | def __init__(self, p=1, t=2, r=1): 15 | """ 16 | :param p: the number of pre-processing Residual Units before splitting into trunk branch and mask branch 17 | :param t: the number of Residual Units in trunk branch 18 | :param r: the number of Residual Units between adjacent pooling layer in the mask branch 19 | """ 20 | self.p = p 21 | self.t = t 22 | self.r = r 23 | 24 | self.residual_block = ResidualBlock() 25 | 26 | def f_prop(self, input, input_channels, scope="attention_module", is_training=True): 27 | """ 28 | f_prop function of attention module 29 | :param input: A Tensor. input data [batch_size, height, width, channel] 30 | :param input_channels: dimension of input channel. 31 | :param scope: str, tensorflow name scope 32 | :param is_training: boolean, whether training step or not(test step) 33 | :return: A Tensor [batch_size, height, width, channel] 34 | """ 35 | with tf.variable_scope(scope): 36 | 37 | # residual blocks(TODO: change this function) 38 | with tf.variable_scope("first_residual_blocks"): 39 | for i in range(self.p): 40 | input = self.residual_block.f_prop(input, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 41 | 42 | with tf.variable_scope("trunk_branch"): 43 | output_trunk = input 44 | for i in range(self.t): 45 | output_trunk = self.residual_block.f_prop(output_trunk, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 46 | 47 | with tf.variable_scope("soft_mask_branch"): 48 | 49 | with tf.variable_scope("down_sampling_1"): 50 | # max pooling 51 | filter_ = [1, 2, 2, 1] 52 | output_soft_mask = tf.nn.max_pool(input, ksize=filter_, strides=filter_, padding='SAME') 53 | 54 | for i in range(self.r): 55 | output_soft_mask = self.residual_block.f_prop(output_soft_mask, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 56 | 57 | with tf.variable_scope("skip_connection"): 58 | # TODO(define new blocks) 59 | output_skip_connection = self.residual_block.f_prop(output_soft_mask, input_channels, is_training=is_training) 60 | 61 | 62 | with tf.variable_scope("down_sampling_2"): 63 | # max pooling 64 | filter_ = [1, 2, 2, 1] 65 | output_soft_mask = tf.nn.max_pool(output_soft_mask, ksize=filter_, strides=filter_, padding='SAME') 66 | 67 | for i in range(self.r): 68 | output_soft_mask = self.residual_block.f_prop(output_soft_mask, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 69 | 70 | with tf.variable_scope("up_sampling_1"): 71 | for i in range(self.r): 72 | output_soft_mask = self.residual_block.f_prop(output_soft_mask, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 73 | 74 | # interpolation 75 | output_soft_mask = UpSampling2D([2, 2])(output_soft_mask) 76 | 77 | # add skip connection 78 | output_soft_mask += output_skip_connection 79 | 80 | with tf.variable_scope("up_sampling_2"): 81 | for i in range(self.r): 82 | output_soft_mask = self.residual_block.f_prop(output_soft_mask, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 83 | 84 | # interpolation 85 | output_soft_mask = UpSampling2D([2, 2])(output_soft_mask) 86 | 87 | 88 | with tf.variable_scope("output"): 89 | output_soft_mask = tf.layers.conv2d(output_soft_mask, filters=input_channels, kernel_size=1) 90 | output_soft_mask = tf.layers.conv2d(output_soft_mask, filters=input_channels, kernel_size=1) 91 | 92 | # sigmoid 93 | output_soft_mask = tf.nn.sigmoid(output_soft_mask) 94 | 95 | with tf.variable_scope("attention"): 96 | output = (1 + output_soft_mask) * output_trunk 97 | 98 | with tf.variable_scope("last_residual_blocks"): 99 | for i in range(self.p): 100 | output = self.residual_block.f_prop(output, input_channels, scope="num_blocks_{}".format(i), is_training=is_training) 101 | 102 | return output 103 | -------------------------------------------------------------------------------- /model/layer/tf_ops/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/__init__.py -------------------------------------------------------------------------------- /model/layer/tf_ops/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/__init__.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/cd/__init__.py -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/cd/__init__.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/tf_nndistance.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import os.path as osp 4 | 5 | base_dir = osp.dirname(osp.abspath(__file__)) 6 | 7 | nn_distance_module = tf.load_op_library(osp.join(base_dir, 'tf_nndistance_so.so')) 8 | 9 | 10 | def nn_distance(xyz1, xyz2): 11 | ''' 12 | Computes the distance of nearest neighbors for a pair of point clouds 13 | input: xyz1: (batch_size,#points_1,3) the first point cloud 14 | input: xyz2: (batch_size,#points_2,3) the second point cloud 15 | output: dist1: (batch_size,#point_1) distance from first to second 16 | output: idx1: (batch_size,#point_1) nearest neighbor from first to second 17 | output: dist2: (batch_size,#point_2) distance from second to first 18 | output: idx2: (batch_size,#point_2) nearest neighbor from second to first 19 | ''' 20 | return nn_distance_module.nn_distance(xyz1,xyz2) 21 | 22 | #@tf.RegisterShape('NnDistance') 23 | @ops.RegisterShape('NnDistance') 24 | def _nn_distance_shape(op): 25 | shape1=op.inputs[0].get_shape().with_rank(3) 26 | shape2=op.inputs[1].get_shape().with_rank(3) 27 | return [tf.TensorShape([shape1.dims[0],shape1.dims[1]]),tf.TensorShape([shape1.dims[0],shape1.dims[1]]), 28 | tf.TensorShape([shape2.dims[0],shape2.dims[1]]),tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 29 | @ops.RegisterGradient('NnDistance') 30 | def _nn_distance_grad(op,grad_dist1,grad_idx1,grad_dist2,grad_idx2): 31 | xyz1=op.inputs[0] 32 | xyz2=op.inputs[1] 33 | idx1=op.outputs[1] 34 | idx2=op.outputs[3] 35 | return nn_distance_module.nn_distance_grad(xyz1,xyz2,grad_dist1,idx1,grad_dist2,idx2) 36 | 37 | 38 | if __name__=='__main__': 39 | import numpy as np 40 | import random 41 | import time 42 | #from tensorflow.python.kernel_tests.gradient_checker import compute_gradient 43 | from tensorflow.python.ops.gradient_checker import compute_gradient 44 | random.seed(100) 45 | np.random.seed(100) 46 | with tf.Session('') as sess: 47 | xyz1=np.random.randn(32,16384,3).astype('float32') 48 | xyz2=np.random.randn(32,1024,3).astype('float32') 49 | with tf.device('/gpu:0'): 50 | inp1=tf.Variable(xyz1) 51 | inp2=tf.constant(xyz2) 52 | reta,retb,retc,retd=nn_distance(inp1,inp2) 53 | loss=tf.reduce_sum(reta)+tf.reduce_sum(retc) 54 | train=tf.train.GradientDescentOptimizer(learning_rate=0.05).minimize(loss) 55 | sess.run(tf.initialize_all_variables()) 56 | t0=time.time() 57 | t1=t0 58 | best=1e100 59 | for i in xrange(100): 60 | trainloss,_=sess.run([loss,train]) 61 | newt=time.time() 62 | best=min(best,newt-t1) 63 | print i,trainloss,(newt-t0)/(i+1),best 64 | t1=newt 65 | #print sess.run([inp1,retb,inp2,retd]) 66 | #grads=compute_gradient([inp1,inp2],[(16,32,3),(16,32,3)],loss,(1,),[xyz1,xyz2]) 67 | #for i,j in grads: 68 | #print i.shape,j.shape,np.mean(np.abs(i-j)),np.mean(np.abs(i)),np.mean(np.abs(j)) 69 | #for i in xrange(10): 70 | #t0=time.time() 71 | #a,b,c,d=sess.run([reta,retb,retc,retd],feed_dict={inp1:xyz1,inp2:xyz2}) 72 | #print 'time',time.time()-t0 73 | #print a.shape,b.shape,c.shape,d.shape 74 | #print a.dtype,b.dtype,c.dtype,d.dtype 75 | #samples=np.array(random.sample(range(xyz2.shape[1]),100),dtype='int32') 76 | #dist1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 77 | #idx1=((xyz1[:,samples,None,:]-xyz2[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 78 | #print np.abs(dist1-a[:,samples]).max() 79 | #print np.abs(idx1-b[:,samples]).max() 80 | #dist2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).min(axis=-1) 81 | #idx2=((xyz2[:,samples,None,:]-xyz1[:,None,:,:])**2).sum(axis=-1).argmin(axis=-1) 82 | #print np.abs(dist2-c[:,samples]).max() 83 | #print np.abs(idx2-d[:,samples]).max() 84 | 85 | -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/tf_nndistance.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/cd/tf_nndistance.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/tf_nndistance_g.cu: -------------------------------------------------------------------------------- 1 | #if GOOGLE_CUDA 2 | #define EIGEN_USE_GPU 3 | #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" 4 | 5 | __global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){ 6 | const int batch=512; 7 | __shared__ float buf[batch*3]; 8 | for (int i=blockIdx.x;ibest){ 120 | result[(i*n+j)]=best; 121 | result_i[(i*n+j)]=best_i; 122 | } 123 | } 124 | __syncthreads(); 125 | } 126 | } 127 | } 128 | void NmDistanceKernelLauncher(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i){ 129 | NmDistanceKernel<<>>(b,n,xyz,m,xyz2,result,result_i); 130 | NmDistanceKernel<<>>(b,m,xyz2,n,xyz,result2,result2_i); 131 | } 132 | __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){ 133 | for (int i=blockIdx.x;i>>(b,n,xyz1,m,xyz2,grad_dist1,idx1,grad_xyz1,grad_xyz2); 156 | NmDistanceGradKernel<<>>(b,m,xyz2,n,xyz1,grad_dist2,idx2,grad_xyz2,grad_xyz1); 157 | } 158 | 159 | #endif 160 | -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/tf_nndistance_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/cd/tf_nndistance_g.cu.o -------------------------------------------------------------------------------- /model/layer/tf_ops/cd/tf_nndistance_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/cd/tf_nndistance_so.so -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/emd/__init__.py -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/emd/__init__.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/tf_approxmatch.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.python.framework import ops 3 | import os.path as osp 4 | 5 | base_dir = osp.dirname(osp.abspath(__file__)) 6 | 7 | approxmatch_module = tf.load_op_library(osp.join(base_dir, 'tf_approxmatch_so.so')) 8 | 9 | 10 | def approx_match(xyz1,xyz2): 11 | ''' 12 | input: 13 | xyz1 : batch_size * #dataset_points * 3 14 | xyz2 : batch_size * #query_points * 3 15 | returns: 16 | match : batch_size * #query_points * #dataset_points 17 | ''' 18 | return approxmatch_module.approx_match(xyz1,xyz2) 19 | ops.NoGradient('ApproxMatch') 20 | #@tf.RegisterShape('ApproxMatch') 21 | @ops.RegisterShape('ApproxMatch') 22 | def _approx_match_shape(op): 23 | shape1=op.inputs[0].get_shape().with_rank(3) 24 | shape2=op.inputs[1].get_shape().with_rank(3) 25 | return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])] 26 | 27 | def match_cost(xyz1,xyz2,match): 28 | ''' 29 | input: 30 | xyz1 : batch_size * #dataset_points * 3 31 | xyz2 : batch_size * #query_points * 3 32 | match : batch_size * #query_points * #dataset_points 33 | returns: 34 | cost : batch_size 35 | ''' 36 | return approxmatch_module.match_cost(xyz1,xyz2,match) 37 | #@tf.RegisterShape('MatchCost') 38 | @ops.RegisterShape('MatchCost') 39 | def _match_cost_shape(op): 40 | shape1=op.inputs[0].get_shape().with_rank(3) 41 | shape2=op.inputs[1].get_shape().with_rank(3) 42 | shape3=op.inputs[2].get_shape().with_rank(3) 43 | return [tf.TensorShape([shape1.dims[0]])] 44 | @tf.RegisterGradient('MatchCost') 45 | def _match_cost_grad(op,grad_cost): 46 | xyz1=op.inputs[0] 47 | xyz2=op.inputs[1] 48 | match=op.inputs[2] 49 | grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match) 50 | 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] 51 | 52 | if __name__=='__main__': 53 | alpha=0.5 54 | beta=2.0 55 | import bestmatch 56 | import numpy as np 57 | import math 58 | import random 59 | import cv2 60 | 61 | import tf_nndistance 62 | 63 | npoint=100 64 | 65 | with tf.device('/gpu:2'): 66 | pt_in=tf.placeholder(tf.float32,shape=(1,npoint*4,3)) 67 | mypoints=tf.Variable(np.random.randn(1,npoint,3).astype('float32')) 68 | match=approx_match(pt_in,mypoints) 69 | loss=tf.reduce_sum(match_cost(pt_in,mypoints,match)) 70 | #match=approx_match(mypoints,pt_in) 71 | #loss=tf.reduce_sum(match_cost(mypoints,pt_in,match)) 72 | #distf,_,distb,_=tf_nndistance.nn_distance(pt_in,mypoints) 73 | #loss=tf.reduce_sum((distf+1e-9)**0.5)*0.5+tf.reduce_sum((distb+1e-9)**0.5)*0.5 74 | #loss=tf.reduce_max((distf+1e-9)**0.5)*0.5*npoint+tf.reduce_max((distb+1e-9)**0.5)*0.5*npoint 75 | 76 | optimizer=tf.train.GradientDescentOptimizer(1e-4).minimize(loss) 77 | with tf.Session('') as sess: 78 | sess.run(tf.initialize_all_variables()) 79 | while True: 80 | meanloss=0 81 | meantrueloss=0 82 | for i in xrange(1001): 83 | #phi=np.random.rand(4*npoint)*math.pi*2 84 | #tpoints=(np.hstack([np.cos(phi)[:,None],np.sin(phi)[:,None],(phi*0)[:,None]])*random.random())[None,:,:] 85 | #tpoints=((np.random.rand(400)-0.5)[:,None]*[0,2,0]+[(random.random()-0.5)*2,0,0]).astype('float32')[None,:,:] 86 | tpoints=np.hstack([np.linspace(-1,1,400)[:,None],(random.random()*2*np.linspace(1,0,400)**2)[:,None],np.zeros((400,1))])[None,:,:] 87 | trainloss,_=sess.run([loss,optimizer],feed_dict={pt_in:tpoints.astype('float32')}) 88 | trainloss,trainmatch=sess.run([loss,match],feed_dict={pt_in:tpoints.astype('float32')}) 89 | #trainmatch=trainmatch.transpose((0,2,1)) 90 | show=np.zeros((400,400,3),dtype='uint8')^255 91 | trainmypoints=sess.run(mypoints) 92 | for i in xrange(len(tpoints[0])): 93 | u=np.random.choice(range(len(trainmypoints[0])),p=trainmatch[0].T[i]) 94 | cv2.line(show, 95 | (int(tpoints[0][i,1]*100+200),int(tpoints[0][i,0]*100+200)), 96 | (int(trainmypoints[0][u,1]*100+200),int(trainmypoints[0][u,0]*100+200)), 97 | cv2.cv.CV_RGB(0,255,0)) 98 | for x,y,z in tpoints[0]: 99 | cv2.circle(show,(int(y*100+200),int(x*100+200)),2,cv2.cv.CV_RGB(255,0,0)) 100 | for x,y,z in trainmypoints[0]: 101 | cv2.circle(show,(int(y*100+200),int(x*100+200)),3,cv2.cv.CV_RGB(0,0,255)) 102 | cost=((tpoints[0][:,None,:]-np.repeat(trainmypoints[0][None,:,:],4,axis=1))**2).sum(axis=2)**0.5 103 | #trueloss=bestmatch.bestmatch(cost)[0] 104 | print trainloss#,trueloss 105 | cv2.imshow('show',show) 106 | cmd=cv2.waitKey(10)%256 107 | if cmd==ord('q'): 108 | break 109 | -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/tf_approxmatch.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/emd/tf_approxmatch.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/tf_approxmatch_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/emd/tf_approxmatch_g.cu.o -------------------------------------------------------------------------------- /model/layer/tf_ops/emd/tf_approxmatch_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/emd/tf_approxmatch_so.so -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/grouping/__init__.py -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/grouping/__init__.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/query_ball_point.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | void query_ball_point_cpu(int b, int n, int m, const float* radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | #include "cuPrintf.cuh" 9 | #include "cuPrintf.cu" 10 | 11 | using namespace std; 12 | using namespace std; 13 | float randomf(){ 14 | return (rand()+0.5)/(RAND_MAX+1.0); 15 | } 16 | static double get_time(){ 17 | timespec tp; 18 | clock_gettime(CLOCK_MONOTONIC,&tp); 19 | return tp.tv_sec+tp.tv_nsec*1e-9; 20 | } 21 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 22 | // output: idx (b,m,nsample) 23 | __global__ void query_ball_point_gpu(int b, int n, int m, const float* radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 24 | for (int i=0;i>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 117 | cudaDeviceSynchronize(); 118 | printf("query_ball_point gpu time %f\n",get_time()-t0); 119 | 120 | t0=get_time(); 121 | group_point_gpu<<<1,1>>>(b,n,c,m,nsample,points,idx,out); 122 | cudaDeviceSynchronize(); 123 | printf("grou_point gpu time %f\n",get_time()-t0); 124 | 125 | t0=get_time(); 126 | group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 127 | cudaDeviceSynchronize(); 128 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 129 | 130 | cudaFree(xyz1); 131 | cudaFree(xyz2); 132 | cudaFree(points); 133 | cudaFree(idx); 134 | cudaFree(out); 135 | cudaFree(grad_out); 136 | cudaFree(grad_points); 137 | return 0; 138 | } 139 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/query_ball_point_block.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, const float *radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int index = threadIdx.x; 21 | xyz1 += n*3*index; 22 | xyz2 += m*3*index; 23 | idx += m*nsample*index; 24 | 25 | for (int j=0;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 113 | cudaDeviceSynchronize(); 114 | printf("query_ball_point gpu time %f\n",get_time()-t0); 115 | 116 | t0=get_time(); 117 | group_point_gpu<<<1,b>>>(b,n,c,m,nsample,points,idx,out); 118 | cudaDeviceSynchronize(); 119 | printf("grou_point gpu time %f\n",get_time()-t0); 120 | 121 | t0=get_time(); 122 | group_point_grad_gpu<<<1,b>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 123 | cudaDeviceSynchronize(); 124 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 125 | 126 | cudaFree(xyz1); 127 | cudaFree(xyz2); 128 | cudaFree(points); 129 | cudaFree(idx); 130 | cudaFree(out); 131 | cudaFree(grad_out); 132 | cudaFree(grad_points); 133 | return 0; 134 | } 135 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/query_ball_point_grid.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | // input: radius (1), nsample (1), xyz1 (b,n,3), xyz2 (b,m,3) 18 | // output: idx (b,m,nsample) 19 | __global__ void query_ball_point_gpu(int b, int n, int m, const float *radius, int nsample, const float *xyz1, const float *xyz2, int *idx) { 20 | int batch_index = blockIdx.x; 21 | xyz1 += n*3*batch_index; 22 | xyz2 += m*3*batch_index; 23 | idx += m*nsample*batch_index; 24 | 25 | int index = threadIdx.x; 26 | int stride = blockDim.x; 27 | 28 | for (int j=index;j>>(b,n,m,radius,nsample,xyz1,xyz2,idx); 123 | cudaDeviceSynchronize(); 124 | printf("query_ball_point gpu time %f\n",get_time()-t0); 125 | 126 | t0=get_time(); 127 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 128 | cudaDeviceSynchronize(); 129 | printf("grou_point gpu time %f\n",get_time()-t0); 130 | 131 | t0=get_time(); 132 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 133 | cudaDeviceSynchronize(); 134 | printf("grou_point_grad gpu time %f\n",get_time()-t0); 135 | 136 | cudaFree(xyz1); 137 | cudaFree(xyz2); 138 | cudaFree(points); 139 | cudaFree(idx); 140 | cudaFree(out); 141 | cudaFree(grad_out); 142 | cudaFree(grad_points); 143 | return 0; 144 | } 145 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/selection_sort.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), val (b,m,n) 20 | void selection_sort_cpu(int b, int n, int m, int k, const float *dist, int *idx, float *val) { 21 | float *p_dist; 22 | float tmp; 23 | int tmpi; 24 | for (int i=0;i 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,k), val (b,m,k) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, float *dist, int *idx, float *val) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | idx+=m*k*batch_index; 24 | val+=m*k*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | float *p_dist; 30 | for (int j=index;j>>(b,n,m,k,dist,idx,val); 68 | cudaDeviceSynchronize(); 69 | printf("selection sort cpu time %f\n",get_time()-t0); 70 | 71 | return 0; 72 | } 73 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/selection_sort_const.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // input: k (1), distance matrix dist (b,m,n) 19 | // output: idx (b,m,n), dist_out (b,m,n) 20 | __global__ void selection_sort_gpu(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 21 | int batch_index = blockIdx.x; 22 | dist+=m*n*batch_index; 23 | outi+=m*n*batch_index; 24 | out+=m*n*batch_index; 25 | 26 | int index = threadIdx.x; 27 | int stride = blockDim.x; 28 | 29 | // copy from dist to dist_out 30 | for (int j=index;j>>(b,n,m,k,dist,idx,dist_out); 84 | cudaDeviceSynchronize(); 85 | printf("selection sort cpu time %f\n",get_time()-t0); 86 | 87 | //for (int i=0;i>>(b,n,m,nsample,radius,xyz1,xyz2,idx,pts_cnt); 127 | //cudaDeviceSynchronize(); 128 | } 129 | void selectionSortLauncher(int b, int n, int m, int k, const float *dist, int *outi, float *out) { 130 | selection_sort_gpu<<>>(b,n,m,k,dist,outi,out); 131 | //cudaDeviceSynchronize(); 132 | } 133 | void groupPointLauncher(int b, int n, int c, int m, int nsample, const float *points, const int *idx, float *out){ 134 | group_point_gpu<<>>(b,n,c,m,nsample,points,idx,out); 135 | //cudaDeviceSynchronize(); 136 | } 137 | void groupPointGradLauncher(int b, int n, int c, int m, int nsample, const float *grad_out, const int *idx, float *grad_points){ 138 | group_point_grad_gpu<<>>(b,n,c,m,nsample,grad_out,idx,grad_points); 139 | //group_point_grad_gpu<<<1,1>>>(b,n,c,m,nsample,grad_out,idx,grad_points); 140 | //cudaDeviceSynchronize(); 141 | } 142 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/tf_grouping_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/grouping/tf_grouping_g.cu.o -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/tf_grouping_op_test.py: -------------------------------------------------------------------------------- 1 | 2 | import tensorflow as tf 3 | import numpy as np 4 | from .tf_grouping import query_ball_point, group_point 5 | 6 | class GroupPointTest(tf.test.TestCase): 7 | def test(self): 8 | pass 9 | 10 | def test_grad(self): 11 | with tf.device('/gpu:0'): 12 | points = tf.constant(np.random.random((1,128,16)).astype('float32')) 13 | print(points) 14 | xyz1 = tf.constant(np.random.random((1,128,3)).astype('float32')) 15 | xyz2 = tf.constant(np.random.random((1,8,3)).astype('float32')) 16 | # radius = tf.constant([0.3], dtype="float32") 17 | radius = [0.3] 18 | nsample = 32 19 | idx, pts_cnt = query_ball_point(radius, nsample, xyz1, xyz2) 20 | grouped_points = group_point(points, idx) 21 | print(grouped_points) 22 | 23 | with self.test_session(): 24 | print("---- Going to compute gradient error") 25 | err = tf.test.compute_gradient_error(points, (1,128,16), grouped_points, (1,8,32,16)) 26 | print(err) 27 | self.assertLess(err, 1e-4) 28 | 29 | if __name__=='__main__': 30 | tf.test.main() 31 | -------------------------------------------------------------------------------- /model/layer/tf_ops/grouping/tf_grouping_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/grouping/tf_grouping_so.so -------------------------------------------------------------------------------- /model/layer/tf_ops/interpolating/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/interpolating/__init__.py -------------------------------------------------------------------------------- /model/layer/tf_ops/interpolating/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/interpolating/__init__.pyc -------------------------------------------------------------------------------- /model/layer/tf_ops/interpolating/interpolate.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include // memset 4 | #include // rand, RAND_MAX 5 | #include // sqrtf 6 | #include 7 | #include 8 | using namespace std; 9 | float randomf(){ 10 | return (rand()+0.5)/(RAND_MAX+1.0); 11 | } 12 | static double get_time(){ 13 | timespec tp; 14 | clock_gettime(CLOCK_MONOTONIC,&tp); 15 | return tp.tv_sec+tp.tv_nsec*1e-9; 16 | } 17 | 18 | // Find three nearest neigbors with square distance 19 | // input: xyz1 (b,n,3), xyz2(b,m,3) 20 | // output: dist (b,n,3), idx (b,n,3) 21 | void threenn_cpu(int b, int n, int m, const float *xyz1, const float *xyz2, float *dist, int *idx) { 22 | for (int i=0;i>paddingLevel)]; 12 | for (int i=blockIdx.x;i>2; 18 | for (int k=threadIdx.x*4;k>2)+(k>>(2+paddingLevel))]=v4; 33 | }else{ 34 | float v=0; 35 | for (int k2=k;k2>2)+(k>>(2+paddingLevel))]=v; 43 | } 44 | } 45 | int u=0; 46 | for (;(2<>(u+1));k+=blockDim.x){ 49 | int i1=(((k<<1)+2)<>paddingLevel; 52 | i2+=i2>>paddingLevel; 53 | buffer[i1]+=buffer[i2]; 54 | } 55 | } 56 | u--; 57 | for (;u>=0;u--){ 58 | __syncthreads(); 59 | for (int k=threadIdx.x;k>(u+1));k+=blockDim.x){ 60 | int i1=(((k<<1)+3)<>paddingLevel; 63 | i2+=i2>>paddingLevel; 64 | buffer[i1]+=buffer[i2]; 65 | } 66 | } 67 | __syncthreads(); 68 | for (int k=threadIdx.x*4;k>2)-1)+(((k>>2)-1)>>paddingLevel); 71 | buffer4[k]+=buffer[k2]; 72 | buffer4[k+1]+=buffer[k2]; 73 | buffer4[k+2]+=buffer[k2]; 74 | buffer4[k+3]+=buffer[k2]; 75 | } 76 | } 77 | __syncthreads(); 78 | for (int k=threadIdx.x;k>paddingLevel)]+runningsum2; 82 | float r2=runningsum+t; 83 | runningsum2=t-(r2-runningsum); 84 | runningsum=r2; 85 | __syncthreads(); 86 | } 87 | } 88 | } 89 | 90 | __global__ void binarysearchKernel(int b,int n,int m,const float * __restrict__ dataset,const float * __restrict__ query, int * __restrict__ result){ 91 | int base=1; 92 | while (base=1;k>>=1) 99 | if (r>=k && dataset[i*n+r-k]>=q) 100 | r-=k; 101 | result[i*m+j]=r; 102 | } 103 | } 104 | } 105 | __global__ void farthestpointsamplingKernel(int b,int n,int m,const float * __restrict__ dataset,float * __restrict__ temp,int * __restrict__ idxs){ 106 | if (m<=0) 107 | return; 108 | const int BlockSize=512; 109 | __shared__ float dists[BlockSize]; 110 | __shared__ int dists_i[BlockSize]; 111 | const int BufferSize=3072; 112 | __shared__ float buf[BufferSize*3]; 113 | for (int i=blockIdx.x;ibest){ 147 | best=d2; 148 | besti=k; 149 | } 150 | } 151 | dists[threadIdx.x]=best; 152 | dists_i[threadIdx.x]=besti; 153 | for (int u=0;(1<>(u+1))){ 156 | int i1=(threadIdx.x*2)<>>(b,n,inp,out); 196 | } 197 | //require b*n working space 198 | void probsampleLauncher(int b,int n,int m,const float * inp_p,const float * inp_r,float * temp,int * out){ 199 | cumsumKernel<<<32,512>>>(b,n,inp_p,temp); 200 | binarysearchKernel<<>>(b,n,m,temp,inp_r,out); 201 | } 202 | //require 32*n working space 203 | void farthestpointsamplingLauncher(int b,int n,int m,const float * inp,float * temp,int * out){ 204 | farthestpointsamplingKernel<<<32,512>>>(b,n,m,inp,temp,out); 205 | } 206 | void gatherpointLauncher(int b,int n,int m,const float * inp,const int * idx,float * out){ 207 | gatherpointKernel<<>>(b,n,m,inp,idx,out); 208 | } 209 | void scatteraddpointLauncher(int b,int n,int m,const float * out_g,const int * idx,float * inp_g){ 210 | scatteraddpointKernel<<>>(b,n,m,out_g,idx,inp_g); 211 | } 212 | -------------------------------------------------------------------------------- /model/layer/tf_ops/sampling/tf_sampling_g.cu.o: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/sampling/tf_sampling_g.cu.o -------------------------------------------------------------------------------- /model/layer/tf_ops/sampling/tf_sampling_so.so: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/VIM-Lab/AttentionDPCR/394cc4dbc1beb6adbe9441ddc0dcc17274f79e2e/model/layer/tf_ops/sampling/tf_sampling_so.so -------------------------------------------------------------------------------- /model/layer/tf_utils.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Created on November 26, 2017 3 | 4 | @author: optas 5 | ''' 6 | 7 | import tensorflow as tf 8 | import numpy as np 9 | 10 | 11 | def expand_scope_by_name(scope, name): 12 | """ expand tf scope by given name. 13 | """ 14 | 15 | if isinstance(scope, basestring): 16 | scope += '/' + name 17 | return scope 18 | 19 | if scope is not None: 20 | return scope.name + '/' + name 21 | else: 22 | return scope 23 | 24 | 25 | def replicate_parameter_for_all_layers(parameter, n_layers): 26 | if parameter is not None and len(parameter) != n_layers: 27 | if len(parameter) != 1: 28 | raise ValueError() 29 | parameter = np.array(parameter) 30 | parameter = parameter.repeat(n_layers).tolist() 31 | return parameter 32 | 33 | 34 | def reset_tf_graph(): 35 | ''' Reset's all variables of default-tf graph. Useful for jupyter. 36 | ''' 37 | if 'sess' in globals() and sess: 38 | sess.close() 39 | tf.reset_default_graph() 40 | 41 | 42 | def leaky_relu(alpha): 43 | if not (alpha < 1 and alpha > 0): 44 | raise ValueError() 45 | 46 | return lambda x: tf.maximum(alpha * x, x) 47 | 48 | 49 | def safe_log(x, eps=1e-12): 50 | return tf.log(tf.maximum(x, eps)) -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from utils.datafetcher_finetune import DataFetcher 3 | from model.dense import Finetune 4 | import time 5 | import os 6 | import sys 7 | from progress.bar import IncrementalBar 8 | import pandas as pd 9 | 10 | 11 | shapenet_id_to_category = { 12 | '02691156': 'airplane', 13 | '02828884': 'bench', 14 | '02933112': 'cabinet', 15 | '02958343': 'car', 16 | '03001627': 'chair', 17 | '03211117': 'monitor', 18 | '03636649': 'lamp', 19 | '03691459': 'speaker', 20 | '04090263': 'rifle', 21 | '04256520': 'sofa', 22 | '04379243': 'table', 23 | '04401088': 'telephone', 24 | '04530566': 'vessel' 25 | } 26 | 27 | def load_ckpt(model, model_dir): 28 | ckpt = tf.train.get_checkpoint_state(model_dir) 29 | if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path): 30 | print('Reloading model parameters') 31 | model.saver.restore(sess, ckpt.model_checkpoint_path) 32 | else: 33 | raise ValueError('No such file:[{}]'.format(model_dir)) 34 | return ckpt 35 | 36 | if __name__ == '__main__': 37 | model_name = sys.argv[1] 38 | model_dir = os.path.join('result', model_name) 39 | 40 | batch_size = 4 41 | 42 | # Load data 43 | data = DataFetcher('test', batch_size = batch_size) 44 | data.setDaemon(True) 45 | data.start() 46 | 47 | #GPU settings 90% memory usage 48 | config = tf.ConfigProto() 49 | config.gpu_options.per_process_gpu_memory_fraction = 0.9 50 | config.gpu_options.allow_growth = True 51 | 52 | all_cd = [] 53 | all_emd = [] 54 | all_cats = [] 55 | csv_name = None 56 | 57 | with tf.Session(config = config) as sess: 58 | model = Finetune(sess, 'test', batch_size) 59 | ckpt = load_ckpt(model, model_dir) 60 | csv_name = '{}.csv'.format(ckpt.model_checkpoint_path.split('.')[0]) 61 | print('Testing starts!') 62 | for cat, batch in data.cats_batches.iteritems(): 63 | print('Testing {}'.format(shapenet_id_to_category[cat])) 64 | bar = IncrementalBar(max = batch) 65 | cat_cd = 0.0 66 | cat_emd = 0.0 67 | for i in range(batch): 68 | image, point = data.fetch() 69 | cd, emd = model.test(image, point) 70 | # cd = model.test(image, point) 71 | #scale metric 72 | cd *= 1000. 73 | emd /= 1000. 74 | cat_cd += cd 75 | cat_emd += emd 76 | bar.next() 77 | bar.finish() 78 | cat_cd /= float(batch) 79 | cat_emd /= float(batch) 80 | all_cd.append(cat_cd) 81 | all_emd.append(cat_emd) 82 | all_cats.append(shapenet_id_to_category[cat]) 83 | print('{} cd: {}'.format(shapenet_id_to_category[cat], cat_cd)) 84 | print('{} emd: {}'.format(shapenet_id_to_category[cat], cat_emd)) 85 | 86 | data.shutdown() 87 | all_cats.append('mean') 88 | all_cd.append(sum(all_cd)/float((len(all_cd)))) 89 | all_emd.append(sum(all_emd)/float((len(all_emd)))) 90 | dataframe = pd.DataFrame({'cat':all_cats, 'cd':all_cd, 'emd':all_emd}) 91 | dataframe.to_csv(csv_name, index = False, sep = ',') 92 | print('Testing finished!') 93 | 94 | -------------------------------------------------------------------------------- /train_dense.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from utils.datafetcher_dense import DataFetcher 3 | from model.dense import Dense 4 | import time 5 | import os 6 | import sys 7 | import re 8 | 9 | #python train.py model_name 10 | 11 | if __name__ == '__main__': 12 | model_name = sys.argv[1] 13 | model_dir = os.path.join('result', model_name) 14 | if not os.path.isdir(model_dir): 15 | os.makedirs(model_dir) 16 | train_log = os.path.join(model_dir,'{}_train.log'.format(model_name, )) 17 | 18 | #Global variables setting 19 | epoch = 20 20 | batch_size = 12 21 | 22 | # Load data 23 | data = DataFetcher('train', epoch = epoch, batch_size = batch_size) 24 | data.setDaemon(True) 25 | data.start() 26 | train_number = data.iter 27 | 28 | #GPU settings 90% memory usage 29 | config = tf.ConfigProto() 30 | config.gpu_options.per_process_gpu_memory_fraction = 0.9 31 | config.gpu_options.allow_growth = True 32 | 33 | with tf.Session(config = config) as sess: 34 | model = Dense(sess, 'train', batch_size) 35 | sess.run(tf.global_variables_initializer()) 36 | start_epoch=0 37 | ckpt = tf.train.get_checkpoint_state(model_dir) 38 | if ckpt is not None: 39 | print ('loading '+ckpt.model_checkpoint_path + ' ....') 40 | model.saver.restore(sess, ckpt.model_checkpoint_path) 41 | start_epoch = int(ckpt.model_checkpoint_path.split('.')[0].split('_')[-1]) 42 | 43 | print('Training starts!') 44 | for e in range(start_epoch, epoch): 45 | print('---- Epoch {}/{} ----'.format(e + 1, epoch)) 46 | model.saver.save(sess, os.path.join(model_dir, 'dense_epoch_{}.ckpt'.format(e + 1))) 47 | for i in range(train_number): 48 | image, point = data.fetch() 49 | loss, cd = model.train(image, point) 50 | 51 | if i % 100 == 0 or i == train_number - 1: 52 | current_time = time.strftime("%m-%d %H:%M:%S", time.localtime()) 53 | print('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 54 | with open(train_log, 'a+') as f: 55 | f.write('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}\n'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 56 | 57 | model.saver.save(sess, os.path.join(model_dir, 'dense_epoch_{}.ckpt'.format(e + 1))) 58 | data.shutdown() 59 | print('Training finished!') 60 | 61 | -------------------------------------------------------------------------------- /train_finetune.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from utils.datafetcher_finetune import DataFetcher 3 | from model.dense import Finetune 4 | import time 5 | import os 6 | import sys 7 | 8 | #python train.py model_name 9 | 10 | def load_ckpt(model): 11 | model_dir = model_dir = os.path.join('result', 'sparse') 12 | ckpt = tf.train.get_checkpoint_state(model_dir) 13 | if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path): 14 | print('Reloading attention_baseline parameters') 15 | model.saver_sparse.restore(sess, ckpt.model_checkpoint_path) 16 | else: 17 | raise ValueError('No such file:[{}]'.format(model_dir)) 18 | model_dir = model_dir = os.path.join('result', 'dense') 19 | ckpt = tf.train.get_checkpoint_state(model_dir) 20 | if ckpt and tf.train.checkpoint_exists(ckpt.model_checkpoint_path): 21 | print('Reloading my_pu_comlex parameters') 22 | model.saver_dense.restore(sess, ckpt.model_checkpoint_path) 23 | else: 24 | raise ValueError('No such file:[{}]'.format(model_dir)) 25 | 26 | 27 | if __name__ == '__main__': 28 | model_name = sys.argv[1] 29 | model_dir = os.path.join('result', model_name) 30 | if not os.path.isdir(model_dir): 31 | os.makedirs(model_dir) 32 | train_log = os.path.join(model_dir,'{}_train.log'.format(model_name, )) 33 | 34 | #Global variables setting 35 | epoch = 40 36 | batch_size = 10 37 | 38 | # Load data 39 | data = DataFetcher('train', epoch = epoch, batch_size = batch_size) 40 | data.setDaemon(True) 41 | data.start() 42 | train_number = data.iter 43 | 44 | #GPU settings 90% memory usage 45 | config = tf.ConfigProto() 46 | config.gpu_options.per_process_gpu_memory_fraction = 0.9 47 | config.gpu_options.allow_growth = True 48 | 49 | with tf.Session(config = config) as sess: 50 | model = Finetune(sess, 'train', batch_size) 51 | sess.run(tf.global_variables_initializer()) 52 | start_epoch = 0 53 | ckpt = tf.train.get_checkpoint_state(model_dir) 54 | if ckpt is not None: 55 | print ( 'loading ' + ckpt.model_checkpoint_path + ' ...') 56 | model.saver.restore(sess, ckpt.model_checkpoint_path) 57 | start_epoch = int(ckpt.model_checkpoint_path.split('.')[0].split('_')[-1]) 58 | print('Training starts!') 59 | for e in range(start_epoch, epoch): 60 | model.saver.save(sess, os.path.join(model_dir, 'finetune_epoch_{}.ckpt'.format(e + 1))) 61 | if e == 0: 62 | print('restoring previous epoches') 63 | load_ckpt(model) 64 | print('---- Epoch {}/{} ----'.format(e + 1, epoch)) 65 | 66 | for i in range(train_number): 67 | image, point1, point2 = data.fetch() 68 | loss, cd = model.train(image, point1, point2) 69 | 70 | if i % 100 == 0 or i == train_number - 1: 71 | current_time = time.strftime("%m-%d %H:%M:%S", time.localtime()) 72 | print('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 73 | with open(train_log, 'a+') as f: 74 | f.write('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}\n'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 75 | 76 | model.saver.save(sess, os.path.join(model_dir, 'finetune_epoch_{}.ckpt'.format(e + 1))) 77 | data.shutdown() 78 | print('Training finished!') 79 | 80 | -------------------------------------------------------------------------------- /train_sparse.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from utils.datafetcher_sparse import DataFetcher 3 | from model.sparse import Sparse 4 | import time 5 | import os 6 | import sys 7 | 8 | #python train.py model_name 9 | 10 | if __name__ == '__main__': 11 | model_name = sys.argv[1] 12 | model_dir = os.path.join('result', model_name) 13 | if not os.path.isdir(model_dir): 14 | os.makedirs(model_dir) 15 | train_log = os.path.join(model_dir,'{}_train.log'.format(model_name, )) 16 | 17 | #Global variables setting 18 | epoch = 20 19 | batch_size = 32 20 | 21 | # Load data 22 | data = DataFetcher('train', batch_size = batch_size, epoch = epoch) 23 | data.setDaemon(True) 24 | data.start() 25 | train_number = data.iter 26 | 27 | #GPU settings 90% memory usage 28 | config = tf.ConfigProto() 29 | config.gpu_options.per_process_gpu_memory_fraction = 0.9 30 | config.gpu_options.allow_growth = True 31 | 32 | with tf.Session(config = config) as sess: 33 | model = Sparse(sess, 'train') 34 | sess.run(tf.global_variables_initializer()) 35 | 36 | print('Training starts!') 37 | for e in range(epoch): 38 | print('---- Epoch {}/{} ----'.format(e + 1, epoch)) 39 | model.saver.save(sess, os.path.join(model_dir, 'sparse_epoch_{}.ckpt'.format(e + 1))) 40 | for i in range(train_number): 41 | image, point = data.fetch() 42 | loss, cd = model.train(image, point) 43 | 44 | if i % 100 == 0 or i == train_number - 1: 45 | 46 | 47 | current_time = time.strftime("%m-%d %H:%M:%S", time.localtime()) 48 | print('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 49 | with open(train_log, 'a+') as f: 50 | f.write('Epoch {} / {} iter {} / {} --- Loss:{} - CD:{} - time:{}\n'.format(e + 1, epoch, i + 1, train_number, loss, cd, current_time)) 51 | 52 | model.saver.save(sess, os.path.join(model_dir, 'sparse_epoch_{}.ckpt'.format(e + 1))) 53 | data.shutdown() 54 | print('Training finished!') 55 | 56 | -------------------------------------------------------------------------------- /utils/datafetcher.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py 4 | import threading 5 | import Queue 6 | import math 7 | 8 | shapenet_category_to_id = { 9 | 'airplane' : '02691156', 10 | 'bench' : '02828884', 11 | 'cabinet' : '02933112', 12 | 'car' : '02958343', 13 | 'chair' : '03001627', 14 | 'lamp' : '03636649', 15 | 'monitor' : '03211117', 16 | 'rifle' : '04090263', 17 | 'sofa' : '04256520', 18 | 'speaker' : '03691459', 19 | 'table' : '04379243', 20 | 'telephone' : '04401088', 21 | 'vessel' : '04530566' 22 | } 23 | 24 | class DataFetcher(threading.Thread): 25 | def __init__(self, mode, batch_size = 32, epoch = 10): 26 | super(DataFetcher, self).__init__() 27 | self.stopped = False 28 | self.epoch = epoch 29 | self.current_epoch = 0 30 | self.queue = Queue.Queue(2) 31 | self.batch_size = batch_size 32 | self.mode = mode 33 | if self.mode == 'train': 34 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_192_256_12' 35 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_256_256_12' 36 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_16384_12' 37 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_1024_12' 38 | else: 39 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_192_256_12' 40 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_256_256_12' 41 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_16384_12' 42 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_1024_12' 43 | self.iter, self.cats_batches = self.calculate_cat_batch_number() 44 | 45 | def calculate_cat_batch_number(self): 46 | count = 0 47 | cats = shapenet_category_to_id.values() 48 | cat_batch_number = [] 49 | for cat in cats: 50 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as f: 51 | batch_number = f['image'].shape[0] / self.batch_size 52 | cat_batch_number.append(batch_number) 53 | count += batch_number 54 | cats_batches = dict(zip(cats, cat_batch_number)) 55 | print(cats_batches) 56 | return count, cats_batches 57 | 58 | def run(self): 59 | if self.mode == 'train': 60 | while self.current_epoch < self.epoch: 61 | for cat, batch in self.cats_batches.iteritems(): 62 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 63 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 64 | for i in range(0, batch * self.batch_size, self.batch_size): 65 | if self.stopped: 66 | break 67 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 68 | self.current_epoch += 1 69 | elif self.mode == 'predict': 70 | # for cat, batch in self.cats_batches.iteritems(): 71 | # with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 72 | # with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 73 | # for i in range(0, batch * self.batch_size, self.batch_size): 74 | # if self.stopped: 75 | # break 76 | # self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 77 | 78 | cat = shapenet_category_to_id['chair'] 79 | batch = self.cats_batches[cat] 80 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 81 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 82 | for i in range(300, batch * self.batch_size, self.batch_size): 83 | if self.stopped: 84 | break 85 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 86 | 87 | 88 | else: 89 | for cat, batch in self.cats_batches.iteritems(): 90 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 91 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 92 | for i in range(0, batch * self.batch_size, self.batch_size): 93 | if self.stopped: 94 | break 95 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 96 | 97 | def fetch(self): 98 | if self.stopped: 99 | return None 100 | return self.queue.get() 101 | 102 | def shutdown(self): 103 | self.stopped = True 104 | while not self.queue.empty(): 105 | self.queue.get() 106 | 107 | 108 | 109 | if __name__ == '__main__': 110 | data = DataFetcher('test',batch_size = 1) 111 | data.start() 112 | image, point = data.fetch() 113 | # current = 0 114 | 115 | # #create white background 116 | # background = np.zeros((256,256,3), dtype = np.uint8) 117 | # background.fill(255) 118 | 119 | # # 1. image (obj rendering) 120 | # img = image[current, ...] * 255 121 | # img = img.astype('uint8') 122 | # img += background 123 | # img = np.where(img > 255 , img - 255, img) 124 | # cv2.imwrite('{:0>4}.png'.format(current), img) 125 | 126 | # # 3. gt_rendering 127 | # gt_rendering = background 128 | # X, Y, Z = point.T 129 | # F = 284 130 | # h = (-Y)/(-Z)*F + 256/2.0 131 | # w = X/(-Z)*F + 256/2.0 132 | # # h = np.minimum(np.maximum(h, 0), 255) 133 | # # w = np.minimum(np.maximum(w, 0), 255) 134 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 0] = 0 135 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 136 | # cv2.imwrite('{:0>4}.jpg'.format(current), gt_rendering) 137 | data.shutdown() 138 | -------------------------------------------------------------------------------- /utils/datafetcher_dense.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py 4 | import threading 5 | import Queue 6 | import math 7 | 8 | shapenet_category_to_id = { 9 | 'airplane' : '02691156', 10 | 'bench' : '02828884', 11 | 'cabinet' : '02933112', 12 | 'car' : '02958343', 13 | 'chair' : '03001627', 14 | 'lamp' : '03636649', 15 | 'monitor' : '03211117', 16 | 'rifle' : '04090263', 17 | 'sofa' : '04256520', 18 | 'speaker' : '03691459', 19 | 'table' : '04379243', 20 | 'telephone' : '04401088', 21 | 'vessel' : '04530566' 22 | } 23 | 24 | #image_path -> input point path 25 | #point_path -> gt point path 26 | 27 | class DataFetcher(threading.Thread): 28 | def __init__(self, mode, batch_size = 32, epoch = 10): 29 | super(DataFetcher, self).__init__() 30 | self.stopped = False 31 | self.epoch = epoch 32 | self.current_epoch = 0 33 | self.queue = Queue.Queue(2) 34 | self.batch_size = batch_size 35 | self.mode = mode 36 | if self.mode == 'train': 37 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_192_256_12' 38 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/point_{}_12'.format(1024) 39 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_{}_12'.format(16384) 40 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_1024_12' 41 | else: 42 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_192_256_12' 43 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/point_{}_12'.format(1024) 44 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_{}_12'.format(16384) 45 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_1024_12' 46 | self.iter, self.cats_batches = self.calculate_cat_batch_number() 47 | 48 | def calculate_cat_batch_number(self): 49 | count = 0 50 | cats = shapenet_category_to_id.values() 51 | cat_batch_number = [] 52 | for cat in cats: 53 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as f: 54 | batch_number = f['point'].shape[0] / self.batch_size 55 | cat_batch_number.append(batch_number) 56 | count += batch_number 57 | cats_batches = dict(zip(cats, cat_batch_number)) 58 | print(cats_batches) 59 | return count, cats_batches 60 | 61 | def run(self): 62 | if self.mode == 'train': 63 | while self.current_epoch < self.epoch: 64 | for cat, batch in self.cats_batches.iteritems(): 65 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 66 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 67 | for i in range(0, batch * self.batch_size, self.batch_size): 68 | if self.stopped: 69 | break 70 | self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 71 | self.current_epoch += 1 72 | elif self.mode == 'predict': 73 | # for cat, batch in self.cats_batches.iteritems(): 74 | # with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 75 | # with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 76 | # for i in range(0, batch * self.batch_size, self.batch_size): 77 | # if self.stopped: 78 | # break 79 | # self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 80 | 81 | cat = shapenet_category_to_id['airplane'] 82 | batch = self.cats_batches[cat] 83 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 84 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 85 | for i in range(0, batch * self.batch_size, self.batch_size): 86 | if self.stopped: 87 | break 88 | self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 89 | 90 | else: 91 | for cat, batch in self.cats_batches.iteritems(): 92 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 93 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 94 | for i in range(0, batch * self.batch_size, self.batch_size): 95 | if self.stopped: 96 | break 97 | self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 98 | 99 | def fetch(self): 100 | if self.stopped: 101 | return None 102 | return self.queue.get() 103 | 104 | def shutdown(self): 105 | self.stopped = True 106 | while not self.queue.empty(): 107 | self.queue.get() 108 | 109 | 110 | 111 | if __name__ == '__main__': 112 | data = DataFetcher('test',batch_size = 1) 113 | data.start() 114 | image, point = data.fetch() 115 | # current = 0 116 | 117 | # #create white background 118 | # background = np.zeros((256,256,3), dtype = np.uint8) 119 | # background.fill(255) 120 | 121 | # # 1. image (obj rendering) 122 | # img = image[current, ...] * 255 123 | # img = img.astype('uint8') 124 | # img += background 125 | # img = np.where(img > 255 , img - 255, img) 126 | # cv2.imwrite('{:0>4}.png'.format(current), img) 127 | 128 | # # 3. gt_rendering 129 | # gt_rendering = background 130 | # X, Y, Z = point.T 131 | # F = 284 132 | # h = (-Y)/(-Z)*F + 256/2.0 133 | # w = X/(-Z)*F + 256/2.0 134 | # # h = np.minimum(np.maximum(h, 0), 255) 135 | # # w = np.minimum(np.maximum(w, 0), 255) 136 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 0] = 0 137 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 138 | # cv2.imwrite('{:0>4}.jpg'.format(current), gt_rendering) 139 | data.shutdown() 140 | -------------------------------------------------------------------------------- /utils/datafetcher_finetune.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py 4 | import threading 5 | import Queue 6 | import math 7 | 8 | shapenet_category_to_id = { 9 | 'airplane' : '02691156', 10 | 'bench' : '02828884', 11 | 'cabinet' : '02933112', 12 | 'car' : '02958343', 13 | 'chair' : '03001627', 14 | 'lamp' : '03636649', 15 | 'monitor' : '03211117', 16 | 'rifle' : '04090263', 17 | 'sofa' : '04256520', 18 | 'speaker' : '03691459', 19 | 'table' : '04379243', 20 | 'telephone' : '04401088', 21 | 'vessel' : '04530566' 22 | } 23 | 24 | #image_path -> input point path 25 | #point_path -> gt point path 26 | 27 | class DataFetcher(threading.Thread): 28 | def __init__(self, mode, batch_size = 32, epoch = 10): 29 | super(DataFetcher, self).__init__() 30 | self.stopped = False 31 | self.epoch = epoch 32 | self.current_epoch = 0 33 | self.queue = Queue.Queue(2) 34 | self.batch_size = batch_size 35 | self.mode = mode 36 | self.hierarchies = [1024,4096,16384] 37 | self.point_path_densepcr = {} 38 | if self.mode == 'train': 39 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_192_256_12' 40 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_256_256_12' 41 | self.point_path_densepcr[1024] = '/media/tree/backup/projects/AttentionBased/data/train/point_1024_12' 42 | self.point_path_densepcr[16384] = '/media/tree/backup/projects/AttentionBased/data/train/point_16384_12' 43 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_1024_12' 44 | else: 45 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_192_256_12' 46 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_256_256_12' 47 | self.point_path_densepcr[1024] = '/media/tree/backup/projects/AttentionBased/data/test/point_1024_12' 48 | self.point_path_densepcr[16384] = '/media/tree/backup/projects/AttentionBased/data/test/point_16384_12' 49 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_1024_12' 50 | self.iter, self.cats_batches = self.calculate_cat_batch_number() 51 | 52 | def calculate_cat_batch_number(self): 53 | count = 0 54 | cats = shapenet_category_to_id.values() 55 | cat_batch_number = [] 56 | for cat in cats: 57 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as f: 58 | batch_number = f['image'].shape[0] / self.batch_size 59 | cat_batch_number.append(batch_number) 60 | count += batch_number 61 | cats_batches = dict(zip(cats, cat_batch_number)) 62 | print(cats_batches) 63 | return count, cats_batches 64 | 65 | def run(self): 66 | if self.mode == 'train': 67 | while self.current_epoch < self.epoch: 68 | for cat, batch in self.cats_batches.iteritems(): 69 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 70 | with h5py.File(os.path.join(self.point_path_densepcr[1024], '{}.h5'.format(cat)), 'r') as fp1: 71 | with h5py.File(os.path.join(self.point_path_densepcr[16384], '{}.h5'.format(cat)), 'r') as fp3: 72 | for i in range(0, batch * self.batch_size, self.batch_size): 73 | if self.stopped: 74 | break 75 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp1['point'][i:i+self.batch_size], fp3['point'][i:i+self.batch_size])) 76 | self.current_epoch += 1 77 | 78 | elif self.mode == 'predict': 79 | cat = shapenet_category_to_id['chair'] 80 | batch = self.cats_batches[cat] 81 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 82 | with h5py.File(os.path.join(self.point_path_densepcr[16384], '{}.h5'.format(cat)), 'r') as fp3: 83 | for i in range(36, batch * self.batch_size, self.batch_size): 84 | if self.stopped: 85 | break 86 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp3['point'][i:i+self.batch_size])) 87 | 88 | 89 | else: 90 | for cat, batch in self.cats_batches.iteritems(): 91 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 92 | with h5py.File(os.path.join(self.point_path_densepcr[16384], '{}.h5'.format(cat)), 'r') as fp3: 93 | for i in range(0, batch * self.batch_size, self.batch_size): 94 | if self.stopped: 95 | break 96 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp3['point'][i:i+self.batch_size])) 97 | 98 | def fetch(self): 99 | if self.stopped: 100 | return None 101 | return self.queue.get() 102 | 103 | def shutdown(self): 104 | self.stopped = True 105 | while not self.queue.empty(): 106 | self.queue.get() 107 | 108 | 109 | 110 | if __name__ == '__main__': 111 | data = DataFetcher('test',batch_size = 1) 112 | data.start() 113 | image, point = data.fetch() 114 | # current = 0 115 | 116 | # #create white background 117 | # background = np.zeros((256,256,3), dtype = np.uint8) 118 | # background.fill(255) 119 | 120 | # # 1. image (obj rendering) 121 | # img = image[current, ...] * 255 122 | # img = img.astype('uint8') 123 | # img += background 124 | # img = np.where(img > 255 , img - 255, img) 125 | # cv2.imwrite('{:0>4}.png'.format(current), img) 126 | 127 | # # 3. gt_rendering 128 | # gt_rendering = background 129 | # X, Y, Z = point.T 130 | # F = 284 131 | # h = (-Y)/(-Z)*F + 256/2.0 132 | # w = X/(-Z)*F + 256/2.0 133 | # # h = np.minimum(np.maximum(h, 0), 255) 134 | # # w = np.minimum(np.maximum(w, 0), 255) 135 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 0] = 0 136 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 137 | # cv2.imwrite('{:0>4}.jpg'.format(current), gt_rendering) 138 | data.shutdown() 139 | -------------------------------------------------------------------------------- /utils/datafetcher_sparse.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import h5py 4 | import threading 5 | import Queue 6 | import math 7 | 8 | shapenet_category_to_id = { 9 | 'airplane' : '02691156', 10 | 'bench' : '02828884', 11 | 'cabinet' : '02933112', 12 | 'car' : '02958343', 13 | 'chair' : '03001627', 14 | 'lamp' : '03636649', 15 | 'monitor' : '03211117', 16 | 'rifle' : '04090263', 17 | 'sofa' : '04256520', 18 | 'speaker' : '03691459', 19 | 'table' : '04379243', 20 | 'telephone' : '04401088', 21 | 'vessel' : '04530566' 22 | } 23 | 24 | class DataFetcher(threading.Thread): 25 | def __init__(self, mode, batch_size = 32, epoch = 10): 26 | super(DataFetcher, self).__init__() 27 | self.stopped = False 28 | self.epoch = epoch 29 | self.current_epoch = 0 30 | self.queue = Queue.Queue(2) 31 | self.batch_size = batch_size 32 | self.mode = mode 33 | if self.mode == 'train': 34 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_192_256_12' 35 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/train/image_256_256_12' 36 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_16384_12' 37 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/train/point_1024_12' 38 | else: 39 | # self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_192_256_12' 40 | self.image_path = '/media/tree/backup/projects/AttentionBased/data/test/image_256_256_12' 41 | self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_16384_12' 42 | # self.point_path = '/media/tree/backup/projects/AttentionBased/data/test/point_1024_12' 43 | self.iter, self.cats_batches = self.calculate_cat_batch_number() 44 | 45 | def calculate_cat_batch_number(self): 46 | count = 0 47 | cats = shapenet_category_to_id.values() 48 | cat_batch_number = [] 49 | for cat in cats: 50 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as f: 51 | batch_number = f['image'].shape[0] / self.batch_size 52 | cat_batch_number.append(batch_number) 53 | count += batch_number 54 | cats_batches = dict(zip(cats, cat_batch_number)) 55 | print(cats_batches) 56 | return count, cats_batches 57 | 58 | def run(self): 59 | if self.mode == 'train': 60 | while self.current_epoch < self.epoch: 61 | for cat, batch in self.cats_batches.iteritems(): 62 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 63 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 64 | for i in range(0, batch * self.batch_size, self.batch_size): 65 | if self.stopped: 66 | break 67 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 68 | self.current_epoch += 1 69 | elif self.mode == 'predict': 70 | # for cat, batch in self.cats_batches.iteritems(): 71 | # with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 72 | # with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 73 | # for i in range(0, batch * self.batch_size, self.batch_size): 74 | # if self.stopped: 75 | # break 76 | # self.queue.put((fi['point'][i:i+self.batch_size], fp['point'][i:i+self.batch_size])) 77 | 78 | cat = shapenet_category_to_id['chair'] 79 | batch = self.cats_batches[cat] 80 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 81 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 82 | for i in range(0, batch * self.batch_size, self.batch_size): 83 | if self.stopped: 84 | break 85 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 86 | 87 | 88 | else: 89 | for cat, batch in self.cats_batches.iteritems(): 90 | with h5py.File(os.path.join(self.image_path, '{}.h5'.format(cat)), 'r') as fi: 91 | with h5py.File(os.path.join(self.point_path, '{}.h5'.format(cat)), 'r') as fp: 92 | for i in range(0, batch * self.batch_size, self.batch_size): 93 | if self.stopped: 94 | break 95 | self.queue.put((fi['image'][i:i+self.batch_size].astype('float32') / 255.0, fp['point'][i:i+self.batch_size])) 96 | 97 | def fetch(self): 98 | if self.stopped: 99 | return None 100 | return self.queue.get() 101 | 102 | def shutdown(self): 103 | self.stopped = True 104 | while not self.queue.empty(): 105 | self.queue.get() 106 | 107 | 108 | 109 | if __name__ == '__main__': 110 | data = DataFetcher('test',batch_size = 1) 111 | data.start() 112 | image, point = data.fetch() 113 | # current = 0 114 | 115 | # #create white background 116 | # background = np.zeros((256,256,3), dtype = np.uint8) 117 | # background.fill(255) 118 | 119 | # # 1. image (obj rendering) 120 | # img = image[current, ...] * 255 121 | # img = img.astype('uint8') 122 | # img += background 123 | # img = np.where(img > 255 , img - 255, img) 124 | # cv2.imwrite('{:0>4}.png'.format(current), img) 125 | 126 | # # 3. gt_rendering 127 | # gt_rendering = background 128 | # X, Y, Z = point.T 129 | # F = 284 130 | # h = (-Y)/(-Z)*F + 256/2.0 131 | # w = X/(-Z)*F + 256/2.0 132 | # # h = np.minimum(np.maximum(h, 0), 255) 133 | # # w = np.minimum(np.maximum(w, 0), 255) 134 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 0] = 0 135 | # gt_rendering[np.round(h).astype(int), np.round(w).astype(int), 2] = 0 136 | # cv2.imwrite('{:0>4}.jpg'.format(current), gt_rendering) 137 | data.shutdown() 138 | -------------------------------------------------------------------------------- /utils/render_balls_so.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | using namespace std; 6 | 7 | struct PointInfo{ 8 | int x,y,z; 9 | float r,g,b; 10 | }; 11 | 12 | extern "C"{ 13 | 14 | void render_ball(int h,int w,unsigned char * show,int n,int * xyzs,float * c0,float * c1,float * c2,int r){ 15 | r=max(r,1); 16 | vector depth(h*w,-2100000000); 17 | vector pattern; 18 | for (int dx=-r;dx<=r;dx++) 19 | for (int dy=-r;dy<=r;dy++) 20 | if (dx*dx+dy*dy=h || y2<0 || y2>=w) && depth[x2*w+y2]0: 116 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0)) 117 | if magnifyBlue>=2: 118 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0)) 119 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1)) 120 | if magnifyBlue>=2: 121 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1)) 122 | if showrot: 123 | cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0)) 124 | cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0)) 125 | cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0)) 126 | changed=True 127 | while True: 128 | if changed: 129 | render() 130 | changed=False 131 | cv2.imshow('show3d',show) 132 | if waittime==0: 133 | cmd=cv2.waitKey(10)%256 134 | else: 135 | cmd=cv2.waitKey(waittime)%256 136 | if cmd==ord('q'): 137 | break 138 | elif cmd==ord('Q'): 139 | sys.exit(0) 140 | if cmd==ord('n'): 141 | zoom*=1.1 142 | changed=True 143 | elif cmd==ord('m'): 144 | zoom/=1.1 145 | changed=True 146 | elif cmd==ord('r'): 147 | zoom=1.0 148 | changed=True 149 | elif cmd==ord('s'): 150 | cv2.imwrite('show3d.png',show) 151 | if waittime!=0: 152 | break 153 | return cmd 154 | if __name__=='__main__': 155 | data = DataFetcher('test',batch_size = 1) 156 | data.start() 157 | image, point = data.fetch() 158 | current = 0 159 | X, Y, Z = point.T 160 | point_t = np.concatenate([-Y,X,Z],1) 161 | gt_rendering = get2D(np.vstack(point_t)) 162 | cv2.imwrite('{:0>4}.jpg'.format(current), gt_rendering) 163 | data.shutdown() 164 | --------------------------------------------------------------------------------