├── 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 | 
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 | 
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 | 
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 |
--------------------------------------------------------------------------------