├── README.md ├── image └── figure.JPG ├── prerprocess ├── Ncut_9 │ ├── .DS_Store │ ├── ICgraph.m │ ├── NcutImage.m │ ├── README.txt │ ├── X2distances.m │ ├── a_times_b_cmplx.cpp │ ├── affinityic.cpp │ ├── affinityic.mexa64 │ ├── affinityic.mexmaci64 │ ├── affinityic.mexw64 │ ├── build_scene.m │ ├── cimgnbmap.cpp │ ├── cimgnbmap.mexa64 │ ├── cimgnbmap.mexmaci64 │ ├── cimgnbmap.mexw64 │ ├── compileDir_simple.m │ ├── computeEdges.m │ ├── computeW.m │ ├── compute_relation.m │ ├── compute_relation2.m │ ├── demm.m │ ├── demoMatrix.asv │ ├── demoMatrix.m │ ├── demoNcutClustering.m │ ├── demoNcutImage.m │ ├── discretisation.m │ ├── discretisationEigenVectorData.m │ ├── doog1.m │ ├── doog2.m │ ├── eigs_new.m │ ├── fft_filt_2.m │ ├── gaussian.m │ ├── imread_ncut.m │ ├── jpg_images │ │ ├── 1.jpg │ │ ├── 2.jpg │ │ ├── 3.jpg │ │ └── Thumbs.db │ ├── make_filterbank_even2.m │ ├── make_filterbank_odd2.m │ ├── mex_w_times_x_symmetric.cpp │ ├── mex_w_times_x_symmetric.mexa64 │ ├── mex_w_times_x_symmetric.mexmaci64 │ ├── mex_w_times_x_symmetric.mexw64 │ ├── ncut.m │ ├── ncutW.m │ ├── quadedgep.m │ ├── segment_affinity_matrix.txt │ ├── showmask.m │ ├── sparsifyc.cpp │ ├── sparsifyc.mexa64 │ ├── sparsifyc.mexmaci64 │ ├── sparsifyc.mexw64 │ ├── spmtimesd.cpp │ ├── spmtimesd.mexa64 │ ├── spmtimesd.mexmaci64 │ ├── spmtimesd.mexw64 │ ├── tree.txt │ └── tree6.txt ├── README.md ├── bat_preprocess.py ├── gen_grass_data.py ├── gen_txt.py ├── min_bounding_rect.py ├── pointcnn │ ├── data_utils.py │ ├── data_utils.pyc │ ├── example_x3_l4.py │ ├── folder for trans │ │ ├── data_utils.py │ │ ├── example_x3_l4.py │ │ ├── pointcnn.py │ │ ├── pointcnn_cls.py │ │ ├── pointfly.py │ │ ├── tf_sampling.py │ │ └── train_val_cls.py │ ├── gen_feature.py │ ├── gen_node_prob_feature.py │ ├── pointcnn.py │ ├── pointcnn_cls.py │ ├── pointfly.py │ ├── tf_sampling.py │ ├── train_val_cls.py │ └── val_cls.py ├── preprocess_all.py ├── preprocess_leaf.py ├── preprocess_merge_leaf.py ├── qhull_2d.py ├── util.py └── util_obb.py └── vdrae ├── README.md ├── bat_test_eval.py ├── evaluation.py ├── grassdata.py ├── grassmodel.py ├── pytorch_tools ├── examples │ └── snli │ │ └── spinn-example.py ├── setup.py └── torchfold │ ├── __init__.py │ ├── __pycache__ │ ├── __init__.cpython-35.pyc │ └── torchfold.cpython-35.pyc │ ├── torchfold.py │ └── torchfold_test.py ├── test_for_eval.py ├── torchfoldext.py ├── train.py ├── util.py └── util_obb.py /README.md: -------------------------------------------------------------------------------- 1 | This is the code repository for ["Hierarchy Denoising Recursive Autoencoders for 3D Scene Layout Prediction"][1] . 2 | 3 | IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2019 4 | 5 | Created by Yifei Shi, Angel Xuan Chang, Zhelun Wu, Manolis Savva and Kai Xu 6 | 7 | ![teaser](image/figure.JPG) 8 | 9 | ## Usage - VDRAE 10 | ### Dependancy 11 | The code includes two components: data preprocessing and VDRAE. 12 | 13 | ### Data preparation 14 | You need to process the raw data of [S3DIS](https://www.dropbox.com/s/e9e6o1iikeg4g8l/s3dis.zip?dl=0) or [Matterport3D](https://www.dropbox.com/s/ynvg1yzsn9m9g37/matterport.zip?dl=0) before feed it to VDRAE, please see [preprocess](https://github.com/yifeishi/HierarchyLayout/tree/master/prerprocess) for details. 15 | 16 | ### VDRAE 17 | The VDRAE is the main network of this project. After data preparation, you can train and test the scene by looking at [vdrae](https://github.com/yifeishi/HierarchyLayout/tree/master/vdrae). 18 | 19 | ## Citation 20 | If you find the code is useful, please cite: 21 | ~~~~ 22 | @inproceedings{shi2019hierarchy, 23 | author = {Yifei Shi and Angel Xuan Chang and and Zhelun Wu and Manolis Savva and Kai Xu}, 24 | booktitle = {Proc. Computer Vision and Pattern Recognition (CVPR), IEEE}, 25 | title = {Hierarchy Denoising Recursive Autoencoders for 3D Scene Layout Prediction}, 26 | year = {2019} 27 | } 28 | ~~~~ 29 | 30 | [1]: https://arxiv.org/pdf/1903.03757.pdf "Hierarchy Denoising Recursive Autoencoders for 3D Scene Layout Prediction" 31 | -------------------------------------------------------------------------------- /image/figure.JPG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/image/figure.JPG -------------------------------------------------------------------------------- /prerprocess/Ncut_9/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/.DS_Store -------------------------------------------------------------------------------- /prerprocess/Ncut_9/ICgraph.m: -------------------------------------------------------------------------------- 1 | function [W,imageEdges] = ICgraph(I,dataW,dataEdgemap); 2 | % [W,imageEdges] = ICgraph(I,dataW,dataEdgemap); 3 | % Input: 4 | % I = gray-level image 5 | % optional parameters: 6 | % dataW.sampleRadius=10; 7 | % dataW.sample_rate=0.3; 8 | % dataW.edgeVariance = 0.1; 9 | % 10 | % dataEdgemap.parametres=[4,3, 21,3];%[number of filter orientations, number of scales, filter size, elongation] 11 | % dataEdgemap.threshold=0.02; 12 | % 13 | % Output: 14 | % W: npixels x npixels similarity matrix based on Intervening Contours 15 | % imageEdges: image showing edges extracted in the image 16 | % 17 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 18 | 19 | 20 | 21 | [p,q] = size(I); 22 | 23 | if (nargin< 2) | isempty(dataW), 24 | dataW.sampleRadius=10; 25 | dataW.sample_rate=0.3; 26 | dataW.edgeVariance = 0.1; 27 | end 28 | 29 | if (nargin<3) | isempty(dataEdgemap), 30 | dataEdgemap.parametres=[4,3, 21,3];%[number of filter orientations, number of scales, filter size, elongation] 31 | dataEdgemap.threshold=0.02; 32 | end 33 | 34 | 35 | edgemap = computeEdges(I,dataEdgemap.parametres,dataEdgemap.threshold); 36 | imageEdges = edgemap.imageEdges; 37 | W = computeW(I,dataW,edgemap.emag,edgemap.ephase); 38 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/NcutImage.m: -------------------------------------------------------------------------------- 1 | function [SegLabel,NcutDiscrete,NcutEigenvectors,NcutEigenvalues,W,imageEdges]= NcutImage(I,nbSegments); 2 | % [SegLabel,NcutDiscrete,NcutEigenvectors,NcutEigenvalues,W,imageEdges]= NcutImage(I); 3 | % Input: I = brightness image 4 | % nbSegments = number of segmentation desired 5 | % Output: SegLable = label map of the segmented image 6 | % NcutDiscrete = Discretized Ncut vectors 7 | % 8 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 9 | 10 | 11 | 12 | if nargin <2, 13 | nbSegments = 10; 14 | end 15 | 16 | [W,imageEdges] = ICgraph(I); 17 | 18 | [NcutDiscrete,NcutEigenvectors,NcutEigenvalues] = ncutW(W,nbSegments); 19 | 20 | %% generate segmentation label map 21 | [nr,nc,nb] = size(I); 22 | 23 | SegLabel = zeros(nr,nc); 24 | for j=1:size(NcutDiscrete,2), 25 | SegLabel = SegLabel + j*reshape(NcutDiscrete(:,j),nr,nc); 26 | end -------------------------------------------------------------------------------- /prerprocess/Ncut_9/README.txt: -------------------------------------------------------------------------------- 1 | %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 2 | % Normalized Cut Segmentation Code % 3 | % % 4 | % Timothee Cour (INRIA), Stella Yu (Berkeley), Jianbo Shi (UPENN) % 5 | % % 6 | %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 7 | 8 | License 9 | This software is made publicly for research use only. It may be modified and redistributed under the terms of the GNU General Public License. 10 | 11 | Citation 12 | Please cite the following if you plan to use the code in your own work: 13 | * Normalized Cuts and Image Segmentation, Jianbo Shi and Jitendra Malik, IEEE Transactions on Pattern Analysis and Machine Intelligence (PAMI) 2000 14 | * Normalized Cut Segmentation Code, Timothee Cour, Stella Yu, Jianbo Shi. Copyright 2004 University of Pennsylvania, Computer and Information Science Department. 15 | 16 | Tested on matlab R2009b. 17 | 18 | Installation Notes : 19 | 20 | 1) After you unzipped the files to mydir, 21 | put the Current Directory in Matlab to mydir 22 | 23 | 2) In the matlab command prompt, 24 | type compileDir_simple to compile the mex files (ignore the error on the C++ non-mex file; needs to be done once) 25 | 26 | 3) You can now try any of the functions 27 | 28 | type demoNcutImage to see a demo of image segmentation 29 | type demoNcutClustering to see a demo of point cloud clustering 30 | 31 | 32 | Other top level functions: 33 | 34 | NcutImage.m: given image "I", segment it into "nbSegments" segments 35 | [SegLabel,NcutDiscrete,NcutEigenvectors,NcutEigenvalues,W]= NcutImage(I,nbSegments); 36 | 37 | ICgraph.m: compute Intervening Contour based pixel similarity matrix W 38 | W = ICgraph(I); 39 | 40 | ncutW.m: Given a similarity graph "W", computes Ncut clustering on the graph into "nbSegments" groups; 41 | [NcutDiscrete,NcutEigenvectors,NcutEigenvalues] = ncutW(W,nbSegments); 42 | 43 | 44 | Release notes: 45 | 46 | 2010, January 22: release of all c++ source mex files compatible with matlab R2009b 47 | 2006, May 04: release version 8: fixed incompatibility issues with new matlab 48 | 2004, June 18: release version 7: initial release 49 | 50 | Maintained by Timothee Cour, timothee dot cour at gmail dot com 51 | 52 | January 22, 2010. 53 | 54 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/X2distances.m: -------------------------------------------------------------------------------- 1 | function distances = X2distances(X,Sigma); 2 | %Timothee Cour, 2004 3 | [n,k] = size(X); 4 | if nargin >= 2 5 | X = X*inv(sqrtm(Sigma)); 6 | end 7 | temp = sum(X.*X,2); 8 | temp = repmat(temp,1,n); 9 | distances = -2*X*X' + temp + temp'; -------------------------------------------------------------------------------- /prerprocess/Ncut_9/affinityic.cpp: -------------------------------------------------------------------------------- 1 | /*================================================================ 2 | * function w = affinityic(emag,ephase,pi,pj,sigma) 3 | * Input: 4 | * emag = edge strength at each pixel 5 | * ephase = edge phase at each pixel 6 | * [pi,pj] = index pair representation for MALTAB sparse matrices 7 | * sigma = sigma for IC energy 8 | * Output: 9 | * w = affinity with IC at [pi,pj] 10 | * 11 | 12 | % test sequence 13 | f = synimg(10); 14 | [i,j] = cimgnbmap(size(f),2); 15 | [ex,ey,egx,egy] = quadedgep(f); 16 | a = affinityic(ex,ey,egx,egy,i,j) 17 | show_dist_w(f,a); 18 | 19 | * Stella X. Yu, Nov 19, 2001. 20 | *=================================================================*/ 21 | 22 | # include "mex.h" 23 | # include "math.h" 24 | 25 | void mexFunction( 26 | int nargout, 27 | mxArray *out[], 28 | int nargin, 29 | const mxArray *in[] 30 | ) 31 | { 32 | /* declare variables */ 33 | int nr, nc, np, total; 34 | int i, j, k, ix, iy, jx, jy, ii, jj, iip1, jjp1, iip2, jjp2, step; 35 | double sigma, di, dj, a, z, maxori, phase1, phase2, slope; 36 | // int *ir, *jc; 37 | mwIndex*ir,*jc; 38 | unsigned int *pi, *pj; 39 | double *emag, *ephase, *w; 40 | 41 | /* check argument */ 42 | if (nargin<4) { 43 | mexErrMsgTxt("Four input arguments required"); 44 | } 45 | if (nargout>1) { 46 | mexErrMsgTxt("Too many output arguments"); 47 | } 48 | 49 | /* get edgel information */ 50 | nr = mxGetM(in[0]); 51 | nc = mxGetN(in[0]); 52 | if ( nr*nc ==0 || nr != mxGetM(in[1]) || nc != mxGetN(in[1]) ) { 53 | mexErrMsgTxt("Edge magnitude and phase shall be of the same image size"); 54 | } 55 | emag = mxGetPr(in[0]); 56 | ephase = mxGetPr(in[1]); 57 | np = nr * nc; 58 | 59 | /* get new index pair */ 60 | if (!mxIsUint32(in[2]) | !mxIsUint32(in[3])) { 61 | mexErrMsgTxt("Index pair shall be of type UINT32"); 62 | } 63 | if (mxGetM(in[3]) * mxGetN(in[3]) != np + 1) { 64 | mexErrMsgTxt("Wrong index representation"); 65 | } 66 | pi = (unsigned int*)mxGetData(in[2]); 67 | pj = (unsigned int*)mxGetData(in[3]); 68 | 69 | /* create output */ 70 | out[0] = mxCreateSparse(np,np,pj[np],mxREAL); 71 | if (out[0]==NULL) { 72 | mexErrMsgTxt("Not enough memory for the output matrix"); 73 | } 74 | w = mxGetPr(out[0]); 75 | ir = mxGetIr(out[0]); 76 | jc = mxGetJc(out[0]); 77 | 78 | /* find my sigma */ 79 | if (nargin<5) { 80 | sigma = 0; 81 | for (k=0; ksigma) { sigma = emag[k]; } 83 | } 84 | sigma = sigma / 6; 85 | printf("sigma = %6.5f",sigma); 86 | } else { 87 | sigma = mxGetScalar(in[4]); 88 | } 89 | a = 0.5 / (sigma * sigma); 90 | 91 | /* computation */ 92 | total = 0; 93 | for (j=0; j= abs(dj)) { 121 | slope = dj / di; 122 | step = (iy>=jy) ? 1 : -1; 123 | 124 | iip1 = jy; 125 | jjp1 = jx; 126 | 127 | 128 | for (ii=0;ii maxori){ 137 | maxori = z; 138 | } 139 | } 140 | 141 | iip1 = iip2; 142 | jjp1 = jjp2; 143 | phase1 = phase2; 144 | } 145 | 146 | /* sample in j direction */ 147 | } else { 148 | slope = di / dj; 149 | step = (ix>=jx) ? 1: -1; 150 | 151 | jjp1 = jx; 152 | iip1 = jy; 153 | 154 | 155 | for (jj=0;jj maxori){ 164 | maxori = z; 165 | } 166 | 167 | } 168 | 169 | iip1 = iip2; 170 | jjp1 = jjp2; 171 | phase1 = phase2; 172 | } 173 | } 174 | 175 | maxori = 0.5 * maxori; 176 | maxori = exp(-maxori * maxori * a); 177 | } 178 | ir[total] = i; 179 | 180 | w[total] = maxori; 181 | total = total + 1; 182 | 183 | } /* i */ 184 | } /* j */ 185 | 186 | jc[np] = total; 187 | } 188 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/affinityic.mexa64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/affinityic.mexa64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/affinityic.mexmaci64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/affinityic.mexmaci64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/affinityic.mexw64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/affinityic.mexw64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/build_scene.m: -------------------------------------------------------------------------------- 1 | function [data,size_cluster] = build_scene(caseid) 2 | % [data,size_cluster] = build_scene(case) 3 | % 4 | % case 1: random gaussian shaped blobs. 5 | % case 2: one circular plus some gaussian blob, one inside, one outside 6 | % case 3: one circular plus some gaussian blob, one inside, two outside 7 | % Jianbo Shi, 1997 8 | 9 | 10 | if caseid ==1, 11 | sigma_h = 2; 12 | sigma_v = 10; 13 | 14 | s_v = 10; 15 | s_h = 30; 16 | 17 | a = [sigma_h*randn(1,40);sigma_v*randn(1,40)]; 18 | b = [s_h;s_v]*ones(1,50) + [sigma_h*randn(1,50);... 19 | sigma_v*randn(1,50)]; 20 | 21 | data = [a,b]; 22 | size_cluster = [40,50]; 23 | 24 | elseif caseid == 2, 25 | num_cluster = 3; 26 | radius = 15; 27 | size_cluster = [80,20,20]; 28 | 29 | raw_data = randn(2,sum(size_cluster)); 30 | tmp = rand(2,size_cluster(1))-0.5; 31 | 32 | [t,idt] = sort(tmp(2,:)); 33 | 34 | r_noise = 4; 35 | raw_data2 = [tmp(1,idt)*r_noise;... 36 | tmp(2,idt)*2]; 37 | disp(raw_data2) 38 | 39 | 40 | data = [(radius-raw_data2(1,1:size_cluster(1))).*... 41 | cos(pi*raw_data2(2,1:size_cluster(1)));... 42 | (radius-raw_data2(1,1:size_cluster(1))).*... 43 | sin(pi*raw_data2(2,1:size_cluster(1)))]; 44 | 45 | 46 | center = [0,0];sig = [1,2]; 47 | % size_cluster_base 48 | scb = size_cluster(1)+1; 49 | scb_next = scb+size_cluster(2)-1; 50 | data = [data,[center(1)+sig(1)*raw_data(1,scb:scb_next);... 51 | center(2)+sig(2)*raw_data(2,scb:scb_next)]]; 52 | 53 | 54 | center = [radius+10,0]; sig = [1,1]; 55 | scb = scb_next+1; 56 | scb_next = scb+size_cluster(3)-1; 57 | data = [data,[center(1)+sig(1)*raw_data(1,scb:scb_next);... 58 | center(2)+sig(2)*raw_data(2,scb:scb_next)]]; 59 | elseif caseid==3, 60 | num_cluster = 4; 61 | radius = 15; 62 | size_cluster = [80,10,20,20]; 63 | 64 | raw_data = randn(2,sum(size_cluster)); 65 | tmp = rand(2,size_cluster(1))-0.5; 66 | 67 | [t,idt] = sort(tmp(2,:)); 68 | r_noise = 4; 69 | raw_data2 = [tmp(1,idt)*r_noise;... 70 | tmp(2,idt)*2]; 71 | 72 | data = [(radius-raw_data2(1,1:size_cluster(1))).*... 73 | cos(pi*raw_data2(2,1:size_cluster(1)));... 74 | (radius-raw_data2(1,1:size_cluster(1))).*... 75 | sin(pi*raw_data2(2,1:size_cluster(1)))]; 76 | 77 | 78 | center = [0,0];sig = [1,2]; 79 | % size_cluster_base 80 | scb = size_cluster(1)+1; 81 | scb_next = scb+size_cluster(2)-1; 82 | data = [data,[center(1)+sig(1)*raw_data(1,scb:scb_next);... 83 | center(2)+sig(2)*raw_data(2,scb:scb_next)]]; 84 | 85 | 86 | center = [radius+25,8]; sig = [1,2.3]; 87 | scb = scb_next+1; 88 | scb_next = scb+size_cluster(3)-1; 89 | data = [data,[center(1)+sig(1)*raw_data(1,scb:scb_next);... 90 | center(2)+sig(2)*raw_data(2,scb:scb_next)]]; 91 | 92 | center = [radius+25,-6]; sig = [1.5,2.4]; 93 | scb = scb_next+1; 94 | scb_next = scb+size_cluster(4)-1; 95 | data = [data,[center(1)+sig(1)*raw_data(1,scb:scb_next);... 96 | center(2)+sig(2)*raw_data(2,scb:scb_next)]]; 97 | elseif caseid == 4, 98 | size_cluster = [100,10,10]; 99 | radius = 10; 100 | tmp = rand(2,size_cluster(1))-0.5; 101 | 102 | [t,idt] = sort(tmp(2,:)); 103 | r_noise = 4; 104 | raw_data2 = [tmp(1,idt)*r_noise;... 105 | tmp(2,idt)*2]; 106 | 107 | data = [(radius-raw_data2(1,1:size_cluster(1))).*... 108 | cos(pi*raw_data2(2,1:size_cluster(1)));... 109 | (radius-raw_data2(1,1:size_cluster(1))).*... 110 | sin(pi*raw_data2(2,1:size_cluster(1)))]; 111 | 112 | 113 | result = zeros(1,size_cluster(1)); 114 | 115 | % for j =1:size_cluster(1), 116 | % result(j) = sum(sum(A(1:j,1:j)))/j; 117 | % end 118 | 119 | end 120 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/cimgnbmap.cpp: -------------------------------------------------------------------------------- 1 | /*================================================================ 2 | * function [i,j] = cimgnbmap([nr,nc], nb_r, sample_rate) 3 | * computes the neighbourhood index matrix of an image, 4 | * with each neighbourhood sampled. 5 | * Input: 6 | * [nr,nc] = image size 7 | * nb_r = neighbourhood radius, could be [r_i,r_j] for i,j 8 | * sample_rate = sampling rate, default = 1 9 | * Output: 10 | * [i,j] = each is a column vector, give indices of neighbour pairs 11 | * UINT32 type 12 | * i is of total length of valid elements, 0 for first row 13 | * j is of length nr * nc + 1 14 | * 15 | * See also: imgnbmap.c, id2cind.m 16 | * 17 | * Examples: 18 | * [i,j] = imgnbmap(10, 20); % [10,10] are assumed 19 | * 20 | * Stella X. Yu, Nov 12, 2001. 21 | 22 | % test sequence: 23 | nr = 15; 24 | nc = 15; 25 | nbr = 1; 26 | [i,j] = cimgnbmap([nr,nc], nbr); 27 | mask = csparse(i,j,ones(length(i),1),nr*nc); 28 | show_dist_w(rand(nr,nc),mask) 29 | 30 | *=================================================================*/ 31 | 32 | # include "mex.h" 33 | # include "math.h" 34 | # include 35 | 36 | void mexFunction( 37 | int nargout, 38 | mxArray *out[], 39 | int nargin, 40 | const mxArray *in[] 41 | ) 42 | { 43 | /* declare variables */ 44 | int nr, nc, np, nb, total; 45 | double *dim, sample_rate; 46 | int r_i, r_j, a1, a2, b1, b2, self, neighbor; 47 | int i, j, k, s, t, nsamp, th_rand, no_sample; 48 | unsigned long *p; 49 | 50 | 51 | /* check argument */ 52 | if (nargin < 2) { 53 | mexErrMsgTxt("Two input arguments required"); 54 | } 55 | if (nargout> 2) { 56 | mexErrMsgTxt("Too many output arguments."); 57 | } 58 | 59 | /* get image size */ 60 | i = mxGetM(in[0]); 61 | j = mxGetN(in[0]); 62 | dim = (double *)mxGetData(in[0]); 63 | nr = (int)dim[0]; 64 | if (j>1 || i>1) { 65 | nc = (int)dim[1]; 66 | } else { 67 | nc = nr; 68 | } 69 | np = nr * nc; 70 | 71 | /* get neighbourhood size */ 72 | i = mxGetM(in[1]); 73 | j = mxGetN(in[1]); 74 | dim = (double*)mxGetData(in[1]); 75 | r_i = (int)dim[0]; 76 | if (j>1 || i>1) { 77 | r_j = (int)dim[1]; 78 | } else { 79 | r_j = r_i; 80 | } 81 | if (r_i<0) { r_i = 0; } 82 | if (r_j<0) { r_j = 0; } 83 | 84 | /* get sample rate */ 85 | if (nargin==3) { 86 | sample_rate = (mxGetM(in[2])==0) ? 1: mxGetScalar(in[2]); 87 | } else { 88 | sample_rate = 1; 89 | } 90 | /* prepare for random number generator */ 91 | if (sample_rate<1) { 92 | srand( (unsigned)time( NULL ) ); 93 | th_rand = (int)ceil((double)RAND_MAX * sample_rate); 94 | no_sample = 0; 95 | } else { 96 | sample_rate = 1; 97 | th_rand = RAND_MAX; 98 | no_sample = 1; 99 | } 100 | 101 | /* figure out neighbourhood size */ 102 | 103 | nb = (r_i + r_i + 1) * (r_j + r_j + 1); 104 | if (nb>np) { 105 | nb = np; 106 | } 107 | nb = (int)ceil((double)nb * sample_rate); 108 | 109 | /* intermediate data structure */ 110 | p = (unsigned long *)mxCalloc(np * (nb+1), sizeof(unsigned long)); 111 | if (p==NULL) { 112 | mexErrMsgTxt("Not enough space for my computation."); 113 | } 114 | 115 | /* computation */ 116 | total = 0; 117 | for (j=0; j=nc) { b2 = nc-1; } 130 | 131 | /* i range */ 132 | a1 = i - r_i; 133 | if (a1<0) { a1 = 0; } 134 | a2 = i + r_i; 135 | if (a2>=nr) { a2 = nr-1; } 136 | 137 | /* number of more samples needed */ 138 | nsamp = nb - p[self]; 139 | 140 | k = 0; 141 | t = b1; 142 | s = i + 1; 143 | if (s>a2) { 144 | s = a1; 145 | t = t + 1; 146 | } 147 | while (ka2) { 160 | s = a1; 161 | t = t + 1; 162 | } 163 | } /* k */ 164 | 165 | total = total + p[self]; 166 | } /* i */ 167 | } /* j */ 168 | 169 | /* i, j */ 170 | out[0] = mxCreateNumericMatrix(total, 1, mxUINT32_CLASS, mxREAL); 171 | out[1] = mxCreateNumericMatrix(np+1, 1, mxUINT32_CLASS, mxREAL); 172 | unsigned int *qi = (unsigned int *)mxGetData(out[0]); 173 | unsigned int *qj = (unsigned int *)mxGetData(out[1]); 174 | if (out[0]==NULL || out[1]==NULL) { 175 | mexErrMsgTxt("Not enough space for the output matrix."); 176 | } 177 | 178 | total = 0; 179 | for (j=0; j threshold); 22 | egx1 = g(:,:,1); 23 | egy1 = g(:,:,2); 24 | eindex = find(edges2); 25 | [ey,ex,values] = find(edges2); 26 | 27 | egx = egx1(eindex); 28 | egy = egy1(eindex); 29 | 30 | edgemap.eindex = eindex; 31 | edgemap.values = values; 32 | edgemap.x = ex; 33 | edgemap.y = ey; 34 | edgemap.gx = egx; 35 | edgemap.gy = egy; 36 | edgemap.emag = emag; 37 | edgemap.ephase = ephase; 38 | edgemap.imageEdges = edges2; 39 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/computeW.m: -------------------------------------------------------------------------------- 1 | function W = computeW(imageX,dataW,emag,ephase) 2 | % W = computeW(imageX,dataW,emag,ephase) 3 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 4 | [p,q] = size(imageX); 5 | 6 | [w_i,w_j] = cimgnbmap([p,q],dataW.sampleRadius,dataW.sample_rate); 7 | 8 | W = affinityic(emag,ephase,w_i,w_j,max(emag(:)) * dataW.edgeVariance); 9 | W = W/max(W(:)); 10 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/compute_relation.m: -------------------------------------------------------------------------------- 1 | function [W,distances] = compute_relation(data,scale_sig,order) 2 | % 3 | % [W,distances] = compute_relation(data,scale_sig) 4 | % Input: data= Feature_dimension x Num_data 5 | % ouput: W = pair-wise data similarity matrix 6 | % Dist = pair-wise Euclidean distance 7 | % 8 | % 9 | % Jianbo Shi, 1997 10 | 11 | 12 | distances = zeros(length(data),length(data)); 13 | for j = 1:length(data), 14 | distances(j,:) = (sqrt((data(1,:)-data(1,j)).^2 +... 15 | (data(2,:)-data(2,j)).^2)); 16 | end 17 | 18 | % distances = X2distances(data'); 19 | 20 | if (~exist('scale_sig')), 21 | scale_sig = 0.05*max(distances(:)); 22 | end 23 | 24 | if (~exist('order')), 25 | order = 2; 26 | end 27 | 28 | tmp = (distances/scale_sig).^order; 29 | 30 | W = exp(-tmp); 31 | 32 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/compute_relation2.m: -------------------------------------------------------------------------------- 1 | function [A,B] = compute_relation2(data,scale_sig,order) 2 | % 3 | % [W,Dist] = compute_relation(data,scale_sig) 4 | % Input: data= Feature_dimension x Num_data 5 | % ouput: W = pair-wise data similarity matrix 6 | % Dist = pair-wise Euclidean distance 7 | % 8 | % 9 | % Jianbo Shi, 1997 10 | 11 | if (~exist('order')), 12 | order = 2; 13 | % order = 1; 14 | end 15 | 16 | n = size(data,2); 17 | 18 | B = zeros(n); 19 | for j = 1:n 20 | B(:,j) = (sqrt((data(1,:)-data(1,j)).^2 +... 21 | (data(2,:)-data(2,j)).^2))'; 22 | end 23 | 24 | if (~exist('scale_sig')), 25 | scale_sig = 0.05*max(B(:)); 26 | % scale_sig = 0.04*max(B(:)); 27 | end 28 | 29 | % kNN = 5; 30 | % B2 = B; 31 | % for j=1:n 32 | % [ignore,ind] = sort(B(:,j)); 33 | % B2(ind(kNN+1:end),j) = Inf; 34 | % B2(ind(1:kNN),j) = B(ind(1:kNN),j) / max(B(ind(1:kNN),j)) * 0.5; 35 | % end 36 | % scale_sig = 1; 37 | % B = min(B2,B2');%(B+B')/2; 38 | 39 | tmp = (B/scale_sig).^order; 40 | 41 | A = exp(-tmp); 42 | 43 | 44 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/demm.m: -------------------------------------------------------------------------------- 1 | function a = demm(a,b) 2 | a 3 | b 4 | end -------------------------------------------------------------------------------- /prerprocess/Ncut_9/demoMatrix.asv: -------------------------------------------------------------------------------- 1 | function demoMatrix; 2 | 3 | fileID = fopen('../segment_affinity_matrix.txt','r'); 4 | outputID = fopen('../tree.txt', 'w'); 5 | data = fscanf(fileID, '%f'); 6 | fclose(fileID); 7 | 8 | fileID = fopen('../segment_affinity_matrix.txt','r'); 9 | leafNum=0; 10 | while ~feof(fileID) 11 | str=fgetl(fileID); 12 | if ~isempty(str) 13 | leafNum=leafNum+1; 14 | end 15 | end 16 | fclose(fileID); 17 | 18 | W = reshape(data,leafNum,[]);% yifei 19 | 20 | for i=1:leafNum 21 | for j=1:leafNum 22 | if W(i,j) == 0 23 | W(i,j) = W(i,j); 24 | end 25 | end 26 | end 27 | 28 | disp(W); 29 | 30 | nbCluster = 2; 31 | rootNum = leafNum*2-2; 32 | valid=ones(rootNum,1) 33 | valid(rootNum)=0 34 | 35 | function [id,valid] = findNextID(valid) 36 | id = length(valid); 37 | for i = length(valid):-1:0 38 | if valid(i)==1 39 | id = i; 40 | valid(id)=0; 41 | break 42 | end 43 | end 44 | end 45 | 46 | %%left childs, right childs, weight matrix 47 | function [valid] = buildTree(curRoot, rootList, valid) 48 | 49 | if size(rootList,2) == 3 50 | rootNum = rootNum - 1; 51 | fprintf(outputID, '%d %d null\n', curRoot, rootNum); 52 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(3)-1, rootList(3)-1); 53 | 54 | fprintf(outputID, '%d %d %d\n', rootNum, rootList(1)-1, rootList(1)-1); 55 | fprintf(outputID, '%d %d %d\n', rootNum, rootList(2)-1, rootList(2)-1); 56 | return; 57 | elseif size(rootList,2) == 2 58 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(1)-1, rootList(1)-1); 59 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(2)-1, rootList(2)-1); 60 | return; 61 | end 62 | 63 | compute = false; 64 | 65 | try 66 | [NcutDiscrete,~,~] = ncutW(W(rootList,rootList),nbCluster); 67 | catch 68 | lx = length(rootList); 69 | half = ceil(lx/2); 70 | lcs = rootList(1:half); 71 | rcs = rootList(half+1:end); 72 | lcount = size(lcs, 2); 73 | rcount = size(rcs, 2); 74 | compute = true; 75 | end 76 | 77 | if ~compute 78 | for j=1:nbCluster 79 | id = find(NcutDiscrete(:,j)); 80 | actualId = rootList(id); 81 | 82 | if j==1 83 | lcount = size(actualId, 2); 84 | disp(lcount); 85 | lcs = actualId; 86 | else 87 | rcount = size(actualId, 2); 88 | disp(rcount); 89 | rcs = actualId; 90 | end 91 | end 92 | end 93 | 94 | if lcount > 1 && rcount > 1 95 | [leftID,valid] = findNextID(valid); 96 | fprintf(outputID, '%d %d null\n', curRoot, leftID); 97 | [rightID,valid] = findNextID(valid); 98 | fprintf(outputID, '%d %d null\n', curRoot, rightID); 99 | valid = buildTree(leftID, lcs, valid); 100 | valid = buildTree(rightID, rcs, valid); 101 | elseif lcount == 1 && rcount == 1 102 | fprintf(outputID, '%d %d %d\n', curRoot, lcs(1)-1, lcs(1)-1); 103 | fprintf(outputID, '%d %d %d\n', curRoot, rcs(1)-1, rcs(1)-1); 104 | elseif lcount == 1 && rcount ~= 1 105 | rootNum = rootNum - 1; 106 | fprintf(outputID, '%d %d %d\n', curRoot, lcs(1)-1, lcs(1)-1); 107 | [rightID,valid] = findNextID(valid); 108 | fprintf(outputID, '%d %d null\n', curRoot, rightID); 109 | buildTree(rightID, rcs, valid); 110 | else 111 | rootNum = rootNum - 1; 112 | fprintf(outputID, '%d %d %d\n', curRoot, rcs(1)-1, rcs(1)-1); 113 | [leftID,valid] = findNextID(valid); 114 | fprintf(outputID, '%d %d null\n', curRoot, leftID); 115 | buildTree(leftID, lcs, valid); 116 | end 117 | 118 | end 119 | 120 | buildTree(rootNum, (1:leafNum), valid)% yifei 121 | disp(rootNum) 122 | end 123 | 124 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/demoMatrix.m: -------------------------------------------------------------------------------- 1 | function output = demoMatrix(affinity_dir,hier_dir); 2 | disp(affinity_dir) 3 | disp(hier_dir) 4 | fileID = fopen(affinity_dir,'r'); 5 | outputID = fopen(hier_dir, 'w'); 6 | data = fscanf(fileID, '%f'); 7 | fclose(fileID); 8 | 9 | fileID = fopen(affinity_dir,'r'); 10 | leafNum=0; 11 | while ~feof(fileID) 12 | str=fgetl(fileID); 13 | if ~isempty(str) 14 | leafNum=leafNum+1; 15 | end 16 | end 17 | fclose(fileID); 18 | 19 | W = reshape(data,leafNum,[]);% yifei 20 | 21 | for i=1:leafNum 22 | for j=1:leafNum 23 | %if W(i,j) < 0.01 24 | % W(i,j) = rand(1)*2; 25 | %else 26 | % W(i,j) = W(i,j) * 100; 27 | %end 28 | W(i,j) = W(i,j); 29 | end 30 | end 31 | 32 | %disp(W); 33 | 34 | nbCluster = 2; 35 | rootNum = leafNum*2-2; 36 | valid=ones(rootNum,1); 37 | valid(rootNum)=0; 38 | fprintf('leafNum %d rootNum %d\n',leafNum, rootNum) 39 | 40 | function [id,valid] = findNextID(valid) 41 | id = length(valid); 42 | for i = length(valid):-1:0 43 | if valid(i)==1 44 | id = i; 45 | valid(id)=0; 46 | break 47 | end 48 | end 49 | end 50 | 51 | %%left childs, right childs, weight matrix 52 | function [valid] = buildTree(curRoot, rootList, valid) 53 | if size(rootList,2) == 3 54 | [id,valid] = findNextID(valid); 55 | fprintf(outputID, '%d %d null\n', curRoot, id); 56 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(3)-1, rootList(3)-1); 57 | 58 | fprintf(outputID, '%d %d %d\n', id, rootList(1)-1, rootList(1)-1); 59 | fprintf(outputID, '%d %d %d\n', id, rootList(2)-1, rootList(2)-1); 60 | return; 61 | elseif size(rootList,2) == 2 62 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(1)-1, rootList(1)-1); 63 | fprintf(outputID, '%d %d %d\n', curRoot, rootList(2)-1, rootList(2)-1); 64 | return; 65 | end 66 | 67 | compute = false; 68 | 69 | try 70 | [NcutDiscrete,~,~] = ncutW(W(rootList,rootList),nbCluster); 71 | catch 72 | lx = length(rootList); 73 | half = ceil(lx/2); 74 | lcs = rootList(1:half); 75 | rcs = rootList(half+1:end); 76 | lcount = size(lcs, 2); 77 | rcount = size(rcs, 2); 78 | compute = true; 79 | if lcount > 10 80 | fprintf('equally split left: %d right: %d\n',lcount, rcount) 81 | end 82 | end 83 | 84 | if ~compute 85 | for j=1:nbCluster 86 | id = find(NcutDiscrete(:,j)); 87 | actualId = rootList(id); 88 | if j==1 89 | lcount = size(actualId, 2); 90 | lcs = actualId; 91 | else 92 | rcount = size(actualId, 2); 93 | rcs = actualId; 94 | end 95 | end 96 | %fprintf('left: %d right: %d\n',lcount, rcount) 97 | end 98 | 99 | if lcount == 0 || rcount == 0 100 | fprintf('zero left: %d right: %d\n',lcount, rcount) 101 | lx = length(rootList); 102 | half = ceil(lx/2); 103 | lcs = rootList(1:half); 104 | rcs = rootList(half+1:end); 105 | lcount = size(lcs, 2); 106 | rcount = size(rcs, 2); 107 | fprintf('corrected left: %d right: %d\n',lcount, rcount) 108 | end 109 | 110 | if lcount > 1 && rcount > 1 111 | [leftID,valid] = findNextID(valid); 112 | fprintf(outputID, '%d %d null\n', curRoot, leftID); 113 | [rightID,valid] = findNextID(valid); 114 | fprintf(outputID, '%d %d null\n', curRoot, rightID); 115 | valid = buildTree(leftID, lcs, valid); 116 | valid = buildTree(rightID, rcs, valid); 117 | elseif lcount == 1 && rcount == 1 118 | fprintf(outputID, '%d %d %d\n', curRoot, lcs(1)-1, lcs(1)-1); 119 | fprintf(outputID, '%d %d %d\n', curRoot, rcs(1)-1, rcs(1)-1); 120 | elseif lcount == 1 && rcount ~= 1 121 | fprintf(outputID, '%d %d %d\n', curRoot, lcs(1)-1, lcs(1)-1); 122 | [rightID,valid] = findNextID(valid); 123 | fprintf(outputID, '%d %d null\n', curRoot, rightID); 124 | valid = buildTree(rightID, rcs, valid); 125 | elseif lcount ~= 1 && rcount == 1 126 | fprintf(outputID, '%d %d %d\n', curRoot, rcs(1)-1, rcs(1)-1); 127 | [leftID,valid] = findNextID(valid); 128 | fprintf(outputID, '%d %d null\n', curRoot, leftID); 129 | valid = buildTree(leftID, lcs, valid); 130 | else 131 | 132 | end 133 | 134 | end 135 | 136 | buildTree(rootNum, (1:leafNum), valid);% yifei 137 | output = 1; 138 | %disp(rootNum); 139 | end 140 | 141 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/demoNcutClustering.m: -------------------------------------------------------------------------------- 1 | function demoNcutClustering; 2 | % demoNcutClustering 3 | % 4 | % demo for NcutClustering 5 | % also initialize matlab paths to subfolders 6 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 7 | 8 | disp('Ncut Clustering demo'); 9 | 10 | %% make up a point data set 11 | caseid = 2;%3 12 | [data,size_cluster] = build_scene(caseid); 13 | figure(1);clf; 14 | plot(data(1,:),data(2,:),'ks', 'MarkerFaceColor','k','MarkerSize',5); axis image; hold on; 15 | 16 | disp('This is the input data points to be clustered, press Enter to continue...'); 17 | pause; 18 | 19 | disp('Compute clustering...'); 20 | 21 | % compute similarity matrix 22 | [W,Dist] = compute_relation(data); 23 | 24 | % clustering graph in 25 | nbCluster = 4; 26 | tic; 27 | [NcutDiscrete,NcutEigenvectors,NcutEigenvalues] = ncutW(W,nbCluster); 28 | disp(['The computation took ' num2str(toc) ' seconds']); 29 | figure(3); 30 | plot(NcutEigenvectors); 31 | disp(NcutDiscrete) 32 | 33 | % display clustering result 34 | cluster_color = ['rgbmyc']; 35 | figure(2);clf; 36 | for j=1:nbCluster, 37 | id = find(NcutDiscrete(:,j)); 38 | plot(data(1,id),data(2,id),[cluster_color(j),'s'], 'MarkerFaceColor',cluster_color(j),'MarkerSize',5); hold on; 39 | end 40 | hold off; axis image; 41 | disp('This is the clustering result'); 42 | disp('The demo is finished.'); 43 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/demoNcutImage.m: -------------------------------------------------------------------------------- 1 | function demoNcutImage; 2 | % demoNcutImage 3 | % 4 | % demo for NcutImage 5 | % also initialize matlab paths to subfolders 6 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 7 | 8 | disp('Ncut Image Segmentation demo'); 9 | 10 | %% read image, change color image to brightness image, resize to 160x160 11 | I = imread_ncut('jpg_images/3.jpg',160,160); 12 | 13 | %% display the image 14 | figure(1);clf; imagesc(I);colormap(gray);axis off; 15 | disp('This is the input image to segment, press Enter to continue...'); 16 | pause; 17 | 18 | %% compute the edges imageEdges, the similarity matrix W based on 19 | %% Intervening Contours, the Ncut eigenvectors and discrete segmentation 20 | nbSegments = 5; 21 | disp('computing Ncut eigenvectors ...'); 22 | tic; 23 | [SegLabel,NcutDiscrete,NcutEigenvectors,NcutEigenvalues,W,imageEdges]= NcutImage(I,nbSegments); 24 | disp(['The computation took ' num2str(toc) ' seconds on the ' num2str(size(I,1)) 'x' num2str(size(I,2)) ' image']); 25 | 26 | 27 | %% display the edges 28 | figure(2);clf; imagesc(imageEdges); axis off 29 | disp('This is the edges computed, press Enter to continue...'); 30 | pause; 31 | 32 | %% display the segmentation 33 | figure(3);clf 34 | bw = edge(SegLabel,0.01); 35 | J1=showmask(I,imdilate(bw,ones(2,2))); imagesc(J1);axis off 36 | disp('This is the segmentation, press Enter to continue...'); 37 | pause; 38 | 39 | %% display Ncut eigenvectors 40 | figure(4);clf;set(gcf,'Position',[100,500,200*(nbSegments+1),200]); 41 | [nr,nc,nb] = size(I); 42 | for i=1:nbSegments 43 | subplot(1,nbSegments,i); 44 | imagesc(reshape(NcutEigenvectors(:,i) , nr,nc));axis('image');axis off; 45 | end 46 | disp('This is the Ncut eigenvectors...'); 47 | disp('The demo is finished.'); 48 | 49 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/discretisation.m: -------------------------------------------------------------------------------- 1 | function [EigenvectorsDiscrete,EigenVectors]=discretisation(EigenVectors) 2 | % 3 | % EigenvectorsDiscrete=discretisation(EigenVectors) 4 | % 5 | % Input: EigenVectors = continuous Ncut vector, size = ndata x nbEigenvectors 6 | % Output EigenvectorsDiscrete = discrete Ncut vector, size = ndata x nbEigenvectors 7 | % 8 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004 9 | 10 | [n,k]=size(EigenVectors); 11 | 12 | vm = sqrt(sum(EigenVectors.*EigenVectors,2)); 13 | EigenVectors = EigenVectors./repmat(vm,1,k); 14 | 15 | R=zeros(k); 16 | R(:,1)=EigenVectors(1+round(rand(1)*(n-1)),:)'; 17 | c=zeros(n,1); 18 | for j=2:k 19 | c=c+abs(EigenVectors*R(:,j-1)); 20 | [minimum,i]=min(c); 21 | R(:,j)=EigenVectors(i,:)'; 22 | end 23 | 24 | lastObjectiveValue=0; 25 | exitLoop=0; 26 | nbIterationsDiscretisation = 0; 27 | nbIterationsDiscretisationMax = 20;%voir 28 | while exitLoop== 0 29 | nbIterationsDiscretisation = nbIterationsDiscretisation + 1 ; 30 | EigenvectorsDiscrete = discretisationEigenVectorData(EigenVectors*R); 31 | [U,S,V] = svd(EigenvectorsDiscrete'*EigenVectors,0); 32 | NcutValue=2*(n-trace(S)); 33 | 34 | if abs(NcutValue-lastObjectiveValue) < eps | nbIterationsDiscretisation > nbIterationsDiscretisationMax 35 | exitLoop=1; 36 | else 37 | lastObjectiveValue = NcutValue; 38 | R=V*U'; 39 | end 40 | end -------------------------------------------------------------------------------- /prerprocess/Ncut_9/discretisationEigenVectorData.m: -------------------------------------------------------------------------------- 1 | function Y = discretisationEigenVectorData(EigenVector) 2 | % Y = discretisationEigenVectorData(EigenVector) 3 | % 4 | % discretizes previously rotated eigenvectors in discretisation 5 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004 6 | 7 | [n,k]=size(EigenVector); 8 | 9 | 10 | [Maximum,J]=max(EigenVector'); 11 | 12 | Y=sparse(1:n,J',1,n,k); 13 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/doog1.m: -------------------------------------------------------------------------------- 1 | function H=doog1(sig,r,th,N); 2 | % H=doog1(sig,r,th,N); 3 | % Serge Belongie 4 | 5 | no_pts=N; % no. of points in x,y grid 6 | 7 | [x,y]=meshgrid(-(N/2)+1/2:(N/2)-1/2,-(N/2)+1/2:(N/2)-1/2); 8 | 9 | phi=pi*th/180; 10 | sigy=sig; 11 | sigx=r*sig; 12 | R=[cos(phi) -sin(phi); sin(phi) cos(phi)]; 13 | C=R*diag([sigx,sigy])*R'; 14 | 15 | X=[x(:) y(:)]; 16 | 17 | Gb=gaussian(X,[0 0]',C); 18 | Gb=reshape(Gb,N,N); 19 | 20 | m=R*[0 sig]'; 21 | 22 | a=1; 23 | b=-1; 24 | 25 | % make odd-symmetric filter 26 | Ga=gaussian(X,m/2,C); 27 | Ga=reshape(Ga,N,N); 28 | Gb=rot90(Ga,2); 29 | H=a*Ga+b*Gb; 30 | 31 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/doog2.m: -------------------------------------------------------------------------------- 1 | function G=doog2(sig,r,th,N); 2 | % G=doog2(sig,r,th,N); 3 | % Make difference of offset gaussians kernel 4 | % theta is in degrees 5 | % (see Malik & Perona, J. Opt. Soc. Amer., 1990) 6 | % 7 | % Example: 8 | % >> imagesc(doog2(1,12,0,64,1)) 9 | % >> colormap(gray) 10 | % 11 | % Serge Belongie 12 | 13 | 14 | no_pts=N; % no. of points in x,y grid 15 | 16 | [x,y]=meshgrid(-(N/2)+1/2:(N/2)-1/2,-(N/2)+1/2:(N/2)-1/2); 17 | 18 | phi=pi*th/180; 19 | sigy=sig; 20 | sigx=r*sig; 21 | R=[cos(phi) -sin(phi); sin(phi) cos(phi)]; 22 | C=R*diag([sigx,sigy])*R'; 23 | 24 | X=[x(:) y(:)]; 25 | 26 | Gb=gaussian(X,[0 0]',C); 27 | Gb=reshape(Gb,N,N); 28 | 29 | m=R*[0 sig]'; 30 | Ga=gaussian(X,m,C); 31 | Ga=reshape(Ga,N,N); 32 | Gc=rot90(Ga,2); 33 | 34 | a=-1; 35 | b=2; 36 | c=-1; 37 | 38 | G = a*Ga + b*Gb + c*Gc; 39 | 40 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/fft_filt_2.m: -------------------------------------------------------------------------------- 1 | function FI=fft_filt_2(V,FB,sf); 2 | % FI=fft_filt_2(V,FB,sf); 3 | % fft-based filtering 4 | % requires image to be called "V" 5 | % and filters to be in FB 6 | % sf is the subsampling factor 7 | % 8 | % FI is the result 9 | % Jianbo Shi, 1997 10 | 11 | [M1,M2,N3]=size(FB); 12 | % prepare FFT of image for filtering 13 | [N1,N2]=size(V); 14 | I=zeros(size(V)+[M1-1 M2-1]); 15 | I(1:N1,1:N2)=V; 16 | N1s=length(1:sf:N1); 17 | N2s=length(1:sf:N2); 18 | IF=fft2(I); 19 | FI=zeros(N1s,N2s,N3); 20 | 21 | % apply filters 22 | for n=1:N3; 23 | f=rot90(FB(:,:,n),2); 24 | fF=fft2(f,N1+M1-1,N2+M2-1); 25 | IfF=IF.*fF; 26 | If=real(ifft2(IfF)); 27 | If=If(ceil((M1+1)/2):ceil((M1+1)/2)+N1-1,ceil((M2+1)/2):ceil((M2+1)/2)+N2-1); 28 | FI(:,:,n)=If(1:sf:N1,1:sf:N2); 29 | end 30 | 31 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/gaussian.m: -------------------------------------------------------------------------------- 1 | function p=gaussian(x,m,C); 2 | % p=gaussian(x,m,C); 3 | % 4 | % Evaluate the multi-variate density with mean vector m and covariance 5 | % matrix C for the input vector x. 6 | % 7 | % p=gaussian(X,m,C); 8 | % 9 | % Vectorized version: Here X is a matrix of column vectors, and p is 10 | % a vector of probabilities for each vector. 11 | % Jianbo Shi, 1997 12 | d=length(m); 13 | 14 | if size(x,1)~=d 15 | x=x'; 16 | end 17 | N=size(x,2); 18 | 19 | detC = det(C); 20 | if rcond(C)1), 13 | I =double(rgb2gray(I)); 14 | else 15 | I = double(I); 16 | end 17 | 18 | I = imresize(I,[nr, nc],'bicubic'); 19 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/jpg_images/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/jpg_images/1.jpg -------------------------------------------------------------------------------- /prerprocess/Ncut_9/jpg_images/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/jpg_images/2.jpg -------------------------------------------------------------------------------- /prerprocess/Ncut_9/jpg_images/3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/jpg_images/3.jpg -------------------------------------------------------------------------------- /prerprocess/Ncut_9/jpg_images/Thumbs.db: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/jpg_images/Thumbs.db -------------------------------------------------------------------------------- /prerprocess/Ncut_9/make_filterbank_even2.m: -------------------------------------------------------------------------------- 1 | function FB = make_filterbank(num_ori,filter_scales,wsz,enlong) 2 | % 3 | % F = make_filterbank(num_ori,num_scale,wsz) 4 | % Jianbo Shi, 1997 5 | 6 | 7 | if nargin<4, 8 | enlong = 3; 9 | end 10 | 11 | enlong = 2*enlong; 12 | 13 | % definine filterbank 14 | %num_ori=6; 15 | %num_scale=3; 16 | 17 | num_scale = length(filter_scales); 18 | 19 | M1=wsz; % size in pixels 20 | M2=M1; 21 | 22 | ori_incr=180/num_ori; 23 | ori_offset=ori_incr/2; % helps with equalizing quantiz. error across filter set 24 | 25 | FB=zeros(M1,M2,num_ori,num_scale); 26 | 27 | % elongated filter set 28 | counter = 1; 29 | 30 | for m=1:num_scale 31 | for n=1:num_ori 32 | % r=12 here is equivalent to Malik's r=3; 33 | f=doog2(filter_scales(m),enlong,ori_offset+(n-1)*ori_incr,M1); 34 | FB(:,:,n,m)=f; 35 | end 36 | end 37 | 38 | FB=reshape(FB,M1,M2,num_scale*num_ori); 39 | total_num_filt=size(FB,3); 40 | 41 | for j=1:total_num_filt, 42 | F = FB(:,:,j); 43 | a = sum(sum(abs(F))); 44 | FB(:,:,j) = FB(:,:,j)/a; 45 | end 46 | 47 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/make_filterbank_odd2.m: -------------------------------------------------------------------------------- 1 | function FB = make_filterbank(num_ori,filter_scales,wsz,enlong) 2 | % 3 | % F = make_filterbank(num_ori,num_scale,wsz) 4 | % Jianbo Shi, 1997 5 | 6 | if nargin<4, 7 | enlong = 3; 8 | end 9 | 10 | enlong = enlong*2; 11 | 12 | % definine filterbank 13 | %num_ori=6; 14 | %num_scale=3; 15 | 16 | num_scale = length(filter_scales); 17 | 18 | M1=wsz; % size in pixels 19 | M2=M1; 20 | 21 | ori_incr=180/num_ori; 22 | ori_offset=ori_incr/2; % helps with equalizing quantiz. error across filter set 23 | 24 | FB=zeros(M1,M2,num_ori,num_scale); 25 | 26 | 27 | % elongated filter set 28 | counter = 1; 29 | 30 | for m=1:num_scale 31 | for n=1:num_ori 32 | % r=12 here is equivalent to Malik's r=3; 33 | f=doog1(filter_scales(m),enlong,ori_offset+(n-1)*ori_incr,M1); 34 | FB(:,:,n,m)=f; 35 | end 36 | end 37 | 38 | FB=reshape(FB,M1,M2,num_scale*num_ori); 39 | total_num_filt=size(FB,3); 40 | 41 | for j=1:total_num_filt, 42 | F = FB(:,:,j); 43 | a = sum(sum(abs(F))); 44 | FB(:,:,j) = FB(:,:,j)/a; 45 | end 46 | 47 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/mex_w_times_x_symmetric.cpp: -------------------------------------------------------------------------------- 1 | /*================================================================ 2 | * mex_w_times_x_symmetric.c = used by ncuthard2.m in eigensolver. 3 | * 4 | * Examples: 5 | * mex_w_times_x_c_symmetric(x,triu(A)) = A*x; 6 | * A is sparse and symmetric, but x is full 7 | * 8 | * Timothee Cour, Oct 12, 2003. 9 | 10 | % test sequence: 11 | m=100; 12 | n=50; 13 | x=rand(n,1); 14 | A=sprand(m,n,0.01); 15 | 16 | y2 = mex_w_times_x_c_symmetric(x,triu(A)); 17 | y1=A*x; 18 | max(abs(y1-y2)) 19 | *=================================================================*/ 20 | 21 | # include "math.h" 22 | # include "mex.h" 23 | # include "a_times_b_cmplx.cpp" 24 | /*# include "a_times_b.c"*/ 25 | 26 | 27 | void mexFunction( 28 | int nargout, 29 | mxArray *out[], 30 | int nargin, 31 | const mxArray *in[] 32 | ) 33 | { 34 | int np, nc; 35 | mwIndex*ir, *jc; 36 | double *x, *y, *pr; 37 | 38 | if (nargin < 2) {//voir 39 | mexErrMsgTxt("Four input arguments required !"); 40 | } 41 | if (nargout>1) { 42 | mexErrMsgTxt("Too many output arguments."); 43 | } 44 | 45 | x = mxGetPr(in[0]); 46 | pr = mxGetPr(in[1]); 47 | ir = mxGetIr(in[1]); 48 | jc = mxGetJc(in[1]); 49 | 50 | np = mxGetM(in[1]); 51 | nc = mxGetN(in[1]); 52 | 53 | out[0] = mxCreateDoubleMatrix(np,1,mxREAL); 54 | y = mxGetPr(out[0]); 55 | 56 | CSRsymm_VecMult_CAB_double(np,nc,pr,ir,jc,x,y); 57 | } 58 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/mex_w_times_x_symmetric.mexa64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/mex_w_times_x_symmetric.mexa64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/mex_w_times_x_symmetric.mexmaci64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/mex_w_times_x_symmetric.mexmaci64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/mex_w_times_x_symmetric.mexw64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/mex_w_times_x_symmetric.mexw64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/ncut.m: -------------------------------------------------------------------------------- 1 | function [Eigenvectors,Eigenvalues] = ncut(W,nbEigenValues,dataNcut); 2 | % function [Eigenvectors,Eigenvalues] = ncut(W,nbEigenValues,dataNcut); 3 | % 4 | % Input: 5 | % W= symmetric similarity matrix 6 | % nbEigenValues= number of Ncut eigenvectors computed 7 | % dataNcut= optional parameters 8 | % 9 | % default parameters for dataNcut: 10 | % dataNcut.offset = 5e-1; offset in the diagonal of W 11 | % dataNcut.verbose = 0; 0 for verbose off mode, 1,2,3 for verbose on modes 12 | % dataNcut.maxiterations = 100; max number of iterations in eigensolver 13 | % dataNcut.eigsErrorTolerance = 1e-6; error tolerance in eigensolver 14 | % dataNcut.valeurMin=1e-6; % truncates any values in W less than valeurMin 15 | % 16 | % Output: 17 | % Eigenvectors= continuouse Ncut eigenvectos, size = length(W) x nbEigenValues 18 | % Eigenvalues= Ncut eigenvalues, size = 1x nbEigenValues 19 | % 20 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004. 21 | 22 | if nargin < 2 23 | nbEigenValues = 8; 24 | end 25 | if nargin < 3 26 | dataNcut.offset = 5e-1; 27 | dataNcut.verbose = 0; 28 | dataNcut.maxiterations = 300; 29 | dataNcut.eigsErrorTolerance = 1e-8; 30 | dataNcut.valeurMin=1e-6; 31 | end 32 | % if nargin < 3 33 | % dataNcut.offset = 5e-1; 34 | % dataNcut.verbose = 0; 35 | % dataNcut.maxiterations = 100; 36 | % dataNcut.eigsErrorTolerance = 1e-6; 37 | % dataNcut.valeurMin=1e-6; 38 | % end 39 | 40 | % make W matrix sparse 41 | W = sparsifyc(W,dataNcut.valeurMin); 42 | 43 | % check for matrix symmetry 44 | if max(max(abs(W-W'))) > 1e-10 %voir (-12) 45 | %disp(max(max(abs(W-W')))); 46 | error('W not symmetric'); 47 | end 48 | 49 | n = size(W,1); 50 | nbEigenValues = min(nbEigenValues,n); 51 | offset = dataNcut.offset; 52 | 53 | 54 | % degrees and regularization 55 | d = sum(abs(W),2); 56 | dr = 0.5 * (d - sum(W,2)); 57 | d = d + offset * 2; 58 | dr = dr + offset; 59 | W = W + spdiags(dr,0,n,n); 60 | 61 | Dinvsqrt = 1./sqrt(d+eps); 62 | P = spmtimesd(W,Dinvsqrt,Dinvsqrt); 63 | clear W; 64 | 65 | options.issym = 1; 66 | 67 | if dataNcut.verbose 68 | options.disp = 3; 69 | else 70 | options.disp = 0; 71 | end 72 | options.maxit = dataNcut.maxiterations; 73 | options.tol = dataNcut.eigsErrorTolerance; 74 | 75 | options.v0 = ones(size(P,1),1); 76 | options.p = max(35,2*nbEigenValues); %voir 77 | options.p = min(options.p,n); 78 | 79 | %warning off 80 | % [vbar,s,convergence] = eigs2(@mex_w_times_x_symmetric,size(P,1),nbEigenValues,'LA',options,tril(P)); 81 | [vbar,s,convergence] = eigs(@mex_w_times_x_symmetric,size(P,1),nbEigenValues,'LA',options,tril(P)); 82 | %warning on 83 | 84 | s = real(diag(s)); 85 | [x,y] = sort(-s); 86 | Eigenvalues = -x; 87 | vbar = vbar(:,y); 88 | Eigenvectors = spdiags(Dinvsqrt,0,n,n) * vbar; 89 | 90 | for i=1:size(Eigenvectors,2) 91 | Eigenvectors(:,i) = (Eigenvectors(:,i) / norm(Eigenvectors(:,i)) )*norm(ones(n,1)); 92 | if Eigenvectors(1,i)~=0 93 | Eigenvectors(:,i) = - Eigenvectors(:,i) * sign(Eigenvectors(1,i)); 94 | end 95 | end 96 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/ncutW.m: -------------------------------------------------------------------------------- 1 | function [NcutDiscrete,NcutEigenvectors,NcutEigenvalues] = ncutW(W,nbcluster); 2 | % [NcutDiscrete,NcutEigenvectors,NcutEigenvalues] = ncutW(W,nbcluster); 3 | % 4 | % Calls ncut to compute NcutEigenvectors and NcutEigenvalues of W with nbcluster clusters 5 | % Then calls discretisation to discretize the NcutEigenvectors into NcutDiscrete 6 | % Timothee Cour, Stella Yu, Jianbo Shi, 2004 7 | 8 | % compute continuous Ncut eigenvectors 9 | [NcutEigenvectors,NcutEigenvalues] = ncut(W,nbcluster); 10 | %disp('NcutDiscrete') 11 | %disp(NcutDiscrete) 12 | %disp('NcutEigenvalues') 13 | %disp(NcutEigenvalues) 14 | % compute discretize Ncut vectors 15 | [NcutDiscrete,NcutEigenvectors] =discretisation(NcutEigenvectors); 16 | 17 | 18 | NcutDiscrete = full(NcutDiscrete); -------------------------------------------------------------------------------- /prerprocess/Ncut_9/quadedgep.m: -------------------------------------------------------------------------------- 1 | % function [x,y,gx,gy,par,threshold,mag,mage,g,FIe,FIo,mago] = quadedgep(I,par,threshold); 2 | % Input: 3 | % I = image 4 | % par = vector for 4 parameters 5 | % [number of filter orientations, number of scales, filter size, elongation] 6 | % To use default values, put 0. 7 | % threshold = threshold on edge strength 8 | % Output: 9 | % [x,y,gx,gy] = locations and gradients of an ordered list of edgels 10 | % x,y could be horizontal or vertical or 45 between pixel sites 11 | % but it is guaranteed that there [floor(y) + (floor(x)-1)*nr] 12 | % is ordered and unique. In other words, each edgel has a unique pixel id. 13 | % par = actual par used 14 | % threshold = actual threshold used 15 | % mag = edge magnitude 16 | % mage = phase map 17 | % g = gradient map at each pixel 18 | % [FIe,FIo] = odd and even filter outputs 19 | % mago = odd filter output of optimum orientation 20 | % 21 | % Stella X. Yu, 2001 22 | 23 | 24 | 25 | function [x,y,gx,gy,par,threshold,mag,mage,g,FIe,FIo,mago] = quadedgep(I,par,threshold); 26 | 27 | if nargin<3 | isempty(threshold), 28 | threshold = 0.2; 29 | end 30 | 31 | [r,c] = size(I); 32 | def_par = [8,1,20,3]; 33 | 34 | % take care of parameters, any missing value is substituted by a default value 35 | if nargin<2 | isempty(par), 36 | par = def_par; 37 | end 38 | par(end+1:4)=0; 39 | par = par(:); 40 | j = (par>0); 41 | have_value = [ j, 1-j ]; 42 | j = 1; n_filter = have_value(j,:) * [par(j); def_par(j)]; 43 | j = 2; n_scale = have_value(j,:) * [par(j); def_par(j)]; 44 | j = 3; winsz = have_value(j,:) * [par(j); def_par(j)]; 45 | j = 4; enlong = have_value(j,:) * [par(j); def_par(j)]; 46 | 47 | % always make filter size an odd number so that the results will not be skewed 48 | j = winsz/2; 49 | if not(j > fix(j) + 0.1), 50 | winsz = winsz + 1; 51 | end 52 | 53 | % filter the image with quadrature filters 54 | FBo = make_filterbank_odd2(n_filter,n_scale,winsz,enlong); 55 | FBe = make_filterbank_even2(n_filter,n_scale,winsz,enlong); 56 | n = ceil(winsz/2); 57 | f = [fliplr(I(:,2:n+1)), I, fliplr(I(:,c-n:c-1))]; 58 | f = [flipud(f(2:n+1,:)); f; flipud(f(r-n:r-1,:))]; 59 | FIo = fft_filt_2(f,FBo,1); 60 | FIo = FIo(n+[1:r],n+[1:c],:); 61 | FIe = fft_filt_2(f,FBe,1); 62 | FIe = FIe(n+[1:r],n+[1:c],:); 63 | 64 | % compute the orientation energy and recover a smooth edge map 65 | % pick up the maximum energy across scale and orientation 66 | % even filter's output: as it is the second derivative, zero cross localize the edge 67 | % odd filter's output: orientation 68 | mag = sqrt(sum(FIo.^2,3)+sum(FIe.^2,3)); 69 | mag_a = sqrt(FIo.^2+FIe.^2); 70 | [tmp,max_id] = max(mag_a,[],3); 71 | base_size = r * c; 72 | id = [1:base_size]'; 73 | mage = reshape(FIe(id+(max_id(:)-1)*base_size),[r,c]); 74 | mage = (mage>0) - (mage<0); 75 | 76 | ori_incr=pi/n_filter; % to convert jshi's coords to conventional image xy 77 | ori_offset=ori_incr/2; 78 | theta = ori_offset+([1:n_filter]-1)*ori_incr; % orientation detectors 79 | % [gx,gy] are image gradient in image xy coords, winner take all 80 | mago = reshape(FIo(id+(max_id(:)-1)*base_size),[r,c]); 81 | ori = theta(max_id); 82 | ori = ori .* (mago>0) + (ori + pi).*(mago<0); 83 | gy = mag .* cos(ori); 84 | gx = -mag .* sin(ori); 85 | g = cat(3,gx,gy); 86 | 87 | % phase map: edges are where the phase changes 88 | mag_th = max(mag(:)) * threshold; 89 | eg = (mag>mag_th); 90 | h = eg & [(mage(2:r,:) ~= mage(1:r-1,:)); zeros(1,c)]; 91 | v = eg & [(mage(:,2:c) ~= mage(:,1:c-1)), zeros(r,1)]; 92 | [y,x] = find(h | v); 93 | k = y + (x-1) * r; 94 | h = h(k); 95 | v = v(k); 96 | y = y + h * 0.5; % i 97 | x = x + v * 0.5; % j 98 | t = h + v * r; 99 | gx = g(k) + g(k+t); 100 | k = k + (r * c); 101 | gy = g(k) + g(k+t); 102 | 103 | 104 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/showmask.m: -------------------------------------------------------------------------------- 1 | function RGB=showmask(V,M,display_flag); 2 | % showmask(V,M); 3 | % 4 | % M is a nonneg. mask 5 | % Jianbo Shi, 1997 6 | 7 | V=V-min(V(:)); 8 | V=V/max(V(:)); 9 | V=.25+0.75*V; %brighten things up a bit 10 | 11 | M=M-min(M(:)); 12 | M=M/max(M(:)); 13 | 14 | H=0.0+zeros(size(V)); 15 | S=M; 16 | RGB=hsv2rgb(H,S,V); 17 | 18 | %if nargin>2 19 | image(RGB) 20 | axis('image') 21 | %end 22 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/sparsifyc.cpp: -------------------------------------------------------------------------------- 1 | /*================================================================= 2 | * syntax: SPMX = SPARSIFY(MX, THRES) 3 | * 4 | * SPARSIFY - sparsify the input matrix, i.e. ignore the values 5 | * of the matrix which are below a threshold 6 | * 7 | * Input: - MX: m-by-n matrix (sparse or full) 8 | * - THRES: threshold value (double) 9 | * 10 | * Output: - SPMX: m-by-n sparse matrix only with values 11 | * whose absolut value is above the given threshold 12 | * 13 | * Written by Mirko Visontai (10/24/2003) 14 | *=================================================================*/ 15 | 16 | 17 | #include 18 | #include "mex.h" 19 | 20 | void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) 21 | { 22 | /* Declare variable */ 23 | int i,m,n,nzmax,newnnz,col,processed,passed; 24 | int starting_row_index, current_row_index, stopping_row_index; 25 | double *in_pr,*in_pi,*out_pr,*out_pi; 26 | mwIndex *in_ir,*in_jc,*out_ir,*out_jc; 27 | double thres; 28 | 29 | /* Check for proper number of input and output arguments */ 30 | if ((nlhs != 1) || (nrhs != 2)){ 31 | mexErrMsgTxt("usage: SPMX = SPARSIFY(MX, THRES)."); 32 | } 33 | /* if matrix is complex threshold the norm of the numbers */ 34 | if (mxIsComplex(prhs[0])){ 35 | /* Check data type of input argument */ 36 | if (mxIsSparse(prhs[0])){ 37 | 38 | /* read input */ 39 | in_pr = mxGetPr(prhs[0]); 40 | in_pi = mxGetPi(prhs[0]); 41 | in_ir = mxGetIr(prhs[0]); 42 | in_jc = mxGetJc(prhs[0]); 43 | nzmax = mxGetNzmax(prhs[0]); 44 | m = mxGetM(prhs[0]); 45 | n = mxGetN(prhs[0]); 46 | thres = mxGetScalar(prhs[1]); 47 | 48 | /* Count new nonzeros */ 49 | newnnz=0; 50 | for(i=0; ithres) {newnnz++;} 52 | } 53 | 54 | if (newnnz>0){ 55 | /* create output */ 56 | plhs[0] = mxCreateSparse(m,n,newnnz,mxCOMPLEX); 57 | if (plhs[0]==NULL) 58 | mexErrMsgTxt("Could not allocate enough memory!\n"); 59 | out_pr = mxGetPr(plhs[0]); 60 | out_pi = mxGetPr(plhs[0]); 61 | out_ir = mxGetIr(plhs[0]); 62 | out_jc = mxGetJc(plhs[0]); 63 | passed = 0; 64 | out_jc[0] = 0; 65 | for (col=0; col thres){ 77 | 78 | out_pr[passed]=in_pr[current_row_index]; 79 | out_pi[passed]=in_pi[current_row_index]; 80 | out_ir[passed]=in_ir[current_row_index]; 81 | out_jc[col+1] = out_jc[col+1]+1; 82 | passed++; 83 | } 84 | } 85 | } 86 | } 87 | } 88 | else{ 89 | plhs[0] = mxCreateSparse(m,n,0,mxCOMPLEX); 90 | } 91 | } 92 | else{ /* for full matrices */ 93 | /* read input */ 94 | in_pr = mxGetPr(prhs[0]); 95 | in_pi = mxGetPr(prhs[0]); 96 | m = mxGetM(prhs[0]); 97 | n = mxGetN(prhs[0]); 98 | thres = mxGetScalar(prhs[1]); 99 | 100 | /* Count new nonzeros */ 101 | newnnz=0; 102 | for(i=0; ithres) {newnnz++;} 104 | } 105 | 106 | if (newnnz>0){ 107 | /* create output */ 108 | plhs[0] = mxCreateSparse(m,n,newnnz,mxCOMPLEX); 109 | if (plhs[0]==NULL) 110 | mexErrMsgTxt("Could not allocate enough memory!\n"); 111 | out_pr = mxGetPr(plhs[0]); 112 | out_pi = mxGetPi(plhs[0]); 113 | out_ir = mxGetIr(plhs[0]); 114 | out_jc = mxGetJc(plhs[0]); 115 | passed = 0; 116 | out_jc[0] = 0; 117 | 118 | for (col=0; col thres){ 123 | 124 | out_pr[passed]=in_pr[current_row_index+m*col]; 125 | out_ir[passed]=current_row_index; 126 | out_jc[col+1] = out_jc[col+1]+1; 127 | passed++; 128 | } 129 | } 130 | } 131 | } 132 | else{ 133 | plhs[0] = mxCreateSparse(m,n,0,mxCOMPLEX); 134 | } 135 | } 136 | } 137 | else { 138 | /* Check data type of input argument */ 139 | if (mxIsSparse(prhs[0])){ 140 | 141 | /* read input */ 142 | in_pr = mxGetPr(prhs[0]); 143 | in_ir = mxGetIr(prhs[0]); 144 | in_jc = mxGetJc(prhs[0]); 145 | nzmax = mxGetNzmax(prhs[0]); 146 | n = mxGetN(prhs[0]); 147 | m = mxGetM(prhs[0]); 148 | thres = mxGetScalar(prhs[1]); 149 | 150 | /* Count new nonzeros */ 151 | newnnz=0; 152 | for(i=0; ithres) {newnnz++;} 154 | } 155 | 156 | if (newnnz>0){ 157 | /* create output */ 158 | plhs[0] = mxCreateSparse(m,n,newnnz,mxREAL); 159 | if (plhs[0]==NULL) 160 | mexErrMsgTxt("Could not allocate enough memory!\n"); 161 | out_pr = mxGetPr(plhs[0]); 162 | out_ir = mxGetIr(plhs[0]); 163 | out_jc = mxGetJc(plhs[0]); 164 | passed = 0; 165 | out_jc[0] = 0; 166 | for (col=0; colthres){ 177 | out_pr[passed]=in_pr[current_row_index]; 178 | out_ir[passed]=in_ir[current_row_index]; 179 | out_jc[col+1] = out_jc[col+1]+1; 180 | passed++; 181 | } 182 | } 183 | } 184 | } 185 | } 186 | else{ 187 | plhs[0] = mxCreateSparse(m,n,0,mxREAL); 188 | } 189 | } 190 | else{ /* for full matrices */ 191 | /* read input */ 192 | in_pr = mxGetPr(prhs[0]); 193 | n = mxGetN(prhs[0]); 194 | m = mxGetM(prhs[0]); 195 | thres = mxGetScalar(prhs[1]); 196 | 197 | /* Count new nonzeros */ 198 | newnnz=0; 199 | for(i=0; ithres) {newnnz++;} 201 | } 202 | 203 | if (newnnz>0){ 204 | /* create output */ 205 | plhs[0] = mxCreateSparse(m,n,newnnz,mxREAL); 206 | if (plhs[0]==NULL) 207 | mexErrMsgTxt("Could not allocate enough memory!\n"); 208 | out_pr = mxGetPr(plhs[0]); 209 | out_ir = mxGetIr(plhs[0]); 210 | out_jc = mxGetJc(plhs[0]); 211 | passed = 0; 212 | out_jc[0] = 0; 213 | 214 | for (col=0; colthres){ 218 | out_pr[passed]=in_pr[current_row_index+m*col]; 219 | out_ir[passed]=current_row_index; 220 | out_jc[col+1] = out_jc[col+1]+1; 221 | passed++; 222 | } 223 | } 224 | } 225 | } 226 | else{ 227 | plhs[0] = mxCreateSparse(m,n,0,mxREAL); 228 | } 229 | } 230 | } 231 | } 232 | 233 | -------------------------------------------------------------------------------- /prerprocess/Ncut_9/sparsifyc.mexa64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/sparsifyc.mexa64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/sparsifyc.mexmaci64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/sparsifyc.mexmaci64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/sparsifyc.mexw64: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/Ncut_9/sparsifyc.mexw64 -------------------------------------------------------------------------------- /prerprocess/Ncut_9/spmtimesd.cpp: -------------------------------------------------------------------------------- 1 | /*================================================================ 2 | * spmtimesd.c 3 | * This routine computes a sparse matrix times a diagonal matrix 4 | * whose diagonal entries are stored in a full vector. 5 | * 6 | * Examples: 7 | * spmtimesd(m,d,[]) = diag(d) * m, 8 | * spmtimesd(m,[],d) = m * diag(d) 9 | * spmtimesd(m,d1,d2) = diag(d1) * m * diag(d2) 10 | * m could be complex, but d is assumed real 11 | * 12 | * Stella X. Yu's first MEX function, Nov 9, 2001. 13 | 14 | % test sequence: 15 | 16 | m = 1000; 17 | n = 2000; 18 | a=sparse(rand(m,n)); 19 | d1 = rand(m,1); 20 | d2 = rand(n,1); 21 | tic; b=spmtimesd(a,d1,d2); toc 22 | tic; bb = spdiags(d1,0,m,m) * a * spdiags(d2,0,n,n); toc 23 | e = (bb-b); 24 | max(abs(e(:))) 25 | 26 | *=================================================================*/ 27 | 28 | # include "mex.h" 29 | 30 | void mexFunction( 31 | int nargout, 32 | mxArray *out[], 33 | int nargin, 34 | const mxArray *in[] 35 | ) 36 | { 37 | /* declare variables */ 38 | int i, j, k, m, n, nzmax, xm, yn; 39 | mwIndex *pir, *pjc, *qir, *qjc; 40 | double *x, *y, *pr, *pi, *qr, *qi; 41 | 42 | /* check argument */ 43 | if (nargin != 3) { 44 | mexErrMsgTxt("Three input arguments required"); 45 | } 46 | if (nargout>1) { 47 | mexErrMsgTxt("Too many output arguments."); 48 | } 49 | if (!(mxIsSparse(in[0]))) { 50 | mexErrMsgTxt("Input argument #1 must be of type sparse"); 51 | } 52 | if ( mxIsSparse(in[1]) || mxIsSparse(in[2]) ) { 53 | mexErrMsgTxt("Input argument #2 & #3 must be of type full"); 54 | } 55 | 56 | /* computation starts */ 57 | m = mxGetM(in[0]); 58 | n = mxGetN(in[0]); 59 | pr = mxGetPr(in[0]); 60 | pi = mxGetPi(in[0]); 61 | pir = mxGetIr(in[0]); 62 | pjc = mxGetJc(in[0]); 63 | 64 | i = mxGetM(in[1]); 65 | j = mxGetN(in[1]); 66 | xm = ((i>j)? i: j); 67 | 68 | i = mxGetM(in[2]); 69 | j = mxGetN(in[2]); 70 | yn = ((i>j)? i: j); 71 | 72 | if ( xm>0 && xm != m) { 73 | mexErrMsgTxt("Row multiplication dimension mismatch."); 74 | } 75 | if ( yn>0 && yn != n) { 76 | mexErrMsgTxt("Column multiplication dimension mismatch."); 77 | } 78 | 79 | 80 | nzmax = mxGetNzmax(in[0]); 81 | mxComplexity cmplx = (pi==NULL ? mxREAL : mxCOMPLEX); 82 | out[0] = mxCreateSparse(m,n,nzmax,cmplx); 83 | if (out[0]==NULL) { 84 | mexErrMsgTxt("Not enough space for the output matrix."); 85 | } 86 | 87 | qr = mxGetPr(out[0]); 88 | qi = mxGetPi(out[0]); 89 | qir = mxGetIr(out[0]); 90 | qjc = mxGetJc(out[0]); 91 | 92 | /* left multiplication */ 93 | x = mxGetPr(in[1]); 94 | if (yn==0) { 95 | for (j=0; j 10: 49 | labels[count] = int(L[16]) 50 | ious[count] = float(L[17]) 51 | if ious[count] < 0.5: 52 | labels[count] = 0 53 | else: 54 | offsets[count,0] = float(L[8]) 55 | offsets[count,1] = float(L[9]) 56 | offsets[count,2] = float(L[10]) 57 | offsets[count,3] = float(L[11]) 58 | offsets[count,4] = float(L[12]) 59 | offsets[count,5] = float(L[13]) 60 | offsets[count,6] = float(L[14]) 61 | offsets[count,7] = float(L[15]) 62 | 63 | tmp1 = os.path.join(scene_dir,'data','merge_leaf_pts','objects_'+str(count)+'_pointcnn_feature.txt') 64 | tmp2 = os.path.join(scene_dir,'data','internal_pts','objects_'+str(count)+'_pointcnn_feature.txt') 65 | if os.path.exists(tmp1) and valids[count] == True: 66 | feature_dir.append(tmp1) 67 | valids[count] = True 68 | counts[0] += 1 69 | elif os.path.exists(tmp2) and valids[count] == True: 70 | feature_dir.append(tmp2) 71 | valids[count] = True 72 | counts[0] += 1 73 | else: 74 | feature_dir.append('empty') 75 | labels[count] = 0 76 | counts[1] += 1 77 | count = count + 1 78 | 79 | hier = open(hier_name) 80 | mapFather = {} 81 | mapChild1 = {} 82 | mapChild2 = {} 83 | isLeaf = {} 84 | 85 | while 1: 86 | line = hier.readline() 87 | if not line: 88 | break 89 | L = line.split() 90 | father = int(L[0]) 91 | child = int(L[1]) 92 | mapFather[child] = father 93 | isLeaf[father] = 'False' 94 | if mapChild1.get(father, 'False') == 'False': 95 | mapChild1[father] = child 96 | else: 97 | mapChild2[father] = child 98 | if L[2] == "null": 99 | isLeaf[child] = 'False' 100 | else: 101 | isLeaf[child] = 'True' 102 | return(feature_dir, obbs, offsets, labels, labels_pointcnn, mapFather, mapChild1, mapChild2, isLeaf) 103 | 104 | 105 | config = util.get_args() 106 | g_path = config.g_path 107 | scene_list_path = config.scene_list_path 108 | pickle_path = config.pickle_path 109 | if not os.path.exists(pickle_path): 110 | os.mkdir(pickle_path) 111 | 112 | f = open(scene_list_path) 113 | lines = f.readlines() 114 | scene_list = [] 115 | for line in lines: 116 | scene_list.append(line.split()[0]) 117 | 118 | count = 0 119 | scene_names = os.listdir(g_path) 120 | scene_names.sort() 121 | for scene_name in scene_names: 122 | if scene_name.split('.')[0] not in scene_list: 123 | continue 124 | scene_dir = os.path.join(g_path,scene_name) 125 | hier_name = os.path.join(scene_dir,'data','files','hier_ncut.txt') 126 | node_info_name = os.path.join(scene_dir,'data','files','node_info_on_points.txt') 127 | all_box_dir = os.path.join(scene_dir,'data','files','pred_box_all_node.txt') 128 | 129 | if not os.path.exists(hier_name): 130 | print('hierarchy not exist',hier_name) 131 | continue 132 | (feature_dir,boxes,boxes_reg,category,category_pred,mapFather, mapChild1, mapChild2, isLeaf) = genGrassData(scene_dir,hier_name,node_info_name,all_box_dir) 133 | mdict={'scene_dir': scene_dir, 'feature_dir': feature_dir, 'boxes': boxes,'boxes_reg': boxes_reg, 'category_pred':category_pred, 'category':category, 'mapFather':mapFather, 'mapChild1':mapChild1, 'mapChild2':mapChild2, 'isLeaf':isLeaf} 134 | with open(os.path.join(pickle_path, 'scene' + '_' + str(count) + '.pkl'), 'wb+') as f: 135 | pickle.dump(mdict,f) 136 | count = count + 1 137 | -------------------------------------------------------------------------------- /prerprocess/gen_txt.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import util 4 | 5 | config = util.get_args() 6 | g_path = config.g_path 7 | 8 | data_dir = os.path.join(g_path,'data') 9 | merge_leaf_dir = os.path.join(data_dir,'merge_leaf_pts') 10 | internal_dir = os.path.join(data_dir,'internal_pts') 11 | files_dir = os.path.join(data_dir,'files') 12 | node_info_dir = os.path.join(files_dir,'node_info_on_points.txt') 13 | pred_box_txt_dir = os.path.join(files_dir,'pred_box_all_node.txt') 14 | 15 | id_pts_dict = {} 16 | id_pts_feature_dict = {} 17 | pts_names = os.listdir(merge_leaf_dir) 18 | pts_names.sort() 19 | for pts_name in pts_names: 20 | if len(pts_name.split('.'))>1 and pts_name.split('.')[1]=='pts': 21 | id = int(pts_name.split('.')[0].split('_')[1]) 22 | pts_file = os.path.join(merge_leaf_dir,pts_name) 23 | pts_feature_file = os.path.join(merge_leaf_dir,pts_name.split('.')[0]+'_pointcnn_feature.txt') 24 | id_pts_dict[id] = pts_file 25 | id_pts_feature_dict[id] = pts_feature_file 26 | internal_pts_names = os.listdir(internal_dir) 27 | internal_pts_names.sort() 28 | for internal_pts_name in internal_pts_names: 29 | if len(internal_pts_name.split('.'))>1 and internal_pts_name.split('.')[1]=='pts': 30 | id = int(internal_pts_name.split('.')[0].split('_')[1]) 31 | pts_file= os.path.join(internal_dir,internal_pts_name) 32 | pts_feature_file = os.path.join(internal_dir,internal_pts_name.split('.')[0]+'_pointcnn_feature.txt') 33 | id_pts_dict[id] = pts_file 34 | id_pts_feature_dict[id] = pts_feature_file 35 | 36 | output = open('Matterport_all_node_pointcnn_input.txt','a') 37 | f = open(node_info_dir,'r') 38 | lines = f.readlines() 39 | for i in range(len(lines)): 40 | L = lines[i].split() 41 | for j in range(len(L)): 42 | L[j] = float(L[j]) 43 | if len(L) > 8: 44 | output.write('%s %f %f %f %f %f %f %f %f %d %s %s\n'\ 45 | %(id_pts_dict[i],L[0],L[1],L[2],L[3],\ 46 | L[4],L[5],L[6],L[7],\ 47 | L[16],pred_box_txt_dir,id_pts_feature_dict[i])) 48 | else: 49 | neg = float(99) 50 | output.write('%s %f %f %f %f %f %f %f %f %d %s %s\n'\ 51 | %(id_pts_dict[i],L[0],L[1],L[2],L[3],\ 52 | L[4],L[5],L[6],L[7],\ 53 | neg,pred_box_txt_dir,id_pts_feature_dict[i])) 54 | 55 | for i in range(64): 56 | L = lines[0].split() 57 | for j in range(len(L)): 58 | L[j] = float(L[j]) 59 | neg = float(99) 60 | pred_box_txt_dir = 'tmp.txt' 61 | output.write('%s %f %f %f %f %f %f %f %f %d %s %s\n'\ 62 | %(id_pts_dict[i],L[0],L[1],L[2],L[3],\ 63 | L[4],L[5],L[6],L[7],\ 64 | neg,pred_box_txt_dir,pred_box_txt_dir)) 65 | output.close() 66 | -------------------------------------------------------------------------------- /prerprocess/min_bounding_rect.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # Find the minimum-area bounding box of a set of 2D points 4 | # 5 | # The input is a 2D convex hull, in an Nx2 numpy array of x-y co-ordinates. 6 | # The first and last points points must be the same, making a closed polygon. 7 | # This program finds the rotation angles of each edge of the convex polygon, 8 | # then tests the area of a bounding box aligned with the unique angles in 9 | # 90 degrees of the 1st Quadrant. 10 | # Returns the 11 | # 12 | # Tested with Python 2.6.5 on Ubuntu 10.04.4 13 | # Results verified using Matlab 14 | 15 | # Copyright (c) 2013, David Butterworth, University of Queensland 16 | # All rights reserved. 17 | # 18 | # Redistribution and use in source and binary forms, with or without 19 | # modification, are permitted provided that the following conditions are met: 20 | # 21 | # * Redistributions of source code must retain the above copyright 22 | # notice, this list of conditions and the following disclaimer. 23 | # * Redistributions in binary form must reproduce the above copyright 24 | # notice, this list of conditions and the following disclaimer in the 25 | # documentation and/or other materials provided with the distribution. 26 | # * Neither the name of the Willow Garage, Inc. nor the names of its 27 | # contributors may be used to endorse or promote products derived from 28 | # this software without specific prior written permission. 29 | # 30 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 31 | # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 32 | # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 33 | # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 34 | # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 35 | # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 36 | # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 37 | # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 38 | # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 39 | # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 40 | # POSSIBILITY OF SUCH DAMAGE. 41 | 42 | from numpy import * 43 | import sys # maxint 44 | 45 | def minBoundingRect(hull_points_2d): 46 | #print "Input convex hull points: " 47 | #print hull_points_2d 48 | 49 | # Compute edges (x2-x1,y2-y1) 50 | edges = zeros( (len(hull_points_2d)-1,2) ) # empty 2 column array 51 | for i in range( len(edges) ): 52 | edge_x = hull_points_2d[i+1,0] - hull_points_2d[i,0] 53 | edge_y = hull_points_2d[i+1,1] - hull_points_2d[i,1] 54 | edges[i] = [edge_x,edge_y] 55 | #print "Edges: \n", edges 56 | 57 | # Calculate edge angles atan2(y/x) 58 | edge_angles = zeros( (len(edges)) ) # empty 1 column array 59 | for i in range( len(edge_angles) ): 60 | edge_angles[i] = math.atan2( edges[i,1], edges[i,0] ) 61 | #print "Edge angles: \n", edge_angles 62 | 63 | # Check for angles in 1st quadrant 64 | for i in range( len(edge_angles) ): 65 | edge_angles[i] = abs( edge_angles[i] % (math.pi/2) ) # want strictly positive answers 66 | #print "Edge angles in 1st Quadrant: \n", edge_angles 67 | 68 | # Remove duplicate angles 69 | edge_angles = unique(edge_angles) 70 | #print "Unique edge angles: \n", edge_angles 71 | 72 | # Test each angle to find bounding box with smallest area 73 | min_bbox = (0, sys.maxint, 0, 0, 0, 0, 0, 0) # rot_angle, area, width, height, min_x, max_x, min_y, max_y 74 | #print "Testing", len(edge_angles), "possible rotations for bounding box... \n" 75 | for i in range( len(edge_angles) ): 76 | 77 | # Create rotation matrix to shift points to baseline 78 | # R = [ cos(theta) , cos(theta-PI/2) 79 | # cos(theta+PI/2) , cos(theta) ] 80 | R = array([ [ math.cos(edge_angles[i]), math.cos(edge_angles[i]-(math.pi/2)) ], [ math.cos(edge_angles[i]+(math.pi/2)), math.cos(edge_angles[i]) ] ]) 81 | #print "Rotation matrix for ", edge_angles[i], " is \n", R 82 | 83 | # Apply this rotation to convex hull points 84 | rot_points = dot(R, transpose(hull_points_2d) ) # 2x2 * 2xn 85 | #print "Rotated hull points are \n", rot_points 86 | 87 | # Find min/max x,y points 88 | min_x = nanmin(rot_points[0], axis=0) 89 | max_x = nanmax(rot_points[0], axis=0) 90 | min_y = nanmin(rot_points[1], axis=0) 91 | max_y = nanmax(rot_points[1], axis=0) 92 | #print "Min x:", min_x, " Max x: ", max_x, " Min y:", min_y, " Max y: ", max_y 93 | 94 | # Calculate height/width/area of this bounding rectangle 95 | width = max_x - min_x 96 | height = max_y - min_y 97 | area = width*height 98 | #print "Potential bounding box ", i, ": width: ", width, " height: ", height, " area: ", area 99 | 100 | # Store the smallest rect found first (a simple convex hull might have 2 answers with same area) 101 | if (area < min_bbox[1]): 102 | min_bbox = ( edge_angles[i], area, width, height, min_x, max_x, min_y, max_y ) 103 | # Bypass, return the last found rect 104 | #min_bbox = ( edge_angles[i], area, width, height, min_x, max_x, min_y, max_y ) 105 | 106 | # Re-create rotation matrix for smallest rect 107 | angle = min_bbox[0] 108 | R = array([ [ math.cos(angle), math.cos(angle-(math.pi/2)) ], [ math.cos(angle+(math.pi/2)), math.cos(angle) ] ]) 109 | #print "Projection matrix: \n", R 110 | 111 | # Project convex hull points onto rotated frame 112 | proj_points = dot(R, transpose(hull_points_2d) ) # 2x2 * 2xn 113 | #print "Project hull points are \n", proj_points 114 | 115 | # min/max x,y points are against baseline 116 | min_x = min_bbox[4] 117 | max_x = min_bbox[5] 118 | min_y = min_bbox[6] 119 | max_y = min_bbox[7] 120 | #print "Min x:", min_x, " Max x: ", max_x, " Min y:", min_y, " Max y: ", max_y 121 | 122 | # Calculate center point and project onto rotated frame 123 | center_x = (min_x + max_x)/2 124 | center_y = (min_y + max_y)/2 125 | center_point = dot( [ center_x, center_y ], R ) 126 | #print "Bounding box center point: \n", center_point 127 | 128 | # Calculate corner points and project onto rotated frame 129 | corner_points = zeros( (4,2) ) # empty 2 column array 130 | corner_points[0] = dot( [ max_x, min_y ], R ) 131 | corner_points[1] = dot( [ min_x, min_y ], R ) 132 | corner_points[2] = dot( [ min_x, max_y ], R ) 133 | corner_points[3] = dot( [ max_x, max_y ], R ) 134 | #print "Bounding box corner points: \n", corner_points 135 | 136 | #print "Angle of rotation: ", angle, "rad ", angle * (180/math.pi), "deg" 137 | 138 | return (angle, min_bbox[1], min_bbox[2], min_bbox[3], center_point, corner_points) # rot_angle, area, width, height, center_point, corner_points 139 | 140 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/data_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import time 6 | import os 7 | import h5py 8 | import plyfile 9 | import random 10 | import numpy as np 11 | from matplotlib import cm 12 | 13 | 14 | def save_ply(points, filename, colors=None, normals=None): 15 | vertex = np.array([tuple(p) for p in points], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 16 | n = len(vertex) 17 | desc = vertex.dtype.descr 18 | 19 | if normals is not None: 20 | vertex_normal = np.array([tuple(n) for n in normals], dtype=[('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4')]) 21 | assert len(vertex_normal) == n 22 | desc = desc + vertex_normal.dtype.descr 23 | 24 | if colors is not None: 25 | vertex_color = np.array([tuple(c * 255) for c in colors], 26 | dtype=[('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]) 27 | assert len(vertex_color) == n 28 | desc = desc + vertex_color.dtype.descr 29 | 30 | vertex_all = np.empty(n, dtype=desc) 31 | 32 | for prop in vertex.dtype.names: 33 | vertex_all[prop] = vertex[prop] 34 | 35 | if normals is not None: 36 | for prop in vertex_normal.dtype.names: 37 | vertex_all[prop] = vertex_normal[prop] 38 | 39 | if colors is not None: 40 | for prop in vertex_color.dtype.names: 41 | vertex_all[prop] = vertex_color[prop] 42 | 43 | ply = plyfile.PlyData([plyfile.PlyElement.describe(vertex_all, 'vertex')], text=False) 44 | if not os.path.exists(os.path.dirname(filename)): 45 | os.makedirs(os.path.dirname(filename)) 46 | ply.write(filename) 47 | 48 | 49 | def save_ply_property(points, property, property_max, filename, cmap_name='tab20'): 50 | point_num = points.shape[0] 51 | colors = np.full(points.shape, 0.5) 52 | cmap = cm.get_cmap(cmap_name) 53 | for point_idx in range(point_num): 54 | if property[point_idx] == 0: 55 | colors[point_idx] = np.array([0, 0, 0]) 56 | else: 57 | colors[point_idx] = cmap(property[point_idx] / property_max)[:3] 58 | save_ply(points, filename, colors) 59 | 60 | 61 | def save_ply_batch(points_batch, file_path, points_num=None): 62 | batch_size = points_batch.shape[0] 63 | if type(file_path) != list: 64 | basename = os.path.splitext(file_path)[0] 65 | ext = '.ply' 66 | for batch_idx in range(batch_size): 67 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 68 | if type(file_path) == list: 69 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx]) 70 | else: 71 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext)) 72 | 73 | 74 | def save_ply_color_batch(points_batch, colors_batch, file_path, points_num=None): 75 | batch_size = points_batch.shape[0] 76 | if type(file_path) != list: 77 | basename = os.path.splitext(file_path)[0] 78 | ext = '.ply' 79 | for batch_idx in range(batch_size): 80 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 81 | if type(file_path) == list: 82 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx], colors_batch[batch_idx][:point_num]) 83 | else: 84 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext), 85 | colors_batch[batch_idx][:point_num]) 86 | 87 | 88 | def save_ply_property_batch(points_batch, property_batch, file_path, points_num=None, property_max=None, 89 | cmap_name='tab20'): 90 | batch_size = points_batch.shape[0] 91 | if type(file_path) != list: 92 | basename = os.path.splitext(file_path)[0] 93 | ext = '.ply' 94 | property_max = np.max(property_batch) if property_max is None else property_max 95 | for batch_idx in range(batch_size): 96 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 97 | if type(file_path) == list: 98 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 99 | property_max, file_path[batch_idx], cmap_name) 100 | else: 101 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 102 | property_max, '%s_%04d%s' % (basename, batch_idx, ext), cmap_name) 103 | 104 | 105 | def save_ply_point_with_normal(data_sample, folder): 106 | for idx, sample in enumerate(data_sample): 107 | filename_pts = os.path.join(folder, '{:08d}.ply'.format(idx)) 108 | save_ply(sample[..., :3], filename_pts, normals=sample[..., 3:]) 109 | 110 | 111 | def grouped_shuffle(inputs): 112 | for idx in range(len(inputs) - 1): 113 | assert (len(inputs[idx]) == len(inputs[idx + 1])) 114 | 115 | shuffle_indices = np.arange(len(inputs[0])) 116 | np.random.shuffle(shuffle_indices) 117 | outputs = [] 118 | for idx in range(len(inputs)): 119 | for i in range(len(inputs[0])): 120 | outputs.append(inputs[idx][shuffle_indices[i], ...]) 121 | outputs.append(inputs[idx][shuffle_indices[i], ...]) 122 | return outputs 123 | 124 | 125 | def load_cls(filelist): 126 | points = [] 127 | labels = [] 128 | 129 | folder = os.path.dirname(filelist) 130 | for line in open(filelist): 131 | filename = os.path.basename(line.rstrip()) 132 | data = h5py.File(os.path.join(folder, filename)) 133 | if 'normal' in data: 134 | points.append(np.concatenate([data['data'][...], data['normal'][...]], axis=-1).astype(np.float32)) 135 | else: 136 | points.append(data['data'][...].astype(np.float32)) 137 | labels.append(np.squeeze(data['label'][:]).astype(np.int64)) 138 | return (np.concatenate(points, axis=0), 139 | np.concatenate(labels, axis=0)) 140 | 141 | 142 | def load_cls_train_val(filelist, filelist_val): 143 | data_train, label_train = grouped_shuffle(load_cls(filelist)) 144 | data_val, label_val = load_cls(filelist_val) 145 | return data_train, label_train, data_val, label_val 146 | 147 | 148 | def is_h5_list(filelist): 149 | return all([line.strip()[-3:] == '.h5' for line in open(filelist)]) 150 | 151 | 152 | def load_seg_list(filelist): 153 | folder = os.path.dirname(filelist) 154 | return [os.path.join(folder, line.strip()) for line in open(filelist)] 155 | 156 | 157 | def load_seg(filelist): 158 | points = [] 159 | labels = [] 160 | point_nums = [] 161 | labels_seg = [] 162 | indices_split_to_full = [] 163 | 164 | folder = os.path.dirname(filelist) 165 | for line in open(filelist): 166 | data = h5py.File(os.path.join(folder, line.strip())) 167 | points.append(data['data'][...].astype(np.float32)) 168 | labels.append(data['label'][...].astype(np.int64)) 169 | point_nums.append(data['data_num'][...].astype(np.int32)) 170 | labels_seg.append(data['label_seg'][...].astype(np.int64)) 171 | if 'indices_split_to_full' in data: 172 | indices_split_to_full.append(data['indices_split_to_full'][...].astype(np.int64)) 173 | 174 | return (np.concatenate(points, axis=0), 175 | np.concatenate(labels, axis=0), 176 | np.concatenate(point_nums, axis=0), 177 | np.concatenate(labels_seg, axis=0), 178 | np.concatenate(indices_split_to_full, axis=0) if indices_split_to_full else None) 179 | 180 | 181 | def balance_classes(labels): 182 | _, inverse, counts = np.unique(labels, return_inverse=True, return_counts=True) 183 | counts_max = np.amax(counts) 184 | repeat_num_avg_unique = counts_max / counts 185 | repeat_num_avg = repeat_num_avg_unique[inverse] 186 | repeat_num_floor = np.floor(repeat_num_avg) 187 | repeat_num_probs = repeat_num_avg - repeat_num_floor 188 | repeat_num = repeat_num_floor + (np.random.rand(repeat_num_probs.shape[0]) < repeat_num_probs) 189 | 190 | return repeat_num.astype(np.int64) 191 | 192 | def load_cls_pts(filelist,k): 193 | 194 | labels = [] 195 | points = [] 196 | for i in range(k): 197 | folder = os.path.dirname(filelist) 198 | for line in open(filelist): 199 | filename , label = line.split(" ",1) 200 | _,filename = filename.split("./") 201 | file_path = os.path.join(folder, filename) 202 | data = [] 203 | count = 0 204 | with open(file_path) as fpts: 205 | while 1: 206 | line = fpts.readline() 207 | if not line: 208 | break; 209 | L = line.split(' ') 210 | L = [float(i) for i in L] 211 | data.append(np.array(L)) 212 | count = count + 1 213 | if(count>=2048): 214 | index = random.sample(range(0,count),2048) 215 | pointcloud = np.zeros((2048, 6),dtype=np.float32) 216 | for i in range(2048): 217 | pointcloud[i] = data[index[i]] 218 | points.append(pointcloud) 219 | labels.append(int(label)) 220 | 221 | labels = np.array(labels) 222 | points = np.array(points) 223 | return points , labels 224 | 225 | 226 | def load_ptsfile(filelist,k=1): 227 | data_train, label_train = grouped_shuffle(load_cls_pts(filelist,k)) 228 | return data_train, label_train 229 | 230 | def load_file(filelist): 231 | 232 | labels = [] 233 | weights = [] 234 | points_path = [] 235 | sizes = [] 236 | len = 0 237 | for line in open(filelist): 238 | size = np.zeros((1,3)) 239 | filename , label, weight, size[0,0], size[0,1], size[0,2] = line.split(" ") 240 | points_path.append(filename) 241 | weights.append(float(weight)) 242 | labels.append(int(label)) 243 | sizes.append(size) 244 | #print(size[0,0], size[0,1], size[0,2]) 245 | len = len +1 246 | labels = np.array(labels) 247 | return points_path, labels, weights, sizes, len 248 | 249 | def read_path(filelist): 250 | 251 | read_paths = [] 252 | write_paths = [] 253 | len = 0 254 | for line in open(filelist): 255 | read_path, write_path = line.split(" ") 256 | write_path = write_path.rstrip('\n') 257 | read_paths.append(read_path) 258 | write_paths.append(write_path) 259 | len = len +1 260 | return read_paths, write_paths, len 261 | 262 | def read_test_path(filelist): 263 | 264 | read_paths = [] 265 | sizes = [] 266 | obbs = [] 267 | gts = [] 268 | write_paths = [] 269 | feature_paths = [] 270 | len = 0 271 | for line in open(filelist): 272 | size = np.zeros((1,3)) 273 | obb = np.zeros(8) 274 | read_path, obb[0], obb[1], obb[2], obb[3], obb[4], obb[5], obb[6], obb[7], gt, write_path, feature_path = line.split(" ") 275 | size[0,0], size[0,1], size[0,2] = obb[0], obb[1], obb[2] 276 | feature_path = feature_path.rstrip('\n') 277 | read_paths.append(read_path) 278 | sizes.append(size) 279 | obbs.append(obb) 280 | gts.append(int(gt)) 281 | write_paths.append(write_path) 282 | feature_paths.append(feature_path) 283 | len = len +1 284 | return read_paths, sizes, obbs, gts, write_paths, feature_paths, len -------------------------------------------------------------------------------- /prerprocess/pointcnn/data_utils.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/prerprocess/pointcnn/data_utils.pyc -------------------------------------------------------------------------------- /prerprocess/pointcnn/example_x3_l4.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | import os 4 | import sys 5 | import math 6 | 7 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 8 | import data_utils 9 | 10 | load_fn = data_utils.load_ptsfile 11 | balance_fn = None 12 | map_fn = None 13 | keep_remainder = True 14 | save_ply_fn = None 15 | num_class = 101 16 | 17 | sample_num = 2048 18 | 19 | batch_size = 32 20 | 21 | num_epochs = 1500 22 | 23 | step_val = 500 24 | 25 | learning_rate_base = 0.01 26 | decay_steps = 50000 27 | decay_rate = 0.5 28 | learning_rate_min = 1e-6 29 | 30 | weight_decay = 1e-5 31 | 32 | jitter = 0.0 33 | jitter_val = 0.0 34 | 35 | #rotation_range = [0, math.pi, 0, 'u'] 36 | rotation_range = [0, 0, math.pi, 'u'] 37 | rotation_range_val = [0, 0, 0, 'u'] 38 | rotation_order = 'rxyz' 39 | 40 | #scaling_range = [0.1, 0.1, 0.1, 'g'] 41 | scaling_range = [0, 0, 0, 'g'] 42 | scaling_range_val = [0, 0, 0, 'u'] 43 | 44 | sample_num_variance = 1 // 8 45 | sample_num_clip = 1 // 4 46 | 47 | x = 3 48 | 49 | xconv_param_name = ('K', 'D', 'P', 'C', 'links') 50 | xconv_params = [dict(zip(xconv_param_name, xconv_param)) for xconv_param in 51 | [(8, 1, -1, 16 * x, []), 52 | (12, 2, 384, 32 * x, []), 53 | (16, 2, 128, 64 * x, []), 54 | (16, 3, 128, 128 * x, [])]] 55 | 56 | with_global = True 57 | 58 | fc_param_name = ('C', 'dropout_rate') 59 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 60 | [(128 * x, 0.0), 61 | (64 * 4, 0.8)]] 62 | # 全连接参数 63 | 64 | sampling = 'random' 65 | 66 | optimizer = 'adam' 67 | epsilon = 1e-2 68 | 69 | data_dim = 6 70 | use_extra_features = False 71 | with_X_transformation = True 72 | sorting_method = None 73 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/folder for trans/data_utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import time 6 | import os 7 | import h5py 8 | import plyfile 9 | import random 10 | import numpy as np 11 | from matplotlib import cm 12 | 13 | 14 | def save_ply(points, filename, colors=None, normals=None): 15 | vertex = np.array([tuple(p) for p in points], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')]) 16 | n = len(vertex) 17 | desc = vertex.dtype.descr 18 | 19 | if normals is not None: 20 | vertex_normal = np.array([tuple(n) for n in normals], dtype=[('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4')]) 21 | assert len(vertex_normal) == n 22 | desc = desc + vertex_normal.dtype.descr 23 | 24 | if colors is not None: 25 | vertex_color = np.array([tuple(c * 255) for c in colors], 26 | dtype=[('red', 'u1'), ('green', 'u1'), ('blue', 'u1')]) 27 | assert len(vertex_color) == n 28 | desc = desc + vertex_color.dtype.descr 29 | 30 | vertex_all = np.empty(n, dtype=desc) 31 | 32 | for prop in vertex.dtype.names: 33 | vertex_all[prop] = vertex[prop] 34 | 35 | if normals is not None: 36 | for prop in vertex_normal.dtype.names: 37 | vertex_all[prop] = vertex_normal[prop] 38 | 39 | if colors is not None: 40 | for prop in vertex_color.dtype.names: 41 | vertex_all[prop] = vertex_color[prop] 42 | 43 | ply = plyfile.PlyData([plyfile.PlyElement.describe(vertex_all, 'vertex')], text=False) 44 | if not os.path.exists(os.path.dirname(filename)): 45 | os.makedirs(os.path.dirname(filename)) 46 | ply.write(filename) 47 | 48 | 49 | def save_ply_property(points, property, property_max, filename, cmap_name='tab20'): 50 | point_num = points.shape[0] 51 | colors = np.full(points.shape, 0.5) 52 | cmap = cm.get_cmap(cmap_name) 53 | for point_idx in range(point_num): 54 | if property[point_idx] == 0: 55 | colors[point_idx] = np.array([0, 0, 0]) 56 | else: 57 | colors[point_idx] = cmap(property[point_idx] / property_max)[:3] 58 | save_ply(points, filename, colors) 59 | 60 | 61 | def save_ply_batch(points_batch, file_path, points_num=None): 62 | batch_size = points_batch.shape[0] 63 | if type(file_path) != list: 64 | basename = os.path.splitext(file_path)[0] 65 | ext = '.ply' 66 | for batch_idx in range(batch_size): 67 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 68 | if type(file_path) == list: 69 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx]) 70 | else: 71 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext)) 72 | 73 | 74 | def save_ply_color_batch(points_batch, colors_batch, file_path, points_num=None): 75 | batch_size = points_batch.shape[0] 76 | if type(file_path) != list: 77 | basename = os.path.splitext(file_path)[0] 78 | ext = '.ply' 79 | for batch_idx in range(batch_size): 80 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 81 | if type(file_path) == list: 82 | save_ply(points_batch[batch_idx][:point_num], file_path[batch_idx], colors_batch[batch_idx][:point_num]) 83 | else: 84 | save_ply(points_batch[batch_idx][:point_num], '%s_%04d%s' % (basename, batch_idx, ext), 85 | colors_batch[batch_idx][:point_num]) 86 | 87 | 88 | def save_ply_property_batch(points_batch, property_batch, file_path, points_num=None, property_max=None, 89 | cmap_name='tab20'): 90 | batch_size = points_batch.shape[0] 91 | if type(file_path) != list: 92 | basename = os.path.splitext(file_path)[0] 93 | ext = '.ply' 94 | property_max = np.max(property_batch) if property_max is None else property_max 95 | for batch_idx in range(batch_size): 96 | point_num = points_batch.shape[1] if points_num is None else points_num[batch_idx] 97 | if type(file_path) == list: 98 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 99 | property_max, file_path[batch_idx], cmap_name) 100 | else: 101 | save_ply_property(points_batch[batch_idx][:point_num], property_batch[batch_idx][:point_num], 102 | property_max, '%s_%04d%s' % (basename, batch_idx, ext), cmap_name) 103 | 104 | 105 | def save_ply_point_with_normal(data_sample, folder): 106 | for idx, sample in enumerate(data_sample): 107 | filename_pts = os.path.join(folder, '{:08d}.ply'.format(idx)) 108 | save_ply(sample[..., :3], filename_pts, normals=sample[..., 3:]) 109 | 110 | 111 | def grouped_shuffle(inputs): 112 | for idx in range(len(inputs) - 1): 113 | assert (len(inputs[idx]) == len(inputs[idx + 1])) 114 | 115 | shuffle_indices = np.arange(len(inputs[0])) 116 | np.random.shuffle(shuffle_indices) 117 | outputs = [] 118 | for idx in range(len(inputs)): 119 | for i in range(len(inputs[0])): 120 | outputs.append(inputs[idx][shuffle_indices[i], ...]) 121 | outputs.append(inputs[idx][shuffle_indices[i], ...]) 122 | return outputs 123 | 124 | 125 | def load_cls(filelist): 126 | points = [] 127 | labels = [] 128 | 129 | folder = os.path.dirname(filelist) 130 | for line in open(filelist): 131 | filename = os.path.basename(line.rstrip()) 132 | data = h5py.File(os.path.join(folder, filename)) 133 | if 'normal' in data: 134 | points.append(np.concatenate([data['data'][...], data['normal'][...]], axis=-1).astype(np.float32)) 135 | else: 136 | points.append(data['data'][...].astype(np.float32)) 137 | labels.append(np.squeeze(data['label'][:]).astype(np.int64)) 138 | return (np.concatenate(points, axis=0), 139 | np.concatenate(labels, axis=0)) 140 | 141 | 142 | def load_cls_train_val(filelist, filelist_val): 143 | data_train, label_train = grouped_shuffle(load_cls(filelist)) 144 | data_val, label_val = load_cls(filelist_val) 145 | return data_train, label_train, data_val, label_val 146 | 147 | 148 | def is_h5_list(filelist): 149 | return all([line.strip()[-3:] == '.h5' for line in open(filelist)]) 150 | 151 | 152 | def load_seg_list(filelist): 153 | folder = os.path.dirname(filelist) 154 | return [os.path.join(folder, line.strip()) for line in open(filelist)] 155 | 156 | 157 | def load_seg(filelist): 158 | points = [] 159 | labels = [] 160 | point_nums = [] 161 | labels_seg = [] 162 | indices_split_to_full = [] 163 | 164 | folder = os.path.dirname(filelist) 165 | for line in open(filelist): 166 | data = h5py.File(os.path.join(folder, line.strip())) 167 | points.append(data['data'][...].astype(np.float32)) 168 | labels.append(data['label'][...].astype(np.int64)) 169 | point_nums.append(data['data_num'][...].astype(np.int32)) 170 | labels_seg.append(data['label_seg'][...].astype(np.int64)) 171 | if 'indices_split_to_full' in data: 172 | indices_split_to_full.append(data['indices_split_to_full'][...].astype(np.int64)) 173 | 174 | return (np.concatenate(points, axis=0), 175 | np.concatenate(labels, axis=0), 176 | np.concatenate(point_nums, axis=0), 177 | np.concatenate(labels_seg, axis=0), 178 | np.concatenate(indices_split_to_full, axis=0) if indices_split_to_full else None) 179 | 180 | 181 | def balance_classes(labels): 182 | _, inverse, counts = np.unique(labels, return_inverse=True, return_counts=True) 183 | counts_max = np.amax(counts) 184 | repeat_num_avg_unique = counts_max / counts 185 | repeat_num_avg = repeat_num_avg_unique[inverse] 186 | repeat_num_floor = np.floor(repeat_num_avg) 187 | repeat_num_probs = repeat_num_avg - repeat_num_floor 188 | repeat_num = repeat_num_floor + (np.random.rand(repeat_num_probs.shape[0]) < repeat_num_probs) 189 | 190 | return repeat_num.astype(np.int64) 191 | 192 | def load_cls_pts(filelist,k): 193 | 194 | labels = [] 195 | points = [] 196 | for i in range(k): 197 | folder = os.path.dirname(filelist) 198 | for line in open(filelist): 199 | filename , label = line.split(" ",1) 200 | _,filename = filename.split("./") 201 | file_path = os.path.join(folder, filename) 202 | data = [] 203 | count = 0 204 | with open(file_path) as fpts: 205 | while 1: 206 | line = fpts.readline() 207 | if not line: 208 | break; 209 | L = line.split(' ') 210 | L = [float(i) for i in L] 211 | data.append(np.array(L)) 212 | count = count + 1 213 | if(count>=2048): 214 | index = random.sample(range(0,count),2048) 215 | pointcloud = np.zeros((2048, 6),dtype=np.float32) 216 | for i in range(2048): 217 | pointcloud[i] = data[index[i]] 218 | points.append(pointcloud) 219 | labels.append(int(label)) 220 | 221 | labels = np.array(labels) 222 | points = np.array(points) 223 | return points , labels 224 | 225 | 226 | def load_ptsfile(filelist,k=1): 227 | data_train, label_train = grouped_shuffle(load_cls_pts(filelist,k)) 228 | return data_train, label_train 229 | 230 | def load_file(filelist): 231 | 232 | labels = [] 233 | weights = [] 234 | points_path = [] 235 | len = 0 236 | for line in open(filelist): 237 | filename , label, weight = line.split(" ") 238 | points_path.append(filename) 239 | weights.append(float(weight)) 240 | labels.append(int(label)) 241 | len = len +1 242 | labels = np.array(labels) 243 | return points_path, labels, weights, len -------------------------------------------------------------------------------- /prerprocess/pointcnn/folder for trans/example_x3_l4.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | import os 4 | import sys 5 | import math 6 | 7 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 8 | import data_utils 9 | 10 | load_fn = data_utils.load_ptsfile 11 | balance_fn = None 12 | map_fn = None 13 | keep_remainder = True 14 | save_ply_fn = None 15 | num_class = 101 16 | 17 | sample_num = 1024 18 | 19 | batch_size = 64 20 | 21 | num_epochs = 10 22 | 23 | step_val = 500 24 | 25 | learning_rate_base = 0.01 26 | decay_steps = 15000 27 | decay_rate = 0.5 28 | learning_rate_min = 1e-6 29 | 30 | weight_decay = 1e-5 31 | 32 | jitter = 0.0 33 | jitter_val = 0.0 34 | 35 | rotation_range = [0, math.pi, 0, 'u'] 36 | rotation_range_val = [0, 0, 0, 'u'] 37 | rotation_order = 'rxyz' 38 | 39 | scaling_range = [0.1, 0.1, 0.1, 'g'] 40 | scaling_range_val = [0, 0, 0, 'u'] 41 | 42 | sample_num_variance = 1 // 8 43 | sample_num_clip = 1 // 4 44 | 45 | x = 3 46 | 47 | xconv_param_name = ('K', 'D', 'P', 'C', 'links') 48 | xconv_params = [dict(zip(xconv_param_name, xconv_param)) for xconv_param in 49 | [(8, 1, -1, 16 * x, []), 50 | (12, 2, 384, 32 * x, []), 51 | (16, 2, 128, 64 * x, []), 52 | (16, 3, 128, 128 * x, [])]] 53 | 54 | with_global = True 55 | 56 | fc_param_name = ('C', 'dropout_rate') 57 | fc_params = [dict(zip(fc_param_name, fc_param)) for fc_param in 58 | [(128 * x, 0.0), 59 | (64 * 4, 0.8)]] 60 | # 全连接参数 61 | 62 | sampling = 'random' 63 | 64 | optimizer = 'adam' 65 | epsilon = 1e-2 66 | 67 | data_dim = 6 68 | use_extra_features = False 69 | with_X_transformation = True 70 | sorting_method = None 71 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/folder for trans/pointcnn.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import math 6 | import pointfly as pf 7 | import tensorflow as tf 8 | 9 | 10 | def xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, depth_multiplier, 11 | sorting_method=None, with_global=False): 12 | _, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True) 13 | indices = indices_dilated[:, :, ::D, :] 14 | 15 | if sorting_method is not None: 16 | indices = pf.sort_points(pts, indices, sorting_method) 17 | 18 | nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3) 19 | nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3) 20 | nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3) 21 | 22 | # Prepare features to be transformed 23 | nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts, tag + 'nn_fts_from_pts_0', is_training) 24 | nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts, tag + 'nn_fts_from_pts', is_training) 25 | if fts is None: 26 | nn_fts_input = nn_fts_from_pts 27 | else: 28 | nn_fts_from_prev = tf.gather_nd(fts, indices, name=tag + 'nn_fts_from_prev') 29 | nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev], axis=-1, name=tag + 'nn_fts_input') 30 | 31 | if with_X_transformation: 32 | ######################## X-transformation ######################### 33 | X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K)) 34 | X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK') 35 | X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K)) 36 | X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK') 37 | X_2 = pf.depthwise_conv2d(X_1_KK, K, tag + 'X_2', is_training, (1, K), activation=None) 38 | X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK') 39 | fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X') 40 | ################################################################### 41 | else: 42 | fts_X = nn_fts_input 43 | 44 | fts_conv = pf.separable_conv2d(fts_X, C, tag + 'fts_conv', is_training, (1, K), depth_multiplier=depth_multiplier) 45 | fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d') 46 | 47 | if with_global: 48 | fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training) 49 | fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global', is_training) 50 | return tf.concat([fts_global, fts_conv_3d], axis=-1, name=tag + 'fts_conv_3d_with_global') 51 | else: 52 | return fts_conv_3d 53 | 54 | 55 | class PointCNN: 56 | def __init__(self, points, features, is_training, setting): 57 | xconv_params = setting.xconv_params 58 | fc_params = setting.fc_params 59 | with_X_transformation = setting.with_X_transformation 60 | sorting_method = setting.sorting_method 61 | N = tf.shape(points)[0] 62 | 63 | if setting.sampling == 'fps': 64 | from sampling import tf_sampling 65 | 66 | self.layer_pts = [points] 67 | if features is None: 68 | self.layer_fts = [features] 69 | else: 70 | features = tf.reshape(features, (N, -1, setting.data_dim - 3), name='features_reshape') 71 | C_fts = xconv_params[0]['C'] // 2 72 | features_hd = pf.dense(features, C_fts, 'features_hd', is_training) 73 | self.layer_fts = [features_hd] 74 | 75 | for layer_idx, layer_param in enumerate(xconv_params): 76 | tag = 'xconv_' + str(layer_idx + 1) + '_' 77 | K = layer_param['K'] 78 | D = layer_param['D'] 79 | P = layer_param['P'] 80 | C = layer_param['C'] 81 | links = layer_param['links'] 82 | if setting.sampling != 'random' and links: 83 | print('Error: flexible links are supported only when random sampling is used!') 84 | exit() 85 | 86 | # get k-nearest points 87 | pts = self.layer_pts[-1] 88 | fts = self.layer_fts[-1] 89 | if P == -1 or (layer_idx > 0 and P == xconv_params[layer_idx - 1]['P']): 90 | qrs = self.layer_pts[-1] 91 | else: 92 | if setting.sampling == 'fps': 93 | fps_indices = tf_sampling.farthest_point_sample(P, pts) 94 | batch_indices = tf.tile(tf.reshape(tf.range(N), (-1, 1, 1)), (1, P, 1)) 95 | indices = tf.concat([batch_indices, tf.expand_dims(fps_indices,-1)], axis=-1) 96 | qrs = tf.gather_nd(pts, indices, name= tag + 'qrs') # (N, P, 3) 97 | elif setting.sampling == 'ids': 98 | indices = pf.inverse_density_sampling(pts, K, P) 99 | qrs = tf.gather_nd(pts, indices) 100 | elif setting.sampling == 'random': 101 | qrs = tf.slice(pts, (0, 0, 0), (-1, P, -1), name=tag + 'qrs') # (N, P, 3) 102 | else: 103 | print('Unknown sampling method!') 104 | exit() 105 | self.layer_pts.append(qrs) 106 | 107 | if layer_idx == 0: 108 | C_pts_fts = C // 2 if fts is None else C // 4 109 | depth_multiplier = 4 110 | else: 111 | C_prev = xconv_params[layer_idx - 1]['C'] 112 | C_pts_fts = C_prev // 4 113 | depth_multiplier = math.ceil(C / C_prev) 114 | with_global = (setting.with_global and layer_idx == len(xconv_params) - 1) 115 | fts_xconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, 116 | depth_multiplier, sorting_method, with_global) 117 | fts_list = [] 118 | for link in links: 119 | fts_from_link = self.layer_fts[link] 120 | if fts_from_link is not None: 121 | fts_slice = tf.slice(fts_from_link, (0, 0, 0), (-1, P, -1), name=tag + 'fts_slice_' + str(-link)) 122 | fts_list.append(fts_slice) 123 | if fts_list: 124 | fts_list.append(fts_xconv) 125 | self.layer_fts.append(tf.concat(fts_list, axis=-1, name=tag + 'fts_list_concat')) 126 | else: 127 | self.layer_fts.append(fts_xconv) 128 | 129 | if hasattr(setting, 'xdconv_params'): 130 | for layer_idx, layer_param in enumerate(setting.xdconv_params): 131 | tag = 'xdconv_' + str(layer_idx + 1) + '_' 132 | K = layer_param['K'] 133 | D = layer_param['D'] 134 | pts_layer_idx = layer_param['pts_layer_idx'] 135 | qrs_layer_idx = layer_param['qrs_layer_idx'] 136 | 137 | pts = self.layer_pts[pts_layer_idx + 1] 138 | fts = self.layer_fts[pts_layer_idx + 1] if layer_idx == 0 else self.layer_fts[-1] 139 | qrs = self.layer_pts[qrs_layer_idx + 1] 140 | fts_qrs = self.layer_fts[qrs_layer_idx + 1] 141 | P = xconv_params[qrs_layer_idx]['P'] 142 | C = xconv_params[qrs_layer_idx]['C'] 143 | C_prev = xconv_params[pts_layer_idx]['C'] 144 | C_pts_fts = C_prev // 4 145 | depth_multiplier = 1 146 | fts_xdconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, 147 | depth_multiplier, sorting_method) 148 | fts_concat = tf.concat([fts_xdconv, fts_qrs], axis=-1, name=tag + 'fts_concat') 149 | fts_fuse = pf.dense(fts_concat, C, tag + 'fts_fuse', is_training) 150 | self.layer_pts.append(qrs) 151 | self.layer_fts.append(fts_fuse) 152 | 153 | self.fc_layers = [self.layer_fts[-1]] 154 | for layer_idx, layer_param in enumerate(fc_params): 155 | C = layer_param['C'] 156 | dropout_rate = layer_param['dropout_rate'] 157 | fc = pf.dense(self.fc_layers[-1], C, 'fc{:d}'.format(layer_idx), is_training) 158 | fc_drop = tf.layers.dropout(fc, dropout_rate, training=is_training, name='fc{:d}_drop'.format(layer_idx)) 159 | self.fc_layers.append(fc_drop) 160 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/folder for trans/pointcnn_cls.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import pointfly as pf 6 | import tensorflow as tf 7 | from pointcnn import PointCNN 8 | 9 | 10 | class Net(PointCNN): 11 | def __init__(self, points, features, is_training, setting): 12 | PointCNN.__init__(self, points, features, is_training, setting) 13 | self.fc_layers[-1] = tf.reduce_mean(self.fc_layers[-1], axis=1, keep_dims=True, name='fc_mean') 14 | # self.fc_layers[-1] = tf.cond(is_training, lambda: self.fc_layers[-1], lambda: fc_mean) 15 | self.logits = pf.dense(self.fc_layers[-1], setting.num_class, 'logits', 16 | is_training, with_bn=False, activation=None) -------------------------------------------------------------------------------- /prerprocess/pointcnn/folder for trans/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 | import numpy as np 11 | import pickle as pickle 12 | 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | sys.path.append(BASE_DIR) 15 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 16 | def prob_sample(inp,inpr): 17 | ''' 18 | input: 19 | batch_size * ncategory float32 20 | batch_size * npoints float32 21 | returns: 22 | batch_size * npoints int32 23 | ''' 24 | return sampling_module.prob_sample(inp,inpr) 25 | ops.NoGradient('ProbSample') 26 | # TF1.0 API requires set shape in C++ 27 | #@tf.RegisterShape('ProbSample') 28 | #def _prob_sample_shape(op): 29 | # shape1=op.inputs[0].get_shape().with_rank(2) 30 | # shape2=op.inputs[1].get_shape().with_rank(2) 31 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 32 | def gather_point(inp,idx): 33 | ''' 34 | input: 35 | batch_size * ndataset * 3 float32 36 | batch_size * npoints int32 37 | returns: 38 | batch_size * npoints * 3 float32 39 | ''' 40 | return sampling_module.gather_point(inp,idx) 41 | #@tf.RegisterShape('GatherPoint') 42 | #def _gather_point_shape(op): 43 | # shape1=op.inputs[0].get_shape().with_rank(3) 44 | # shape2=op.inputs[1].get_shape().with_rank(2) 45 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 46 | @tf.RegisterGradient('GatherPoint') 47 | def _gather_point_grad(op,out_g): 48 | inp=op.inputs[0] 49 | idx=op.inputs[1] 50 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 51 | def farthest_point_sample(npoint,inp): 52 | ''' 53 | input: 54 | int32 55 | batch_size * ndataset * 3 float32 56 | returns: 57 | batch_size * npoint int32 58 | ''' 59 | return sampling_module.farthest_point_sample(inp, npoint) 60 | ops.NoGradient('FarthestPointSample') 61 | 62 | 63 | if __name__=='__main__': 64 | 65 | 66 | batch_size = 3 67 | 68 | #np.random.seed(100) 69 | triangles=np.random.rand(batch_size,5,3,3).astype('float32') 70 | #pts=np.random.rand(batch_size,1024,3,3).astype('float32') 71 | 72 | inp=tf.constant(triangles) 73 | tria=inp[:,:,0,:] 74 | trib=inp[:,:,1,:] 75 | tric=inp[:,:,2,:] 76 | 77 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 78 | randomnumbers=tf.random_uniform((batch_size,8192))#(N,8192) 79 | triids=prob_sample(areas,randomnumbers) 80 | tria_sample=gather_point(tria,triids) 81 | trib_sample=gather_point(trib,triids) 82 | tric_sample=gather_point(tric,triids) 83 | us=tf.random_uniform((batch_size,8192)) 84 | vs=tf.random_uniform((batch_size,8192)) 85 | uplusv=1-tf.abs(us+vs-1) 86 | uminusv=us-vs 87 | us=(uplusv+uminusv)*0.5 88 | vs=(uplusv-uminusv)*0.5 89 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 90 | test = farthest_point_sample(1024,pt_sample) 91 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 92 | 93 | with tf.Session() as sess: 94 | ret=sess.run(reduced_sample) 95 | pt = sess.run(pt_sample) 96 | 97 | print("tria:",tria.shape) 98 | print("areas:",areas.shape) 99 | print("triids:",triids.shape) 100 | print("tria_sample:",tria_sample.shape) 101 | print("pt_sample:",pt.shape,pt.dtype) 102 | print("test:",test.shape) 103 | print("reduced_sample",ret.shape,ret.dtype) 104 | 105 | 106 | #pickle.dump(ret,open('1.pkl','wb'),-1) 107 | print("done") 108 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/pointcnn.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import math 6 | import pointfly as pf 7 | import tensorflow as tf 8 | 9 | 10 | def xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, depth_multiplier, 11 | sorting_method=None, with_global=False): 12 | _, indices_dilated = pf.knn_indices_general(qrs, pts, K * D, True) 13 | indices = indices_dilated[:, :, ::D, :] 14 | 15 | if sorting_method is not None: 16 | indices = pf.sort_points(pts, indices, sorting_method) 17 | 18 | nn_pts = tf.gather_nd(pts, indices, name=tag + 'nn_pts') # (N, P, K, 3) 19 | nn_pts_center = tf.expand_dims(qrs, axis=2, name=tag + 'nn_pts_center') # (N, P, 1, 3) 20 | nn_pts_local = tf.subtract(nn_pts, nn_pts_center, name=tag + 'nn_pts_local') # (N, P, K, 3) 21 | 22 | # Prepare features to be transformed 23 | nn_fts_from_pts_0 = pf.dense(nn_pts_local, C_pts_fts, tag + 'nn_fts_from_pts_0', is_training) 24 | nn_fts_from_pts = pf.dense(nn_fts_from_pts_0, C_pts_fts, tag + 'nn_fts_from_pts', is_training) 25 | if fts is None: 26 | nn_fts_input = nn_fts_from_pts 27 | else: 28 | nn_fts_from_prev = tf.gather_nd(fts, indices, name=tag + 'nn_fts_from_prev') 29 | nn_fts_input = tf.concat([nn_fts_from_pts, nn_fts_from_prev], axis=-1, name=tag + 'nn_fts_input') 30 | 31 | if with_X_transformation: 32 | ######################## X-transformation ######################### 33 | X_0 = pf.conv2d(nn_pts_local, K * K, tag + 'X_0', is_training, (1, K)) 34 | X_0_KK = tf.reshape(X_0, (N, P, K, K), name=tag + 'X_0_KK') 35 | X_1 = pf.depthwise_conv2d(X_0_KK, K, tag + 'X_1', is_training, (1, K)) 36 | X_1_KK = tf.reshape(X_1, (N, P, K, K), name=tag + 'X_1_KK') 37 | X_2 = pf.depthwise_conv2d(X_1_KK, K, tag + 'X_2', is_training, (1, K), activation=None) 38 | X_2_KK = tf.reshape(X_2, (N, P, K, K), name=tag + 'X_2_KK') 39 | fts_X = tf.matmul(X_2_KK, nn_fts_input, name=tag + 'fts_X') 40 | ################################################################### 41 | else: 42 | fts_X = nn_fts_input 43 | 44 | fts_conv = pf.separable_conv2d(fts_X, C, tag + 'fts_conv', is_training, (1, K), depth_multiplier=depth_multiplier) 45 | fts_conv_3d = tf.squeeze(fts_conv, axis=2, name=tag + 'fts_conv_3d') 46 | 47 | if with_global: 48 | fts_global_0 = pf.dense(qrs, C // 4, tag + 'fts_global_0', is_training) 49 | fts_global = pf.dense(fts_global_0, C // 4, tag + 'fts_global', is_training) 50 | return tf.concat([fts_global, fts_conv_3d], axis=-1, name=tag + 'fts_conv_3d_with_global') 51 | else: 52 | return fts_conv_3d 53 | 54 | 55 | class PointCNN: 56 | def __init__(self, points, features, is_training, setting): 57 | xconv_params = setting.xconv_params 58 | fc_params = setting.fc_params 59 | with_X_transformation = setting.with_X_transformation 60 | sorting_method = setting.sorting_method 61 | N = tf.shape(points)[0] 62 | 63 | if setting.sampling == 'fps': 64 | from sampling import tf_sampling 65 | 66 | self.layer_pts = [points] 67 | if features is None: 68 | self.layer_fts = [features] 69 | else: 70 | features = tf.reshape(features, (N, -1, setting.data_dim - 3), name='features_reshape') 71 | C_fts = xconv_params[0]['C'] // 2 72 | features_hd = pf.dense(features, C_fts, 'features_hd', is_training) 73 | self.layer_fts = [features_hd] 74 | 75 | for layer_idx, layer_param in enumerate(xconv_params): 76 | tag = 'xconv_' + str(layer_idx + 1) + '_' 77 | K = layer_param['K'] 78 | D = layer_param['D'] 79 | P = layer_param['P'] 80 | C = layer_param['C'] 81 | links = layer_param['links'] 82 | if setting.sampling != 'random' and links: 83 | print('Error: flexible links are supported only when random sampling is used!') 84 | exit() 85 | 86 | # get k-nearest points 87 | pts = self.layer_pts[-1] 88 | fts = self.layer_fts[-1] 89 | if P == -1 or (layer_idx > 0 and P == xconv_params[layer_idx - 1]['P']): 90 | qrs = self.layer_pts[-1] 91 | else: 92 | if setting.sampling == 'fps': 93 | fps_indices = tf_sampling.farthest_point_sample(P, pts) 94 | batch_indices = tf.tile(tf.reshape(tf.range(N), (-1, 1, 1)), (1, P, 1)) 95 | indices = tf.concat([batch_indices, tf.expand_dims(fps_indices,-1)], axis=-1) 96 | qrs = tf.gather_nd(pts, indices, name= tag + 'qrs') # (N, P, 3) 97 | elif setting.sampling == 'ids': 98 | indices = pf.inverse_density_sampling(pts, K, P) 99 | qrs = tf.gather_nd(pts, indices) 100 | elif setting.sampling == 'random': 101 | qrs = tf.slice(pts, (0, 0, 0), (-1, P, -1), name=tag + 'qrs') # (N, P, 3) 102 | else: 103 | print('Unknown sampling method!') 104 | exit() 105 | self.layer_pts.append(qrs) 106 | 107 | if layer_idx == 0: 108 | C_pts_fts = C // 2 if fts is None else C // 4 109 | depth_multiplier = 4 110 | else: 111 | C_prev = xconv_params[layer_idx - 1]['C'] 112 | C_pts_fts = C_prev // 4 113 | depth_multiplier = math.ceil(C / C_prev) 114 | with_global = (setting.with_global and layer_idx == len(xconv_params) - 1) 115 | fts_xconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, 116 | depth_multiplier, sorting_method, with_global) 117 | fts_list = [] 118 | for link in links: 119 | fts_from_link = self.layer_fts[link] 120 | if fts_from_link is not None: 121 | fts_slice = tf.slice(fts_from_link, (0, 0, 0), (-1, P, -1), name=tag + 'fts_slice_' + str(-link)) 122 | fts_list.append(fts_slice) 123 | if fts_list: 124 | fts_list.append(fts_xconv) 125 | self.layer_fts.append(tf.concat(fts_list, axis=-1, name=tag + 'fts_list_concat')) 126 | else: 127 | self.layer_fts.append(fts_xconv) 128 | 129 | if hasattr(setting, 'xdconv_params'): 130 | for layer_idx, layer_param in enumerate(setting.xdconv_params): 131 | tag = 'xdconv_' + str(layer_idx + 1) + '_' 132 | K = layer_param['K'] 133 | D = layer_param['D'] 134 | pts_layer_idx = layer_param['pts_layer_idx'] 135 | qrs_layer_idx = layer_param['qrs_layer_idx'] 136 | 137 | pts = self.layer_pts[pts_layer_idx + 1] 138 | fts = self.layer_fts[pts_layer_idx + 1] if layer_idx == 0 else self.layer_fts[-1] 139 | qrs = self.layer_pts[qrs_layer_idx + 1] 140 | fts_qrs = self.layer_fts[qrs_layer_idx + 1] 141 | P = xconv_params[qrs_layer_idx]['P'] 142 | C = xconv_params[qrs_layer_idx]['C'] 143 | C_prev = xconv_params[pts_layer_idx]['C'] 144 | C_pts_fts = C_prev // 4 145 | depth_multiplier = 1 146 | fts_xdconv = xconv(pts, fts, qrs, tag, N, K, D, P, C, C_pts_fts, is_training, with_X_transformation, 147 | depth_multiplier, sorting_method) 148 | fts_concat = tf.concat([fts_xdconv, fts_qrs], axis=-1, name=tag + 'fts_concat') 149 | fts_fuse = pf.dense(fts_concat, C, tag + 'fts_fuse', is_training) 150 | self.layer_pts.append(qrs) 151 | self.layer_fts.append(fts_fuse) 152 | 153 | self.fc_layers = [self.layer_fts[-1]] 154 | for layer_idx, layer_param in enumerate(fc_params): 155 | C = layer_param['C'] 156 | dropout_rate = layer_param['dropout_rate'] 157 | fc = pf.dense(self.fc_layers[-1], C, 'fc{:d}'.format(layer_idx), is_training) 158 | fc_drop = tf.layers.dropout(fc, dropout_rate, training=is_training, name='fc{:d}_drop'.format(layer_idx)) 159 | self.fc_layers.append(fc_drop) 160 | -------------------------------------------------------------------------------- /prerprocess/pointcnn/pointcnn_cls.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | 5 | import pointfly as pf 6 | import tensorflow as tf 7 | from pointcnn import PointCNN 8 | 9 | 10 | class Net(PointCNN): 11 | def __init__(self, points, features, is_training, setting): 12 | PointCNN.__init__(self, points, features, is_training, setting) 13 | self.fc_layers[-1] = tf.reduce_mean(self.fc_layers[-1], axis=1, keep_dims=True, name='fc_mean') 14 | # self.fc_layers[-1] = tf.cond(is_training, lambda: self.fc_layers[-1], lambda: fc_mean) 15 | self.logits = pf.dense(self.fc_layers[-1], setting.num_class, 'logits', 16 | is_training, with_bn=False, activation=None) -------------------------------------------------------------------------------- /prerprocess/pointcnn/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 | import numpy as np 11 | import pickle as pickle 12 | 13 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 14 | sys.path.append(BASE_DIR) 15 | sampling_module=tf.load_op_library(os.path.join(BASE_DIR, 'tf_sampling_so.so')) 16 | def prob_sample(inp,inpr): 17 | ''' 18 | input: 19 | batch_size * ncategory float32 20 | batch_size * npoints float32 21 | returns: 22 | batch_size * npoints int32 23 | ''' 24 | return sampling_module.prob_sample(inp,inpr) 25 | ops.NoGradient('ProbSample') 26 | # TF1.0 API requires set shape in C++ 27 | #@tf.RegisterShape('ProbSample') 28 | #def _prob_sample_shape(op): 29 | # shape1=op.inputs[0].get_shape().with_rank(2) 30 | # shape2=op.inputs[1].get_shape().with_rank(2) 31 | # return [tf.TensorShape([shape2.dims[0],shape2.dims[1]])] 32 | def gather_point(inp,idx): 33 | ''' 34 | input: 35 | batch_size * ndataset * 3 float32 36 | batch_size * npoints int32 37 | returns: 38 | batch_size * npoints * 3 float32 39 | ''' 40 | return sampling_module.gather_point(inp,idx) 41 | #@tf.RegisterShape('GatherPoint') 42 | #def _gather_point_shape(op): 43 | # shape1=op.inputs[0].get_shape().with_rank(3) 44 | # shape2=op.inputs[1].get_shape().with_rank(2) 45 | # return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[2]])] 46 | @tf.RegisterGradient('GatherPoint') 47 | def _gather_point_grad(op,out_g): 48 | inp=op.inputs[0] 49 | idx=op.inputs[1] 50 | return [sampling_module.gather_point_grad(inp,idx,out_g),None] 51 | def farthest_point_sample(npoint,inp): 52 | ''' 53 | input: 54 | int32 55 | batch_size * ndataset * 3 float32 56 | returns: 57 | batch_size * npoint int32 58 | ''' 59 | return sampling_module.farthest_point_sample(inp, npoint) 60 | ops.NoGradient('FarthestPointSample') 61 | 62 | 63 | if __name__=='__main__': 64 | 65 | 66 | batch_size = 3 67 | 68 | #np.random.seed(100) 69 | triangles=np.random.rand(batch_size,5,3,3).astype('float32') 70 | #pts=np.random.rand(batch_size,1024,3,3).astype('float32') 71 | 72 | inp=tf.constant(triangles) 73 | tria=inp[:,:,0,:] 74 | trib=inp[:,:,1,:] 75 | tric=inp[:,:,2,:] 76 | 77 | areas=tf.sqrt(tf.reduce_sum(tf.cross(trib-tria,tric-tria)**2,2)+1e-9) 78 | randomnumbers=tf.random_uniform((batch_size,8192))#(N,8192) 79 | triids=prob_sample(areas,randomnumbers) 80 | tria_sample=gather_point(tria,triids) 81 | trib_sample=gather_point(trib,triids) 82 | tric_sample=gather_point(tric,triids) 83 | us=tf.random_uniform((batch_size,8192)) 84 | vs=tf.random_uniform((batch_size,8192)) 85 | uplusv=1-tf.abs(us+vs-1) 86 | uminusv=us-vs 87 | us=(uplusv+uminusv)*0.5 88 | vs=(uplusv-uminusv)*0.5 89 | pt_sample=tria_sample+(trib_sample-tria_sample)*tf.expand_dims(us,-1)+(tric_sample-tria_sample)*tf.expand_dims(vs,-1) 90 | test = farthest_point_sample(1024,pt_sample) 91 | reduced_sample=gather_point(pt_sample,farthest_point_sample(1024,pt_sample)) 92 | 93 | with tf.Session() as sess: 94 | ret=sess.run(reduced_sample) 95 | pt = sess.run(pt_sample) 96 | 97 | print("tria:",tria.shape) 98 | print("areas:",areas.shape) 99 | print("triids:",triids.shape) 100 | print("tria_sample:",tria_sample.shape) 101 | print("pt_sample:",pt.shape,pt.dtype) 102 | print("test:",test.shape) 103 | print("reduced_sample",ret.shape,ret.dtype) 104 | 105 | 106 | #pickle.dump(ret,open('1.pkl','wb'),-1) 107 | print("done") 108 | -------------------------------------------------------------------------------- /prerprocess/preprocess_merge_leaf.py: -------------------------------------------------------------------------------- 1 | import json 2 | import random 3 | import numpy as np 4 | import os 5 | import util 6 | from sklearn.neighbors import KDTree 7 | from open3d import * 8 | import math 9 | 10 | def getPTS(pts_file,colorPara=1): 11 | fpts = open(pts_file) 12 | count = 0 13 | while 1: 14 | line = fpts.readline() 15 | if not line: 16 | break 17 | count = count + 1 18 | if count==0: 19 | return np.zeros(9) 20 | points = np.zeros((count,9)) 21 | count = 0 22 | fpts = open(pts_file) 23 | while 1: 24 | line = fpts.readline() 25 | if not line: 26 | break 27 | L = line.split(' ') 28 | points[count,0] = float(L[0]) 29 | points[count,1] = float(L[1]) 30 | points[count,2] = float(L[2]) 31 | points[count,3] = float(L[3])/colorPara 32 | points[count,4] = float(L[4])/colorPara 33 | points[count,5] = float(L[5])/colorPara 34 | points[count,6] = float(L[6]) 35 | points[count,7] = float(L[7]) 36 | points[count,8] = float(L[8]) 37 | count = count + 1 38 | return points 39 | 40 | def computeColorBin(pts): 41 | color_bin = np.zeros(180)#r 5,g 5,b 5 42 | for i in range(0,pts.shape[0]): 43 | id0 = int(pts[i,3]*255/52) 44 | id1 = int(pts[i,4]*255/52) 45 | id2 = int(pts[i,5]*255/52) 46 | id = id0*25+id1*5+id2 47 | color_bin[id] += 1 48 | return color_bin 49 | 50 | def computeNormalAffinity(normal0,normal1,center0,center1): 51 | affinity = abs(np.dot(normal0,normal1)) 52 | if normal0[0]*normal0[0]+normal0[1]*normal0[1]+normal0[2]*normal0[2] != 0: 53 | affinity = affinity/math.sqrt(normal0[0]*normal0[0]+normal0[1]*normal0[1]+normal0[2]*normal0[2]) 54 | else: 55 | affinity = 0 56 | if normal1[0]*normal1[0]+normal1[1]*normal1[1]+normal1[2]*normal1[2] != 0: 57 | affinity = affinity/math.sqrt(normal1[0]*normal1[0]+normal1[1]*normal1[1]+normal1[2]*normal1[2]) 58 | else: 59 | affinity = 0 60 | if np.dot(center0-center1,normal0) > 0: 61 | affinity *= affinity*0.1 + 0.9 62 | else: 63 | affinity = affinity 64 | return affinity 65 | 66 | def chiSquareDis(bin0,bin1): 67 | dis = 0 68 | for i in range(bin0.shape[0]): 69 | if bin0[i] + bin1[i] > 0: 70 | dis += (bin0[i] - bin1[i])*(bin0[i] - bin1[i])/(bin0[i] + bin1[i])*0.5 71 | return dis 72 | 73 | def computeColorAffinity(bin0,bin1): 74 | dis = chiSquareDis(bin0,bin1) 75 | if dis == 0: 76 | affinity = 5 77 | else: 78 | affinity = 1/chiSquareDis(bin0,bin1)*20 79 | return affinity 80 | 81 | config = util.get_args() 82 | g_path = config.g_path 83 | scene_dir = os.path.join(g_path,'scene_normal.txt') 84 | overseg_dir = os.path.join(g_path,'overseg.txt') 85 | gt_info_dir = os.path.join(g_path,'gt_info.txt') 86 | data_dir = os.path.join(g_path,'data') 87 | leaf_dir = os.path.join(data_dir,'leaf_pts') 88 | merge_leaf_dir = os.path.join(data_dir,'merge_leaf_pts') 89 | internal_dir = os.path.join(data_dir,'internal_pts') 90 | files_dir = os.path.join(data_dir,'files') 91 | 92 | ########################################################################## 93 | ## generate merge leaf pts 94 | pts = getPTS(scene_dir) 95 | z_max = pts[:,2].max() 96 | z_min = pts[:,2].min() 97 | 98 | pts_names = os.listdir(merge_leaf_dir) 99 | pts_names.sort() 100 | count=0 101 | for pts_name in pts_names: 102 | if len(pts_name.split('.'))>1 and pts_name.split('.')[1]=='pts': 103 | count=count+1 104 | 105 | segment_pts_list = [] 106 | segment_name_list = [] 107 | segment_id_list = [] 108 | segment_pts_type_list = [] #'other':0,'floor':1,'ceiling':2 109 | trees = [] 110 | for pts_name in pts_names: 111 | if len(pts_name.split('.'))>1 and pts_name.split('.')[1]=='pts': 112 | id = int(pts_name.split('.')[0].split('_')[1]) 113 | pts = getPTS(os.path.join(merge_leaf_dir,pts_name)) 114 | if pts.shape[0] < 10: 115 | trees.append(None) 116 | segment_pts_list.append('None') 117 | segment_name_list.append(pts_name) 118 | segment_id_list.append(id) 119 | segment_pts_type_list.append(np.zeros(pts.shape[0])) 120 | continue 121 | tree = KDTree(pts[:,:3], leaf_size=2) 122 | trees.append(tree) 123 | segment_pts_list.append(pts) 124 | segment_name_list.append(pts_name) 125 | segment_id_list.append(id) 126 | segment_pts_type_list.append(np.zeros(pts.shape[0])) 127 | 128 | for i in range(len(segment_pts_list)): 129 | segment_pts = segment_pts_list[i] 130 | if segment_pts is 'None': 131 | continue 132 | for j in range(segment_pts.shape[0]): 133 | if segment_pts[j,2] > z_min and segment_pts[j,2] < z_min + 0.2: 134 | segment_pts_type_list[i][j] = 1 135 | elif segment_pts[j,2] > z_max - 0.5 and segment_pts[j,2] < z_max: 136 | segment_pts_type_list[i][j] = 2 137 | 138 | segment_type_list = [] 139 | for i in range(len(segment_pts_list)): 140 | segment_pts_type = segment_pts_type_list[i] 141 | type0 = np.where(segment_pts_type==0,1,0) 142 | type1 = np.where(segment_pts_type==1,1,0) 143 | type2 = np.where(segment_pts_type==2,1,0) 144 | num0 = type0.sum() 145 | num1 = type1.sum() 146 | num2 = type2.sum() 147 | if max(num0,num1,num2) == num0: 148 | segment_type_list.append(0) 149 | elif max(num0,num1,num2) == num1: 150 | segment_type_list.append(1) 151 | elif max(num0,num1,num2) == num2: 152 | segment_type_list.append(2) 153 | 154 | ########################################################################## 155 | ## compute affinity matrix 156 | dis_matrix = np.zeros((len(trees),len(trees))) 157 | for i in range(len(trees)): 158 | if segment_type_list[i] != 0: 159 | continue 160 | tree = trees[i] 161 | if tree == None: 162 | continue 163 | for j in range(len(segment_pts_list)): 164 | pts = segment_pts_list[j] 165 | if segment_type_list[j] != 0 or i >= j or pts is 'None': 166 | continue 167 | min_dis = 9999 168 | for k in range(pts.shape[0]): 169 | sample_prob = 0.1 170 | if random.random() > sample_prob: 171 | continue 172 | dist,ind = tree.query(pts[k:k+1,:3], k=1) 173 | if min_dis > dist[0]: 174 | min_dis = dist[0] 175 | dis_matrix[segment_id_list[i],segment_id_list[j]] = min_dis 176 | dis_matrix[segment_id_list[j],segment_id_list[i]] = min_dis 177 | 178 | segment_size_list = [] 179 | segment_cen_list = [] 180 | for i in range(len(segment_pts_list)): 181 | if segment_pts_list[i] is 'None': 182 | segment_size_list.append(0) 183 | continue 184 | segment_size_list.append(segment_pts_list[i].shape[0]) 185 | segment_cen_list.append(segment_pts_list[i].mean(0)[0:3]) 186 | 187 | segment_color_bin_list = [] 188 | for i in range(len(segment_pts_list)): 189 | if segment_pts_list[i] is 'None': 190 | segment_color_bin_list.append(np.zeros(180)) 191 | continue 192 | color_bin = computeColorBin(segment_pts_list[i]) 193 | segment_color_bin_list.append(color_bin) 194 | 195 | 196 | affinity_matrix = np.zeros((len(segment_pts_list),len(segment_pts_list))) 197 | for i in range(len(segment_pts_list)): 198 | for j in range(len(segment_pts_list)): 199 | id0 = segment_id_list[i] 200 | id1 = segment_id_list[j] 201 | if id0 >= id1: 202 | continue 203 | if segment_pts_list[i] is 'None' or segment_pts_list[j] is 'None': 204 | continue 205 | if dis_matrix[id0,id1] == 0: 206 | affinity_matrix[id0,id1] = 0 207 | affinity_matrix[id1,id0] = affinity_matrix[id0,id1] 208 | continue 209 | affinity_matrix[id0,id1] = computeColorAffinity(segment_color_bin_list[id0],segment_color_bin_list[id1]) / dis_matrix[id0,id1] * 1000 210 | if segment_cen_list[id0][2] > 2: 211 | affinity_matrix[id0,id1] *= 0.5 212 | if segment_cen_list[id1][2] > 2: 213 | affinity_matrix[id0,id1] *= 0.5 214 | affinity_matrix[id1,id0] = affinity_matrix[id0,id1] 215 | 216 | output = open(os.path.join(files_dir,'merge_segment_affinity_matrix.txt'),'w') 217 | for i in range(0,affinity_matrix.shape[0]): 218 | for j in range(0,affinity_matrix.shape[1]): 219 | if affinity_matrix[i,j] == 0: 220 | output.write('1 ') 221 | else: 222 | output.write('%f '%(affinity_matrix[i,j])) 223 | output.write('\n') 224 | output.close() -------------------------------------------------------------------------------- /prerprocess/qhull_2d.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | 3 | # Compute the convex hull of a set of 2D points 4 | # A Python implementation of the qhull algorithm 5 | # 6 | # Tested with Python 2.6.5 on Ubuntu 10.04.4 7 | 8 | # Copyright (c) 2008 Dave (www.literateprograms.org) 9 | # 10 | # Permission is hereby granted, free of charge, to any person obtaining a copy 11 | # of this software and associated documentation files (the "Software"), to deal 12 | # in the Software without restriction, including without limitation the rights 13 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 14 | # copies of the Software, and to permit persons to whom the Software is 15 | # furnished to do so, subject to the following conditions: 16 | # 17 | # The above copyright notice and this permission notice shall be included in 18 | # all copies or substantial portions of the Software. 19 | # 20 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 21 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 22 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 23 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 24 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 25 | # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS 26 | # IN THE SOFTWARE. 27 | 28 | from __future__ import division 29 | from numpy import * 30 | 31 | link = lambda a,b: concatenate((a,b[1:])) 32 | edge = lambda a,b: concatenate(([a],[b])) 33 | 34 | def qhull2D(sample): 35 | def dome(sample,base): 36 | h, t = base 37 | dists = dot(sample-h, dot(((0,-1),(1,0)),(t-h))) 38 | outer = repeat(sample, dists>0, 0) 39 | if len(outer): 40 | pivot = sample[argmax(dists)] 41 | return link(dome(outer, edge(h, pivot)), 42 | dome(outer, edge(pivot, t))) 43 | else: 44 | return base 45 | if len(sample) > 2: 46 | axis = sample[:,0] 47 | base = take(sample, [argmin(axis), argmax(axis)], 0) 48 | return link(dome(sample, base), dome(sample, base[::-1])) 49 | else: 50 | return sample 51 | 52 | -------------------------------------------------------------------------------- /prerprocess/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | from argparse import ArgumentParser 3 | 4 | def get_args(): 5 | parser = ArgumentParser(description='preprocess') 6 | parser.add_argument('--g_path', type=str, default='') 7 | parser.add_argument('--pickle_path', type=str, default='') 8 | parser.add_argument('--g_gt_path', type=str, default='') 9 | parser.add_argument('--data_path', type=str, default='') 10 | parser.add_argument('--output_path', type=str, default='') 11 | parser.add_argument('--scene_list_path', type=str, default='') 12 | args = parser.parse_args() 13 | return args 14 | -------------------------------------------------------------------------------- /vdrae/README.md: -------------------------------------------------------------------------------- 1 | # VDRAE 2 | These are the code for VDRAE. 3 | 4 | ### Dependancy 5 | The code depends on Pytorch 3. [pytorch-tools][2] should be installed. 6 | 7 | ### Pretrained model 8 | You can download the pretrained models for S3DIS and Matterport3D [here](https://www.dropbox.com/s/53fyv743lcduul0/pretrained_model.zip?dl=0). 9 | 10 | ### Data preparation 11 | You need to process the raw data before feed it to VDRAE, please see [preprocess](https://github.com/yifeishi/HierarchyLayout/tree/master/preprocess) for details. You can also download the preprocessed data [here](https://www.dropbox.com/s/kdty4kn7tlwsusv/processed_data.zip?dl=0). 12 | 13 | ### Training 14 | Once you have the preprocessed data, you can train a model by run: 15 | ~~~~ 16 | python3 train.py --data_path PROCESSED_DATA_PATH 17 | ~~~~ 18 | 19 | Arguments: 20 | ``` 21 | '--epochs' (number of epochs; default=5000) 22 | '--batch_size' (batch size; default=1) 23 | '--save_snapshot' (save snapshots of trained model; default=True) 24 | '--save_snapshot_every' (save training log for every X frames; default=100) 25 | '--no_cuda' (use cpu only; default=False) 26 | '--gpu' (device id of GPU to run cuda; default=0) 27 | '--data_path' (path of the preprocessed data) 28 | '--save_path' (trained model path; default='models') 29 | ``` 30 | 31 | ### Testing 32 | To perform the inference, run: 33 | ~~~~ 34 | python3 test.py --data_path PROCESSED_DATA_PATH --g_path DATASET_PATH --pretrained_model PRETRAEINED_MODEL_PATH 35 | ~~~~ 36 | 37 | Arguments: 38 | ``` 39 | '--no_cuda' (use cpu only; default=False) 40 | '--gpu' (device id of GPU to run cuda; default=0) 41 | '--data_path' (path of the preprocessed data) 42 | '--g_path' (path of the dataset) 43 | '--pretrained_model' (pretrained model path; default='models') 44 | ``` 45 | 46 | ### Evaluation 47 | To evaluate the results , run: 48 | ~~~~ 49 | python evaluation.py --g_path DATASET_PATH --g_gt_path DATASET_GROUNDTRUTH_PATH --ap_category THE_CATEGORY_TO_BE_EVALUATED --IOU IOU 50 | ~~~~ 51 | 52 | Arguments: 53 | ``` 54 | '--g_path' (path of the dataset) 55 | '--g_gt_path' (path of the groud-truth data) 56 | '--ap_category' (the category of evaluation) 57 | '--IOU' (iou to be used for evaluation; default=0.5) 58 | ``` 59 | 60 | ### Testing and evaluation script 61 | An example about how to run the code of testing and evaluation can be found in bat_test_eval.py 62 | 63 | [1]: https://arxiv.org/pdf/1903.03757.pdf "Hierarchy Denoising Recursive Autoencoders for 3D Scene Layout Prediction" 64 | [2]: https://github.com/nearai/torchfold "torchfold" 65 | -------------------------------------------------------------------------------- /vdrae/bat_test_eval.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | dataset = 'Matterport' 4 | data_path = '/home/net663/Downloads/yifeis/' 5 | g_path = os.path.join(data_path,dataset,'data_release') 6 | g_gt_path = os.path.join(data_path,dataset,'gt') 7 | test_scene_list_path = os.path.join(data_path,dataset,'task','test.txt') 8 | pretrained_model = config.pretrained_model 9 | cmd = 'python3 test_for_eval.py --g_path %s --pretrained_model %s --data_path %s'%(g_path, pretrained_model,'processed_data_test') 10 | os.system('%s' %cmd) 11 | 12 | if dataset == 'S3DIS': 13 | categories = [6,8,9,14] #chair,board,table,sofa 14 | elif dataset == 'Matterport': 15 | categories = [3,5,7,8,11,15,18,22,25,28] 16 | elif dataset == 'ScanNet': 17 | categories = [3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39] 18 | IOUs = [0.5] 19 | for IOU in IOUs: 20 | cmdList = [] 21 | for category in categories: 22 | cmd = 'python evaluation.py --g_path %s --g_gt_path %s --ap_category %d --IOU %f --scene_list_path %s'%(g_path,g_gt_path,category,IOU,test_scene_list_path) 23 | cmdList.append(cmd) 24 | #os.system('%s' %cmd) 25 | pool = Pool(len(categories)*len(IOUs)) 26 | pool.map(os.system, cmdList) 27 | pool.close() 28 | pool.join() -------------------------------------------------------------------------------- /vdrae/evaluation.py: -------------------------------------------------------------------------------- 1 | import util 2 | import numpy as np 3 | import os 4 | from util_obb import * 5 | from sklearn.neighbors import KDTree 6 | 7 | config = util.get_args() 8 | g_path = config.g_path 9 | g_gt_path = config.g_gt_path 10 | 11 | def computePtsIou(pts0, pts1, obbA, obbB): 12 | iou = 0 13 | aabbA = obb2Aabb(obbA) 14 | aabbB = obb2Aabb(obbB) 15 | iou_aabb = computeAabbIou(aabbA,aabbB) 16 | if iou_aabb < 0.1: 17 | iou = 0 18 | return iou 19 | tree0 = KDTree(pts0[:,:3], leaf_size=2) 20 | tree1 = KDTree(pts1[:,:3], leaf_size=2) 21 | count_in0 = 0 22 | count_in1 = 0 23 | count_all0 = 0 24 | count_all1 = 0 25 | for i in range(pts0.shape[0]): 26 | if random.random() > 0.99: 27 | continue 28 | dist,ind = tree1.query(pts0[i:i+1,:3], k=1) 29 | if dist[0] < 0.1: 30 | count_in0 += 1 31 | count_all0 += 1 32 | for i in range(pts1.shape[0]): 33 | if random.random() > 0.99: 34 | continue 35 | dist,ind = tree0.query(pts1[i:i+1,:3], k=1) 36 | if dist[0] < 0.1: 37 | count_in1 += 1 38 | count_all1 += 1 39 | intersection = (count_in0 + count_in1) * 0.5 40 | union = count_all0 + count_all1 - intersection 41 | if union == 0: 42 | iou = 0 43 | else: 44 | iou = float(intersection)/union 45 | return iou 46 | 47 | def box2AP(gt_pts_dirs_list,node_pts_dirs_list,gt_info_path_list,pred_box_path_list): 48 | gt_obbs_list = [] 49 | pred_obbs_list = [] 50 | pred_probs_list = [] 51 | valid_gt_pts_dirs_list = [] 52 | valid_pred_pts_dirs_list = [] 53 | prob_list = [] 54 | for i in range(len(gt_info_path_list)): 55 | gt_info_path = gt_info_path_list[i] 56 | pred_box_path = pred_box_path_list[i] 57 | gt_pts_dirs = gt_pts_dirs_list[i] 58 | node_pts_dirs = node_pts_dirs_list[i] 59 | 60 | lines = open(gt_info_path,'r').readlines() 61 | gt_num = 0 62 | for line in lines: 63 | L = line.split() 64 | if int(L[14]) == config.ap_category: 65 | gt_num += 1 66 | 67 | lines = open(pred_box_path,'r').readlines() 68 | pred_num = 0 69 | for line in lines: 70 | L = line.split() 71 | if float(L[9+config.ap_category]) > 0.001: 72 | pred_num += 1 73 | prob_list.append(float(L[9+config.ap_category])) 74 | 75 | gt_obbs = np.zeros((gt_num,8)) 76 | pred_obbs = np.zeros((pred_num,8)) 77 | pred_probs = np.zeros(pred_num) 78 | valid_gt_pts_dirs = []# 79 | valid_pred_pts_dirs = []# 80 | 81 | lines = open(gt_info_path,'r').readlines() 82 | gt_num = 0 83 | count = 0 84 | for line in lines: 85 | L = line.split() 86 | if int(L[14]) == config.ap_category: 87 | gt_obbs[gt_num,0] = float(L[6]) 88 | gt_obbs[gt_num,1] = float(L[7]) 89 | gt_obbs[gt_num,2] = float(L[8]) 90 | gt_obbs[gt_num,3] = float(L[9]) 91 | gt_obbs[gt_num,4] = float(L[10]) 92 | gt_obbs[gt_num,5] = float(L[11]) 93 | gt_obbs[gt_num,6] = float(L[12]) 94 | gt_obbs[gt_num,7] = float(L[13]) 95 | valid_gt_pts_dirs.append(gt_pts_dirs[count]) 96 | gt_num += 1 97 | count += 1 98 | 99 | lines = open(pred_box_path,'r').readlines() 100 | pred_num = 0 101 | count = 0 102 | for line in lines: 103 | L = line.split() 104 | if float(L[9+config.ap_category]) > 0.001: 105 | pred_obbs[pred_num,0] = float(L[0]) 106 | pred_obbs[pred_num,1] = float(L[1]) 107 | pred_obbs[pred_num,2] = float(L[2]) 108 | pred_obbs[pred_num,3] = float(L[3]) 109 | pred_obbs[pred_num,4] = float(L[4]) 110 | pred_obbs[pred_num,5] = float(L[5]) 111 | pred_obbs[pred_num,6] = float(L[6]) 112 | pred_obbs[pred_num,7] = float(L[7]) 113 | pred_probs[pred_num] = float(L[9+config.ap_category]) 114 | valid_pred_pts_dirs.append(node_pts_dirs[count]) 115 | pred_num += 1 116 | count += 1 117 | gt_obbs_list.append(gt_obbs) 118 | pred_obbs_list.append(pred_obbs) 119 | pred_probs_list.append(pred_probs) 120 | valid_gt_pts_dirs_list.append(valid_gt_pts_dirs) 121 | valid_pred_pts_dirs_list.append(valid_pred_pts_dirs) 122 | 123 | prob_list = np.array(prob_list) 124 | prob_list.sort() 125 | prob_list = np.unique(prob_list) 126 | 127 | # comput iou 128 | iou_matrix_list = [] 129 | gt_num_all = 0 130 | for j in range(len(gt_obbs_list)): 131 | gt_obbs = gt_obbs_list[j] 132 | pred_obbs = pred_obbs_list[j] 133 | pred_probs = pred_probs_list[j] 134 | valid_gt_pts_dirs = valid_gt_pts_dirs_list[j] 135 | valid_pred_pts_dirs = valid_pred_pts_dirs_list[j] 136 | iou_matrix = np.zeros((gt_obbs.shape[0],pred_obbs.shape[0])) 137 | for k in range(gt_obbs.shape[0]): 138 | for m in range(pred_obbs.shape[0]): 139 | iou = computePtsIou(getPTS(valid_gt_pts_dirs[k]),getPTS(valid_pred_pts_dirs[m]),gt_obbs[k,:],pred_obbs[m,:]) 140 | iou_matrix[k,m] = iou 141 | iou_matrix_list.append(iou_matrix) 142 | gt_num_all += len(gt_obbs) 143 | 144 | recall_list = [] 145 | precision_list = [] 146 | for i in range(-1,prob_list.shape[0]+1): 147 | TP = 0 148 | FP = 0 149 | FN = 0 150 | gt_box_num = 0 151 | pred_box_num = 0 152 | if i == -1: 153 | threshold = prob_list[i+1] * 0.5 154 | elif i == prob_list.shape[0]: 155 | threshold = (1-prob_list[i-1])*0.5+prob_list[i-1] 156 | else: 157 | threshold = prob_list[i] 158 | 159 | for j in range(len(gt_obbs_list)): 160 | iou_matrix = iou_matrix_list[j] 161 | gt_obbs = gt_obbs_list[j] 162 | pred_obbs = pred_obbs_list[j] 163 | pred_probs = pred_probs_list[j] 164 | gt_box_num += gt_obbs_list[j].shape[0] 165 | pred_box_num += pred_obbs_list[j].shape[0] 166 | 167 | # true-pos 168 | for k in range(gt_obbs.shape[0]): 169 | for m in range(pred_obbs.shape[0]): 170 | iou = iou_matrix[k,m] 171 | if iou > config.IOU and pred_probs[m] >= threshold: 172 | TP += 1 173 | break 174 | 175 | # false-pos 176 | for m in range(pred_obbs.shape[0]): 177 | if pred_probs[m] < threshold: 178 | continue 179 | FPflag = True 180 | for k in range(gt_obbs.shape[0]): 181 | iou = iou_matrix[k,m] 182 | if iou > config.IOU: 183 | FPflag = False 184 | if FPflag: 185 | FP += 1 186 | recall = float(TP)/gt_box_num 187 | if TP+FP>0: 188 | precision = float(TP)/(TP+FP) 189 | recall_list.append(recall) 190 | precision_list.append(precision) 191 | recall_list.append(0) 192 | precision_list.append(1) 193 | ap = 0 194 | length = 0 195 | for i in range(0, len(recall_list)-1): 196 | r = recall_list[i] 197 | p = precision_list[i] 198 | r_n = recall_list[i+1] 199 | p_n = precision_list[i+1] 200 | ap += abs(r - r_n)*p_n 201 | length += (r - r_n) 202 | print(config.ap_category,ap,gt_num_all) 203 | fw = open('ap.txt','a') 204 | fw.write('category:%d iou:%f ap:%f\n'%(config.ap_category,config.IOU,ap)) 205 | fw.close() 206 | 207 | scene_names = os.listdir(g_path) 208 | scene_names.sort() 209 | gt_info_path_list = [] 210 | pred_box_path_list = [] 211 | 212 | node_pts_dirs_list = [] 213 | gt_pts_dirs_list = [] 214 | 215 | f = open(scene_list_path) 216 | lines = f.readlines() 217 | scene_list = [] 218 | for line in lines: 219 | scene_list.append(line.split()[0]) 220 | 221 | for scene_name in scene_names: 222 | if scene_name.split('.')[0] not in scene_list: 223 | continue 224 | scene_dir = os.path.join(g_path,scene_name) 225 | box_dir = os.path.join(scene_dir,'data','files','pred_box_vdrae_for_eval.txt') 226 | merge_leaf_pts_dir = os.path.join(scene_dir,'data','merge_leaf_pts') 227 | internal_pts_dir = os.path.join(scene_dir,'data','internal_pts') 228 | gt_info_path = os.path.join(scene_dir,'gt_info.txt') 229 | if not os.path.exists(box_dir): 230 | continue 231 | scene_count += 1 232 | gt_info_path_list.append(gt_info_path) 233 | pred_box_path_list.append(box_dir) 234 | 235 | node_pts_dirs = [] 236 | gt_pts_dirs = [] 237 | node_pts_num = len(open(box_dir,'r').readlines()) 238 | for line in open(box_dir,'r').readlines(): 239 | L = line.split() 240 | if os.path.exists(os.path.join(merge_leaf_pts_dir,'objects_'+str(L[49])+'.pts')): 241 | node_pts_dirs.append(os.path.join(merge_leaf_pts_dir,'objects_'+str(L[49])+'.pts')) 242 | elif os.path.exists(os.path.join(internal_pts_dir,'objects_'+str(L[49])+'.pts')): 243 | node_pts_dirs.append(os.path.join(internal_pts_dir,'objects_'+str(L[49])+'.pts')) 244 | 245 | gt_pts_num = len(open(gt_info_path,'r').readlines()) 246 | for i in range(gt_pts_num): 247 | gt_pts_dirs.append(os.path.join(g_gt_path, scene_name.split('.')[0], scene_name.split('.')[1], 'Annotations', 'objects_'+str(i)+'.pts')) 248 | node_pts_dirs_list.append(node_pts_dirs) 249 | gt_pts_dirs_list.append(gt_pts_dirs) 250 | box2AP(gt_pts_dirs_list,node_pts_dirs_list,gt_info_path_list,pred_box_path_list) 251 | -------------------------------------------------------------------------------- /vdrae/grassdata.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils import data 3 | from scipy.io import loadmat 4 | from enum import Enum 5 | import numpy as np 6 | import pickle 7 | import util 8 | import numpy as np 9 | import os 10 | 11 | config = util.get_args() 12 | 13 | class Tree(object): 14 | class NodeType(Enum): 15 | Leaf = 0 16 | Internal = 1 17 | 18 | class Node(object): 19 | def __init__(self, box=None, pc_feature=None, left=None, right=None, node_type=None, sym=None, index=None, category=None, category_pred=None, reg=None): 20 | box_np = np.zeros((1,8)) 21 | reg_np = np.zeros((1,8)) 22 | for i in range(8): 23 | box_np[:,i] = box[i] 24 | reg_np[:,i] = reg[i] 25 | box = torch.from_numpy(box_np).float() 26 | reg = torch.from_numpy(reg_np).float() 27 | pc_feature = torch.from_numpy(pc_feature).float() 28 | 29 | self.box = box 30 | self.reg = reg 31 | self.pc_feature = pc_feature 32 | self.encode_feature = None 33 | self.decode_feature = None 34 | self.left = left 35 | self.right = right 36 | self.node_type = node_type 37 | self.label = torch.LongTensor([self.node_type.value]) 38 | self.index = index 39 | self.category = torch.LongTensor([category]) 40 | self.category_pred = torch.FloatTensor([category_pred]) 41 | 42 | def is_leaf(self): 43 | return self.node_type == Tree.NodeType.Leaf and self.box is not None 44 | 45 | def is_internal(self): 46 | return self.node_type == Tree.NodeType.Internal 47 | 48 | def LoadFeature(self,file): 49 | feature = np.zeros((1,256)) 50 | if file == 'empty': 51 | print('empty') 52 | return feature 53 | f=open(file) 54 | line = f.readline() 55 | L = line.split() 56 | for i in range(len(L)): 57 | feature[:,i] = float(L[i]) 58 | return feature 59 | 60 | def __init__(self, scene_dir, feature_dir, boxes, boxes_reg, categories, categories_pred, mapFather, mapChild1, mapChild2, isLeaf, index): 61 | np.set_printoptions(threshold=np.inf) 62 | validNode = np.zeros(boxes.shape[0]) 63 | mapNode = {} 64 | while 1: 65 | for id in range(boxes.shape[0]): 66 | if isLeaf[id] == 'True' and validNode[id] == 0.0: 67 | pc_feature = self.LoadFeature(feature_dir[id]) 68 | node = Tree.Node(node_type=Tree.NodeType.Leaf, box = boxes[id,:], reg=boxes_reg[id,:], pc_feature = pc_feature, index=id, category=categories[id], category_pred=categories_pred[id]) 69 | mapNode[id] = node 70 | validNode[id] = 1 71 | elif isLeaf[id] == 'False' and validNode[id] == 0.0: 72 | child1 = mapChild1[id] 73 | child2 = mapChild2[id] 74 | if mapNode.get(child1, 'False') != 'False' and mapNode.get(child2, 'False') != 'False': 75 | pc_feature = self.LoadFeature(feature_dir[id]) 76 | node = Tree.Node(node_type=Tree.NodeType.Internal, left=mapNode[child1], right=mapNode[child2], box = boxes[id,:], reg=boxes_reg[id,:], pc_feature = pc_feature, index=id, category=categories[id], category_pred=categories_pred[id]) 77 | mapNode[id] = node 78 | validNode[id] = 1 79 | stopFlag = True 80 | count = 0 81 | for id in range(boxes.shape[0]): 82 | if validNode[id] == 0: 83 | stopFlag = False 84 | count = count + 1 85 | if stopFlag: 86 | break 87 | 88 | root = [] 89 | for id in range(boxes.shape[0]): 90 | if mapFather.get(id, 'False') == 'False': 91 | root.append(id) 92 | 93 | self.root = mapNode[root[0]] 94 | self.scene_dir = scene_dir 95 | self.index = index 96 | 97 | class StrToBytes: 98 | def __init__(self, fileobj): 99 | self.fileobj = fileobj 100 | def read(self, size): 101 | return self.fileobj.read(size).encode() 102 | def readline(self, size=-1): 103 | return self.fileobj.readline(size).encode() 104 | 105 | class GRASSDataset(data.Dataset): 106 | def __init__(self, dir, transform=None): 107 | self.dir = dir 108 | self.trees = [] 109 | pkl_names = os.listdir(self.dir) 110 | new_pkl_names = [] 111 | for pkl_name in pkl_names: 112 | with open(os.path.join(self.dir,pkl_name), 'rb') as f: 113 | data = pickle.load(f) 114 | scene_dir = data['scene_dir'] 115 | new_pkl_names.append(pkl_name) 116 | self.pkl_names = new_pkl_names 117 | 118 | def __getitem__(self, index): 119 | with open(os.path.join(self.dir,self.pkl_names[index]), 'rb') as f: 120 | data = pickle.load(f) 121 | 122 | scene_dir = data['scene_dir'] 123 | feature_dir = data['feature_dir'] 124 | boxes = data['boxes'] 125 | boxes_reg = data['boxes_reg'] 126 | categories = data['category'] 127 | categories_pred = data['category_pred'] 128 | 129 | mapFather = data['mapFather'] 130 | mapChild1 = data['mapChild1'] 131 | mapChild2 = data['mapChild2'] 132 | isLeaf = data['isLeaf'] 133 | tree = Tree(scene_dir, feature_dir, boxes, boxes_reg, categories,categories_pred, mapFather, mapChild1, mapChild2, isLeaf, index) 134 | return tree 135 | 136 | def __len__(self): 137 | return len(self.pkl_names) 138 | 139 | 140 | class GRASSDatasetTest(data.Dataset): 141 | def __init__(self, dir, transform=None): 142 | self.dir = dir 143 | self.trees = [] 144 | pkl_names = os.listdir(self.dir) 145 | new_pkl_names = [] 146 | for pkl_name in pkl_names: 147 | with open(os.path.join(self.dir,pkl_name), 'rb') as f: 148 | data = pickle.load(f) 149 | scene_dir = data['scene_dir'] 150 | new_pkl_names.append(pkl_name) 151 | self.pkl_names = new_pkl_names 152 | 153 | def __getitem__(self, index): 154 | with open(os.path.join(self.dir,self.pkl_names[index]), 'rb') as f: 155 | data = pickle.load(f) 156 | 157 | scene_dir = data['scene_dir'] 158 | feature_dir = data['feature_dir'] 159 | boxes = data['boxes'] 160 | boxes_reg = data['boxes_reg'] 161 | categories = data['category'] 162 | categories_pred = data['category_pred'] 163 | 164 | mapFather = data['mapFather'] 165 | mapChild1 = data['mapChild1'] 166 | mapChild2 = data['mapChild2'] 167 | isLeaf = data['isLeaf'] 168 | tree = Tree(scene_dir, feature_dir, boxes, boxes_reg, categories, categories_pred, mapFather, mapChild1, mapChild2, isLeaf, index) 169 | return tree 170 | 171 | def __len__(self): 172 | return len(self.pkl_names) 173 | 174 | 175 | -------------------------------------------------------------------------------- /vdrae/pytorch_tools/examples/snli/spinn-example.py: -------------------------------------------------------------------------------- 1 | """This is just example model, not real implementation !!!""" 2 | 3 | import time 4 | import argparse 5 | import sys 6 | 7 | import torch 8 | from torch.autograd import Variable 9 | import torch.nn as nn 10 | from torch import optim 11 | 12 | from torchtext import data 13 | from torchtext import datasets 14 | 15 | import torchfold 16 | 17 | 18 | parser = argparse.ArgumentParser(description='SPINN') 19 | parser.add_argument('--fold', action='store_true', default=False) 20 | parser.add_argument('--no-cuda', action='store_true', default=False) 21 | parser.add_argument('--batch_size', type=int, default=128) 22 | args, _ = parser.parse_known_args(sys.argv) 23 | args.cuda = not args.no_cuda and torch.cuda.is_available() 24 | 25 | 26 | class TreeLSTM(nn.Module): 27 | def __init__(self, num_units): 28 | super(TreeLSTM, self).__init__() 29 | self.num_units = num_units 30 | self.left = nn.Linear(num_units, 5 * num_units) 31 | self.right = nn.Linear(num_units, 5 * num_units) 32 | 33 | def forward(self, left_in, right_in): 34 | lstm_in = self.left(left_in[0]) 35 | lstm_in += self.right(right_in[0]) 36 | a, i, f1, f2, o = lstm_in.chunk(5, 1) 37 | c = (a.tanh() * i.sigmoid() + f1.sigmoid() * left_in[1] + 38 | f2.sigmoid() * right_in[1]) 39 | h = o.sigmoid() * c.tanh() 40 | return h, c 41 | 42 | 43 | class SPINN(nn.Module): 44 | 45 | def __init__(self, n_classes, size, n_words): 46 | super(SPINN, self).__init__() 47 | self.size = size 48 | self.tree_lstm = TreeLSTM(size) 49 | self.embeddings = nn.Embedding(n_words, size) 50 | self.out = nn.Linear(size, n_classes) 51 | 52 | def leaf(self, word_id): 53 | return self.embeddings(word_id), Variable(torch.FloatTensor(word_id.size()[0], self.size)) 54 | 55 | def children(self, left_h, left_c, right_h, right_c): 56 | return self.tree_lstm((left_h, left_c), (right_h, right_c)) 57 | 58 | def logits(self, encoding): 59 | return self.out(encoding) 60 | 61 | 62 | def encode_tree_regular(model, tree): 63 | def encode_node(node): 64 | if node.is_leaf(): 65 | return model.leaf(Variable(torch.LongTensor([node.id]))) 66 | else: 67 | left_h, left_c = encode_node(node.left) 68 | right_h, right_c = encode_node(node.right) 69 | return model.children(left_h, left_c, right_h, right_c) 70 | encoding, _ = encode_node(tree.root) 71 | return model.logits(encoding) 72 | 73 | 74 | def encode_tree_fold(fold, tree): 75 | def encode_node(node): 76 | if node.is_leaf(): 77 | return fold.add('leaf', node.id).split(2) 78 | else: 79 | left_h, left_c = encode_node(node.left) 80 | right_h, right_c = encode_node(node.right) 81 | return fold.add('children', left_h, left_c, right_h, right_c).split(2) 82 | encoding, _ = encode_node(tree.root) 83 | return fold.add('logits', encoding) 84 | 85 | 86 | class Tree(object): 87 | class Node(object): 88 | def __init__(self, leaf=None, left=None, right=None): 89 | self.id = leaf 90 | self.left = left 91 | self.right = right 92 | 93 | def is_leaf(self): 94 | return self.id is not None 95 | 96 | def __repr__(self): 97 | return str(self.id) if self.is_leaf() else "(%s, %s)" % (self.left, self.right) 98 | 99 | def __init__(self, example, inputs_vocab, answer_vocab): 100 | self.label = answer_vocab.stoi[example.label] - 1 101 | queue = [] 102 | idx, transition_idx = 0, 0 103 | while transition_idx < len(example.premise_transitions): 104 | t = example.premise_transitions[transition_idx] 105 | transition_idx += 1 106 | if t == 'shift': 107 | queue.append(Tree.Node(leaf=inputs_vocab.stoi[example.premise[idx]])) 108 | idx += 1 109 | else: 110 | n_left = queue.pop() 111 | n_right = queue.pop() 112 | queue.append(Tree.Node(left=n_left, right=n_right)) 113 | assert len(queue) == 1 114 | self.root = queue[0] 115 | 116 | 117 | def main(): 118 | inputs = datasets.snli.ParsedTextField(lower=True) 119 | transitions = datasets.snli.ShiftReduceField() 120 | answers = data.Field(sequential=False) 121 | 122 | train, dev, test = datasets.SNLI.splits(inputs, answers, transitions) 123 | inputs.build_vocab(train, dev, test) 124 | answers.build_vocab(train) 125 | train_iter, dev_iter, test_iter = data.BucketIterator.splits( 126 | (train, dev, test), batch_size=args.batch_size, device=0 if args.cuda else -1) 127 | 128 | model = SPINN(3, 500, 1000) 129 | criterion = nn.CrossEntropyLoss() 130 | opt = optim.Adam(model.parameters(), lr=0.01) 131 | 132 | for epoch in range(10): 133 | start = time.time() 134 | iteration = 0 135 | for batch_idx, batch in enumerate(train_iter): 136 | opt.zero_grad() 137 | 138 | all_logits, all_labels = [], [] 139 | fold = torchfold.Fold(cuda=args.cuda) 140 | for example in batch.dataset: 141 | tree = Tree(example, inputs.vocab, answers.vocab) 142 | if args.fold: 143 | all_logits.append(encode_tree_fold(fold, tree)) 144 | else: 145 | all_logits.append(encode_tree_regular(model, tree)) 146 | all_labels.append(tree.label) 147 | 148 | if args.fold: 149 | res = fold.apply(model, [all_logits, all_labels]) 150 | loss = criterion(res[0], res[1]) 151 | else: 152 | loss = criterion(torch.cat(all_logits, 0), Variable(torch.LongTensor(all_labels))) 153 | loss.backward(); opt.step() 154 | 155 | iteration += 1 156 | if iteration % 10 == 0: 157 | print("Avg. Time: %fs" % ((time.time() - start) / iteration)) 158 | # iteration = 0 159 | # start = time.time() 160 | 161 | 162 | if __name__ == "__main__": 163 | main() 164 | -------------------------------------------------------------------------------- /vdrae/pytorch_tools/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | from os import path 3 | import io 4 | 5 | VERSION = "0.1.0" 6 | 7 | here = path.abspath(path.dirname(__file__)) 8 | with io.open(path.join(here, 'README.md'), encoding='utf-8') as f: 9 | long_description = f.read() 10 | 11 | 12 | setup( 13 | name='torchfold', 14 | version=VERSION, 15 | description='Dynamic Batching with PyTorch', 16 | long_description=long_description, 17 | long_description_content_type='text/markdown', 18 | packages=find_packages(exclude=["*_test.py"]), 19 | license='Apache License, Version 2.0', 20 | author='Illia Polosukhin, NEAR Inc', 21 | author_email="illia@near.ai", 22 | project_urls={ 23 | 'Blog Post': "http://near.ai/articles/2017-09-06-PyTorch-Dynamic-Batching/", 24 | 'Source': "https://github.com/nearai/torchfold", 25 | }, 26 | ) 27 | 28 | -------------------------------------------------------------------------------- /vdrae/pytorch_tools/torchfold/__init__.py: -------------------------------------------------------------------------------- 1 | from .torchfold import Fold, Unfold 2 | -------------------------------------------------------------------------------- /vdrae/pytorch_tools/torchfold/__pycache__/__init__.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/vdrae/pytorch_tools/torchfold/__pycache__/__init__.cpython-35.pyc -------------------------------------------------------------------------------- /vdrae/pytorch_tools/torchfold/__pycache__/torchfold.cpython-35.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yifeishi/HierarchyLayout/df71040a75d5b08a5ec2bb0d8294dbfa7369487d/vdrae/pytorch_tools/torchfold/__pycache__/torchfold.cpython-35.pyc -------------------------------------------------------------------------------- /vdrae/pytorch_tools/torchfold/torchfold.py: -------------------------------------------------------------------------------- 1 | import collections 2 | 3 | import torch 4 | from torch.autograd import Variable 5 | 6 | 7 | class Fold(object): 8 | 9 | class Node(object): 10 | def __init__(self, op, step, index, *args): 11 | self.op = op 12 | self.step = step 13 | self.index = index 14 | self.args = args 15 | self.split_idx = -1 16 | self.batch = True 17 | 18 | def split(self, num): 19 | """Split resulting node, if function returns multiple values.""" 20 | nodes = [] 21 | for idx in range(num): 22 | nodes.append(Fold.Node( 23 | self.op, self.step, self.index, *self.args)) 24 | nodes[-1].split_idx = idx 25 | return tuple(nodes) 26 | 27 | def nobatch(self): 28 | self.batch = False 29 | return self 30 | 31 | def get(self, values): 32 | return values[self.step][self.op].get(self.index, self.split_idx) 33 | 34 | def __repr__(self): 35 | return "[%d:%d]%s" % ( 36 | self.step, self.index, self.op) 37 | 38 | class ComputedResult(object): 39 | def __init__(self, batch_size, batched_result): 40 | self.batch_size = batch_size 41 | self.result = batched_result 42 | if isinstance(self.result, tuple): 43 | self.result = list(self.result) 44 | 45 | def try_get_batched(self, nodes): 46 | all_are_nodes = all(isinstance(n, Fold.Node) for n in nodes) 47 | num_nodes_is_equal = len(nodes) == self.batch_size 48 | if not all_are_nodes or not num_nodes_is_equal: 49 | return None 50 | 51 | valid_node_sequence = all( 52 | nodes[i].index < nodes[i + 1].index # Indices are ordered 53 | and nodes[i].split_idx == nodes[i + 1].split_idx # Same split index 54 | and nodes[i].step == nodes[i + 1].step # Same step 55 | and nodes[i].op == nodes[i + 1].op # Same op 56 | for i in range(len(nodes) - 1)) 57 | if not valid_node_sequence: 58 | return None 59 | 60 | if nodes[0].split_idx == -1 and not isinstance(self.result, tuple): 61 | return self.result 62 | elif nodes[0].split_idx >= 0 and not isinstance(self.result[nodes[0].split_idx], tuple): 63 | return self.result[nodes[0].split_idx] 64 | else: 65 | # This result was already chunked. 66 | return None 67 | 68 | def get(self, index, split_idx=-1): 69 | if split_idx == -1: 70 | if not isinstance(self.result, tuple): 71 | self.result = torch.chunk(self.result, self.batch_size) 72 | return self.result[index] 73 | else: 74 | if not isinstance(self.result[split_idx], tuple): 75 | self.result[split_idx] = torch.chunk(self.result[split_idx], self.batch_size) 76 | return self.result[split_idx][index] 77 | 78 | def __init__(self, volatile=False, cuda=False): 79 | self.steps = collections.defaultdict( 80 | lambda: collections.defaultdict(list)) 81 | self.cached_nodes = collections.defaultdict(dict) 82 | self.total_nodes = 0 83 | self.volatile = volatile 84 | self._cuda = cuda 85 | 86 | def cuda(self): 87 | self._cuda = True 88 | return self 89 | 90 | def add(self, op, *args): 91 | """Add op to the fold.""" 92 | self.total_nodes += 1 93 | if not all([isinstance(arg, ( 94 | Fold.Node, int, torch.tensor._TensorBase, Variable)) for arg in args]): 95 | raise ValueError( 96 | "All args should be Tensor, Variable, int or Node, got: %s" % str(args)) 97 | if args not in self.cached_nodes[op]: 98 | step = max([0] + [arg.step + 1 for arg in args 99 | if isinstance(arg, Fold.Node)]) 100 | node = Fold.Node(op, step, len(self.steps[step][op]), *args) 101 | self.steps[step][op].append(args) 102 | self.cached_nodes[op][args] = node 103 | return self.cached_nodes[op][args] 104 | 105 | def _batch_args(self, arg_lists, values): 106 | res = [] 107 | for arg in arg_lists: 108 | r = [] 109 | if all(isinstance(arg_item, Fold.Node) for arg_item in arg): 110 | assert all(arg[0].batch == arg_item.batch 111 | for arg_item in arg[1:]) 112 | 113 | if arg[0].batch: 114 | batched_arg = values[arg[0].step][arg[0].op].try_get_batched(arg) 115 | if batched_arg is not None: 116 | res.append(batched_arg) 117 | else: 118 | res.append( 119 | torch.cat([arg_item.get(values) 120 | for arg_item in arg], 0)) 121 | else: 122 | for arg_item in arg[1:]: 123 | if arg_item != arg[0]: 124 | raise ValueError("Can not use more then one of nobatch argument, got: %s." % str(arg_item)) 125 | res.append(arg[0].get(values)) 126 | elif all(isinstance(arg_item, int) for arg_item in arg): 127 | if self._cuda: 128 | var = Variable( 129 | torch.cuda.LongTensor(arg), volatile=self.volatile) 130 | else: 131 | var = Variable( 132 | torch.LongTensor(arg), volatile=self.volatile) 133 | res.append(var) 134 | else: 135 | for arg_item in arg: 136 | if isinstance(arg_item, Fold.Node): 137 | assert arg_item.batch 138 | r.append(arg_item.get(values)) 139 | elif isinstance(arg_item, (torch.tensor._TensorBase, Variable)): 140 | r.append(arg_item) 141 | else: 142 | raise ValueError( 143 | 'Not allowed to mix Fold.Node/Tensor with int') 144 | res.append(torch.cat(r, 0)) 145 | return res 146 | 147 | def apply(self, nn, nodes): 148 | """Apply current fold to given neural module.""" 149 | values = {} 150 | for step in sorted(self.steps.keys()): 151 | values[step] = {} 152 | for op in self.steps[step]: 153 | func = getattr(nn, op) 154 | try: 155 | batched_args = self._batch_args( 156 | zip(*self.steps[step][op]), values) 157 | except Exception: 158 | print("Error while executing node %s[%d] with args: %s" % ( 159 | op, step, self.steps[step][op][0])) 160 | raise 161 | if batched_args: 162 | arg_size = batched_args[0].size()[0] 163 | else: 164 | arg_size = 1 165 | res = func(*batched_args) 166 | values[step][op] = Fold.ComputedResult(arg_size, res) 167 | try: 168 | return self._batch_args(nodes, values) 169 | except Exception: 170 | print("Retrieving %s" % nodes) 171 | for lst in nodes: 172 | if isinstance(lst[0], Fold.Node): 173 | print(', '.join([str(x.get(values).size()) for x in lst])) 174 | raise 175 | 176 | def __str__(self): 177 | result = '' 178 | for step in sorted(self.steps.keys()): 179 | result += '%d step:\n' % step 180 | for op in self.steps[step]: 181 | first_el = '' 182 | for arg in self.steps[step][op][0]: 183 | if first_el: first_el += ', ' 184 | if isinstance(arg, (torch.tensor._TensorBase, Variable)): 185 | first_el += str(arg.size()) 186 | else: 187 | first_el += str(arg) 188 | result += '\t%s = %d x (%s)\n' % ( 189 | op, len(self.steps[step][op]), first_el) 190 | return result 191 | 192 | def __repr__(self): 193 | return str(self) 194 | 195 | 196 | class Unfold(object): 197 | """Replacement of Fold for debugging, where it does computation right away.""" 198 | 199 | class Node(object): 200 | 201 | def __init__(self, tensor): 202 | self.tensor = tensor 203 | 204 | def __repr__(self): 205 | return str(self.tensor) 206 | 207 | def nobatch(self): 208 | return self 209 | 210 | def split(self, num): 211 | return [Unfold.Node(self.tensor[i]) for i in range(num)] 212 | 213 | def __init__(self, nn, volatile=False, cuda=False): 214 | self.nn = nn 215 | self.volatile = volatile 216 | self._cuda = cuda 217 | 218 | def cuda(self): 219 | self._cuda = True 220 | return self 221 | 222 | def _arg(self, arg): 223 | if isinstance(arg, Unfold.Node): 224 | return arg.tensor 225 | elif isinstance(arg, int): 226 | if self._cuda: 227 | return Variable(torch.cuda.LongTensor([arg]), volatile=self.volatile) 228 | else: 229 | return Variable(torch.LongTensor([arg]), volatile=self.volatile) 230 | else: 231 | return arg 232 | 233 | def add(self, op, *args): 234 | values = [] 235 | for arg in args: 236 | values.append(self._arg(arg)) 237 | res = getattr(self.nn, op)(*values) 238 | return Unfold.Node(res) 239 | 240 | def apply(self, nn, nodes): 241 | if nn != self.nn: 242 | raise ValueError("Expected that nn argument passed to constructor and passed to apply would match.") 243 | result = [] 244 | for n in nodes: 245 | result.append(torch.cat([self._arg(a) for a in n])) 246 | return result 247 | -------------------------------------------------------------------------------- /vdrae/pytorch_tools/torchfold/torchfold_test.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Variable 3 | import torch.nn as nn 4 | 5 | import torchfold 6 | 7 | import mock 8 | import unittest 9 | 10 | 11 | class TestEncoder(nn.Module): 12 | 13 | def __init__(self): 14 | super(TestEncoder, self).__init__() 15 | self.embed = nn.Embedding(10, 10) 16 | self.out = nn.Linear(20, 10) 17 | 18 | def concat(self, *nodes): 19 | return torch.cat(nodes, 0) 20 | 21 | def value(self, idx): 22 | return self.embed(idx) 23 | 24 | def value2(self, idx): 25 | return self.embed(idx), self.embed(idx) 26 | 27 | def attr(self, left, right): 28 | return self.out(torch.cat([left, right], 1)) 29 | 30 | def logits(self, enc, embed): 31 | return torch.mm(enc, embed.t()) 32 | 33 | 34 | class TorchFoldTest(unittest.TestCase): 35 | 36 | def test_rnn(self): 37 | f = torchfold.Fold() 38 | v1, _ = f.add('value2', 1).split(2) 39 | v2, _ = f.add('value2', 2).split(2) 40 | r = v1 41 | for i in range(1000): 42 | r = f.add('attr', v1, v2) 43 | r = f.add('attr', r, v2) 44 | 45 | te = TestEncoder() 46 | enc = f.apply(te, [[r]]) 47 | self.assertEqual(enc[0].size(), (1, 10)) 48 | 49 | def test_nobatch(self): 50 | f = torchfold.Fold() 51 | v = [] 52 | for i in range(15): 53 | v.append(f.add('value', i % 10)) 54 | d = f.add('concat', *v).nobatch() 55 | res = [] 56 | for i in range(100): 57 | res.append(f.add('logits', v[i % 10], d)) 58 | 59 | te = TestEncoder() 60 | enc = f.apply(te, [res]) 61 | self.assertEqual(len(enc), 1) 62 | self.assertEqual(enc[0].size(), (100, 15)) 63 | 64 | 65 | class RNNEncoder(nn.Module): 66 | 67 | def __init__(self, num_units, input_size): 68 | super(RNNEncoder, self).__init__() 69 | self.num_units = num_units 70 | self.input_size = input_size 71 | self.encoder = nn.GRUCell(self.input_size, self.num_units) 72 | 73 | def encode(self, input_, state): 74 | return self.encoder(input_, state) 75 | 76 | 77 | class TestRNNBatching(unittest.TestCase): 78 | 79 | def setUp(self): 80 | torch.manual_seed(42) 81 | self.input_size = 5 82 | self.num_units = 4 83 | 84 | def _generate_variable(self, dim): 85 | t = torch.Tensor(1, dim).uniform_(0, 1) 86 | return Variable(t) 87 | 88 | def test_rnn_optimized_chunking(self): 89 | seq_lengths = [2, 3, 5] 90 | 91 | states = [] 92 | for _ in xrange(len(seq_lengths)): 93 | states.append(self._generate_variable(self.num_units)) 94 | 95 | f = torchfold.Fold() 96 | for seq_ind in xrange(len(seq_lengths)): 97 | for _ in xrange(seq_lengths[seq_ind]): 98 | states[seq_ind] = f.add('encode', self._generate_variable(self.input_size), states[seq_ind]) 99 | 100 | enc = RNNEncoder(self.num_units, self.input_size) 101 | with mock.patch.object(torch, 'chunk', wraps=torch.chunk) as wrapped_chunk: 102 | result = f.apply(enc, [states]) 103 | # torch.chunk is called 3 times instead of max(seq_lengths)=5. 104 | self.assertEquals(3, wrapped_chunk.call_count) 105 | self.assertEqual(len(result), 1) 106 | self.assertEqual(result[0].size(), (len(seq_lengths), self.num_units)) 107 | 108 | 109 | if __name__ == "__main__": 110 | unittest.main() 111 | -------------------------------------------------------------------------------- /vdrae/test_for_eval.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import util 3 | import grassmodel 4 | from grassdata import GRASSDatasetTest 5 | from grassmodel import GRASSEncoderDecoder 6 | import os 7 | import random 8 | import pickle 9 | from util_obb import * 10 | from sklearn.neighbors import KDTree 11 | 12 | 13 | def NMS(predictions,obbs,labels): 14 | visited = np.zeros(len(obbs)) 15 | flags = np.ones(len(obbs)) 16 | 17 | insidePoints = [] 18 | for i in range(len(obbs)): 19 | insidePoints.append(samplePoints(obbs[i])) 20 | 21 | while True: 22 | max_probs = np.zeros(len(obbs)) 23 | for i in range(len(predictions)): 24 | prediction = predictions[i] 25 | label = np.where(prediction == np.max(prediction))[0] 26 | if visited[i] == 0: 27 | max_probs[i] = prediction[label] 28 | index = np.where(max_probs == np.max(max_probs))[0] 29 | index_one = -1 30 | for i in range(len(index)): 31 | if visited[index[i]] == 0: 32 | index_one = index[i] 33 | break 34 | for i in range(len(predictions)): 35 | if i == index_one: 36 | flags[i] = 1 37 | visited[i] = 1 38 | continue 39 | if visited[i] == 1: 40 | continue 41 | if labels[i] != labels[index_one]: 42 | continue 43 | iou = computeObbIou(obbs[index_one],obbs[i],insidePoints[index_one],insidePoints[i]) 44 | IOU = 0.5 45 | if iou > IOU: 46 | flags[i] = 0 47 | visited[i] = 1 48 | stopFlag = 1 49 | num = 0 50 | for i in range(len(predictions)): 51 | if visited[i] == 0: 52 | stopFlag = 0 53 | num += 1 54 | if stopFlag: 55 | break 56 | return flags 57 | 58 | def computePtsIou(pts0, pts1, obbA, obbB): 59 | iou = 0 60 | aabbA = obb2Aabb(obbA) 61 | aabbB = obb2Aabb(obbB) 62 | iou_aabb = computeAabbIou(aabbA,aabbB) 63 | if iou_aabb < 0.1: 64 | iou = 0 65 | return iou 66 | tree0 = KDTree(pts0[:,:3], leaf_size=2) 67 | tree1 = KDTree(pts1[:,:3], leaf_size=2) 68 | count_in0 = 0 69 | count_in1 = 0 70 | count_all0 = 0 71 | count_all1 = 0 72 | for i in range(pts0.shape[0]): 73 | if random.random() > 0.05: 74 | continue 75 | dist,ind = tree1.query(pts0[i:i+1,:3], k=1) 76 | if dist[0] < 0.1: 77 | count_in0 += 1 78 | count_all0 += 1 79 | for i in range(pts1.shape[0]): 80 | if random.random() > 0.05: 81 | continue 82 | dist,ind = tree0.query(pts1[i:i+1,:3], k=1) 83 | if dist[0] < 0.1: 84 | count_in1 += 1 85 | count_all1 += 1 86 | intersection = (count_in0 + count_in1) * 0.5 87 | union = count_all0 + count_all1 - intersection 88 | if union == 0: 89 | iou = 0 90 | else: 91 | iou = float(intersection)/union 92 | return iou 93 | 94 | def NMS_on_points(predictions,obbs,pts_dirs,labels): 95 | visited = np.zeros(len(obbs)) 96 | flags = np.ones(len(obbs)) 97 | 98 | countLabelCorrect = 0 99 | for i in range(len(labels)): 100 | prediction = predictions[i] 101 | validFlag = False 102 | #print('prediction',prediction) 103 | for category in categories: 104 | if prediction[category] > 0.01: 105 | validFlag = True 106 | break 107 | if validFlag == False: 108 | flags[i] = 0 109 | visited[i] = 1 110 | else: 111 | countLabelCorrect += 1 112 | print('node before NMS',countLabelCorrect) 113 | if len(obbs) == 0 or countLabelCorrect == 0: 114 | flags = np.zeros(len(obbs)) 115 | return flags 116 | 117 | while True: 118 | #print('while') 119 | max_probs = np.zeros(len(obbs)) 120 | for i in range(len(predictions)): 121 | prediction = predictions[i] 122 | label = int(labels[i]) 123 | if visited[i] == 0: 124 | max_probs[i] = prediction[label] 125 | index = np.where(max_probs == np.max(max_probs))[0] 126 | index_one = -1 127 | for i in range(len(index)): 128 | if visited[index[i]] == 0: 129 | index_one = index[i] 130 | break 131 | for i in range(len(predictions)): 132 | if i == index_one: 133 | flags[i] = 1 134 | visited[i] = 1 135 | continue 136 | if visited[i] == 1: 137 | continue 138 | if labels[i] != labels[index_one]: 139 | continue 140 | #iou = computeObbIou(obbs[index_one],obbs[i],insidePoints[index_one],insidePoints[i]) 141 | iou = computePtsIou(getPTS(pts_dirs[index_one]),getPTS(pts_dirs[i]),obbs[index_one],obbs[i]) 142 | IOU = 0.5 143 | if iou > IOU: 144 | flags[i] = 0 145 | visited[i] = 1 146 | stopFlag = 1 147 | num = 0 148 | for i in range(len(predictions)): 149 | if visited[i] == 0: 150 | stopFlag = 0 151 | num += 1 152 | if stopFlag: 153 | break 154 | return flags 155 | 156 | class StrToBytes: 157 | def __init__(self, fileobj): 158 | self.fileobj = fileobj 159 | def read(self, size): 160 | return self.fileobj.read(size).encode() 161 | def readline(self, size=-1): 162 | return self.fileobj.readline(size).encode() 163 | 164 | def pickle2predbox(pickle_dir,box_dir,merge_leaf_pts_dir,internal_pts_dir): 165 | fo = open(pickle_dir, 'rb+') 166 | data = pickle.load(fo,encoding = 'ios-8859-1') 167 | gts = data['gts'] 168 | predictions = data['predictions'] 169 | obbs = data['obbs'] 170 | ids = data['ids'] 171 | labels = data['labels'] 172 | count = data['count'] 173 | 174 | pts_dirs = [] 175 | for i in range(count): 176 | if os.path.exists(os.path.join(merge_leaf_pts_dir,'objects_'+str(ids[i])+'.pts')): 177 | pts_dirs.append(os.path.join(merge_leaf_pts_dir,'objects_'+str(ids[i])+'.pts')) 178 | elif os.path.exists(os.path.join(internal_pts_dir,'objects_'+str(ids[i])+'.pts')): 179 | pts_dirs.append(os.path.join(internal_pts_dir,'objects_'+str(ids[i])+'.pts')) 180 | flags = NMS_on_points(predictions,obbs,pts_dirs,labels) 181 | print('node after NMS', int(flags.sum())) 182 | fw = open(box_dir,'w') 183 | correct = wrong = 0 184 | count1 = 0 185 | for i in range(count): 186 | if flags[i] == 0: 187 | continue 188 | prediction = predictions[i] 189 | segment_id = ids[i] 190 | pred_array = prediction[0:40] 191 | label = np.where(pred_array == np.max(pred_array))[0] 192 | obb = obbs[i] 193 | if label[0] == gts[i]: 194 | correct += 1 195 | else: 196 | wrong += 1 197 | fw.write('%f %f %f %f %f %f %f %f %d %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %f %d\n'\ 198 | %(obb[0], obb[1], obb[2], obb[3], obb[4], obb[5], obb[6], obb[7], gts[i],\ 199 | pred_array[0], pred_array[1], pred_array[2], pred_array[3], pred_array[4], 200 | pred_array[5], pred_array[6], pred_array[7], pred_array[8], pred_array[9], 201 | pred_array[10], pred_array[11], pred_array[12], pred_array[13], pred_array[14], 202 | pred_array[15], pred_array[16], pred_array[17], pred_array[18], pred_array[19], 203 | pred_array[20], pred_array[21], pred_array[22], pred_array[23], pred_array[24], 204 | pred_array[25], pred_array[26], pred_array[27], pred_array[28], pred_array[29], 205 | pred_array[30], pred_array[31], pred_array[32], pred_array[33], pred_array[34], 206 | pred_array[35], pred_array[36], pred_array[37], pred_array[38], pred_array[39], segment_id)) 207 | count1 += 1 208 | 209 | 210 | config = util.get_args() 211 | dataset = config.dataset 212 | if dataset == 'S3DIS': 213 | categories = [6,8,9,14] #chair,board,table,sofa 214 | elif dataset == 'Matterport': 215 | categories = [3,5,7,8,11,15,18,22,25,28] 216 | elif dataset == 'ScanNet': 217 | categories = [3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39] 218 | 219 | g_path = config.g_path 220 | config.cuda = not config.no_cuda 221 | if config.gpu<0 and config.cuda: 222 | config.gpu = 0 223 | if config.cuda: 224 | torch.cuda.set_device(config.gpu) 225 | pretrained_model = config.pretrained_model 226 | encoder_decoder = torch.load(pretrained_model) 227 | if config.cuda: 228 | encoder_decoder.cuda(config.gpu) 229 | 230 | grass_data = GRASSDatasetTest(config.data_path) 231 | def my_collate(batch): 232 | return batch 233 | test_iter = torch.utils.data.DataLoader(grass_data, batch_size=1, shuffle=False, collate_fn=my_collate) 234 | 235 | for batch_idx, batch in enumerate(test_iter): 236 | print(os.path.join(g_path,batch[0].scene_dir.split('/')[-1],'data','files','pred.pickle')) 237 | pred_pickle_path = os.path.join(g_path,batch[0].scene_dir.split('/')[-1],'data','files','pred.pickle') 238 | gts, predictions, obbs, labels, ids, count = grassmodel.encode_decode_structure_eval(encoder_decoder, batch[0]) 239 | mdict={'gts': gts, 'predictions': predictions, 'obbs': obbs, 'labels': labels,'ids':ids, 'count': count} 240 | with open(pred_pickle_path, 'wb+') as f: 241 | pickle.dump(mdict,f) 242 | pred_box_path = os.path.join(g_path,batch[0].scene_dir.split('/')[-1],'data','files','pred_box_vdrae_for_eval.txt') 243 | merge_leaf_pts_dir = os.path.join(g_path,batch[0].scene_dir.split('/')[-1],'data','merge_leaf_pts') 244 | internal_pts_dir = os.path.join(g_path,batch[0].scene_dir.split('/')[-1],'data','internal_pts') 245 | pickle2predbox(pred_pickle_path,pred_box_path,merge_leaf_pts_dir,internal_pts_dir) -------------------------------------------------------------------------------- /vdrae/torchfoldext.py: -------------------------------------------------------------------------------- 1 | from pytorch_tools import torchfold 2 | from pytorch_tools.torchfold import Fold 3 | import torch 4 | from torch.autograd import Variable 5 | 6 | 7 | class FoldExt(Fold): 8 | 9 | def __init__(self, volatile=False, cuda=False): 10 | Fold.__init__(self, volatile, cuda) 11 | 12 | 13 | def add(self, op, *args): 14 | """Add op to the fold.""" 15 | self.total_nodes += 1 16 | if not all([isinstance(arg, ( 17 | Fold.Node, int, torch.Tensor, torch.FloatTensor, torch.LongTensor, Variable)) for arg in args]): 18 | raise ValueError( 19 | "All args should be Tensor, Variable, int or Node, got: %s" % str(args)) 20 | if args not in self.cached_nodes[op]: 21 | step = max([0] + [arg.step + 1 for arg in args 22 | if isinstance(arg, Fold.Node)]) 23 | node = Fold.Node(op, step, len(self.steps[step][op]), *args) 24 | self.steps[step][op].append(args) 25 | self.cached_nodes[op][args] = node 26 | return self.cached_nodes[op][args] 27 | 28 | 29 | def _batch_args(self, arg_lists, values): 30 | res = [] 31 | for arg in arg_lists: 32 | r = [] 33 | # print('arg[0]') 34 | # print(arg[0]) 35 | # print('Fold.Node') 36 | # print(Fold.Node) 37 | if isinstance(arg[0], Fold.Node): 38 | if arg[0].batch: 39 | for x in arg: 40 | r.append(x.get(values)) 41 | res.append(torch.cat(r, 0)) 42 | else: 43 | for i in range(2, len(arg)): 44 | if arg[i] != arg[0]: 45 | raise ValueError("Can not use more then one of nobatch argument, got: %s." % str(arg)) 46 | x = arg[0] 47 | res.append(x.get(values)) 48 | else: 49 | # Below is what this extension changes against the original version: 50 | # We make Fold handle float tensor 51 | try: 52 | if (isinstance(arg[0], Variable)): 53 | var = torch.cat(arg, 0) 54 | else: 55 | var = Variable(torch.cat(arg, 0), volatile=self.volatile) 56 | if self._cuda: 57 | var = var.cuda() 58 | res.append(var) 59 | except: 60 | print("Constructing float tensor from %s" % str(arg)) 61 | raise 62 | return res -------------------------------------------------------------------------------- /vdrae/train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | from time import gmtime, strftime 4 | from datetime import datetime 5 | import torch 6 | from torch import nn 7 | from torch.autograd import Variable 8 | import torch.utils.data 9 | from torchfoldext import FoldExt 10 | from grassdata import GRASSDataset 11 | from grassmodel import GRASSEncoderDecoder 12 | import grassmodel 13 | import util 14 | 15 | config = util.get_args() 16 | config.cuda = not config.no_cuda 17 | if config.gpu < 0 and config.cuda: 18 | config.gpu = 0 19 | torch.cuda.set_device(config.gpu) 20 | if config.cuda and torch.cuda.is_available(): 21 | print("Using CUDA on GPU ", config.gpu) 22 | else: 23 | print("Not using CUDA.") 24 | 25 | #pretrained_model = config.pretrained_model 26 | #encoder_decoder = torch.load(pretrained_model) 27 | encoder_decoder = GRASSEncoderDecoder(config) 28 | if config.cuda: 29 | encoder_decoder.cuda(config.gpu) 30 | 31 | print("Loading data ......") 32 | grass_data = GRASSDataset(config.data_path) 33 | def my_collate(batch): 34 | return batch 35 | train_iter = torch.utils.data.DataLoader(grass_data, batch_size=config.batch_size, shuffle=True, collate_fn=my_collate) 36 | 37 | print("Start training ......") 38 | start = time.time() 39 | 40 | if config.save_snapshot: 41 | if not os.path.exists(config.save_path): 42 | os.makedirs(config.save_path) 43 | snapshot_folder = os.path.join(config.save_path, 'snapshots_'+strftime("%Y-%m-%d_%H-%M-%S",gmtime())+'_'+config.nick_name) 44 | if not os.path.exists(snapshot_folder): 45 | os.makedirs(snapshot_folder) 46 | 47 | header = ' Time Epoch Iteration Progress(%) LR TotalLoss' 48 | print(header) 49 | log_template = ' '.join('{:>9s},{:>5.0f}/{:<5.0f},{:>5.0f}/{:<5.0f},{:>9.1f}%, {:>11.9f},{:>10.20f}'.split(',')) 50 | 51 | learning_rate = config.lr 52 | for epoch in range(config.epochs): 53 | encoder_decoder_opt = torch.optim.Adam(encoder_decoder.parameters()) 54 | for batch_idx, batch in enumerate(train_iter): 55 | enc_fold = FoldExt(cuda=config.cuda) 56 | enc_dec_fold_nodes = [] 57 | enc_dec_recon_fold_nodes = [] 58 | enc_dec_label_fold_nodes = [] 59 | for example in batch: 60 | enc_dec_fold_nodes.append(grassmodel.encode_decode_structure_fold(enc_fold, example)) 61 | total_loss = enc_fold.apply(encoder_decoder, [enc_dec_fold_nodes]) 62 | sum_loss = total_loss[0].sum() 63 | encoder_decoder_opt.zero_grad() 64 | sum_loss.backward() 65 | encoder_decoder_opt.step() 66 | if batch_idx % config.show_log_every == 0: 67 | print(log_template.format(strftime("%H:%M:%S",time.gmtime(time.time()-start)), 68 | epoch, config.epochs, 1+batch_idx, len(train_iter), 69 | 100. * (1+batch_idx+len(train_iter)*epoch) / (len(train_iter)*config.epochs), 70 | learning_rate, sum_loss)) 71 | if config.save_snapshot and (epoch+0) % config.save_snapshot_every == 0 : 72 | print("Saving snapshots of the models ...... ") 73 | torch.save(encoder_decoder, snapshot_folder+'/encoder_decoder_model_epoch_{}_loss_{:.4f}.pkl'.format(epoch+0, sum_loss.data[0])) 74 | 75 | print("Saving final models ...... ") 76 | torch.save(encoder_decoder, snapshot_folder+'/encoder_decoder_model.pkl') 77 | print("DONE") 78 | -------------------------------------------------------------------------------- /vdrae/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | from argparse import ArgumentParser 3 | 4 | def get_args(): 5 | parser = ArgumentParser(description='vdrae') 6 | parser.add_argument('--obj_code_size', type=int, default=256+8) 7 | parser.add_argument('--box_code_size', type=int, default=8) 8 | parser.add_argument('--box_size', type=int, default=8) 9 | parser.add_argument('--feature_size', type=int, default=1000) 10 | parser.add_argument('--pc_feature_size', type=int, default=256) 11 | parser.add_argument('--hidden_size', type=int, default=1000) 12 | parser.add_argument('--layout_input_size', type=int, default=8) 13 | parser.add_argument('--layout_code_size', type=int, default=7) 14 | parser.add_argument('--layout_feature_size', type=int, default=100) 15 | parser.add_argument('--symmetry_size', type=int, default=8) 16 | parser.add_argument('--max_box_num', type=int, default=50) 17 | parser.add_argument('--max_sym_num', type=int, default=10) 18 | 19 | parser.add_argument('--epochs', type=int, default=5000) 20 | parser.add_argument('--batch_size', type=int, default=1) 21 | parser.add_argument('--show_log_every', type=int, default=1) 22 | parser.add_argument('--save_log', action='store_true', default=False) 23 | parser.add_argument('--save_log_every', type=int, default=20) 24 | parser.add_argument('--save_snapshot', action='store_true', default=True) 25 | parser.add_argument('--save_snapshot_every', type=int, default=5) 26 | parser.add_argument('--no_plot', action='store_true', default=True) 27 | parser.add_argument('--lr', type=float, default=.00001) 28 | parser.add_argument('--lr_decay_by', type=float, default=1) 29 | parser.add_argument('--lr_decay_every', type=float, default=20) 30 | parser.add_argument('--validate', action='store_true', default=True) 31 | parser.add_argument('--validate_every', type=float, default=20) 32 | parser.add_argument('--no_cuda', action='store_true', default=False) 33 | parser.add_argument('--gpu', type=int, default=0) 34 | 35 | parser.add_argument('--g_path', type=str, default='data') 36 | parser.add_argument('--data_path', type=str, default='processed_data') 37 | parser.add_argument('--scene_list_path', type=str, default='test_scene_list_path') 38 | parser.add_argument('--save_path', type=str, default='models') 39 | parser.add_argument('--nick_name', type=str, default='') 40 | parser.add_argument('--resume_snapshot', type=str, default='') 41 | parser.add_argument('--pretrained_model', type=str, default='./models/snapshots_2019-04-04_15-09-06_/encoder_decoder_model_epoch_260_loss_0.0659.pkl') 42 | parser.add_argument('--room_type', type=str, default='Area_1') 43 | parser.add_argument('--ap_category', type=int, default=6) 44 | parser.add_argument('--IOU', type=float, default=0.5) 45 | 46 | args = parser.parse_args() 47 | return args 48 | --------------------------------------------------------------------------------