00001 
00002 
00003 
00004 
00005 
00006 
00007 
00008 
00009 
00010 
00011 
00012 
00013 
00014 
00015 
00016 
00017 
00018 
00019 
00020 
00021 
00022 
00023 
00024 
00025 
00026 
00027 
00028 
00029 
00030 import sys, optparse, os
00031 import time
00032 
00033 import math, numpy as np 
00034 import scipy.ndimage as ni
00035 import copy
00036 
00037 import hrl_lib.util as ut, hrl_lib.transforms as tr
00038 
00039 
00040 
00041 
00042 
00043 
00044 
00045 
00046 
00047 
00048 
00049 
00050 
00051 
00052 
00053 def subtract(og1,og2):
00054 
00055     if np.all(og1.resolution==og2.resolution) == False:
00056         print 'occupancy_grid_3d.subtract: The resolution of the two grids is not the same.'
00057         print 'res1, res2:', og1.resolution.A1.tolist(), og2.resolution.A1.tolist()
00058         print 'Exiting...'
00059         sys.exit()
00060 
00061     sub_tlb = og2.tlb
00062     sub_brf = og2.brf
00063 
00064     sub_tlb_idx = np.round((sub_tlb-og1.brf)/og1.resolution)
00065     sub_brf_idx = np.round((sub_brf-og1.brf)/og1.resolution)
00066 
00067     x_s,y_s,z_s = int(sub_brf_idx[0,0]),int(sub_brf_idx[1,0]),int(sub_brf_idx[2,0])
00068     x_e,y_e,z_e = int(sub_tlb_idx[0,0]),int(sub_tlb_idx[1,0]),int(sub_tlb_idx[2,0])
00069 
00070     x_e = min(x_e+1,og1.grid_shape[0,0])
00071     y_e = min(y_e+1,og1.grid_shape[1,0])
00072     z_e = min(z_e+1,og1.grid_shape[2,0])
00073     sub_grid = og1.grid[x_s:x_e,y_s:y_e,z_s:z_e]
00074 
00075     if np.any(og1.grid_shape!=og2.grid_shape):
00076         print '#############################################################################'
00077         print 'WARNING: occupancy_grid_3d.subtract has not been tested for grids of different sizes.'
00078         print '#############################################################################'
00079     sub_grid = sub_grid-og2.grid
00080     sub_grid = np.abs(sub_grid) 
00081     og1.grid[x_s:x_e,y_s:y_e,z_s:z_e] = sub_grid
00082     idxs = np.where(sub_grid>=1)
00083     shp = og2.grid_shape
00084     list_idxs = (idxs[0]+idxs[1]*shp[0,0]+idxs[2]*shp[0,0]*shp[1,0]).tolist()
00085 
00086     og1_list_idxs = (idxs[0]+x_s+(idxs[1]+y_s)*shp[0,0]+(idxs[2]+z_s)*shp[0,0]*shp[1,0]).tolist()
00087     og1_list_len = len(og1.grid_points_list)
00088     for og1_pts_idxs,pts_idxs in zip(og1_list_idxs,list_idxs):
00089         if og1_pts_idxs<og1_list_len:
00090             og1.grid_points_list[og1_pts_idxs] += og2.grid_points_list[pts_idxs]
00091 
00092 
00093 class occupancy_grid_3d():
00094 
00095     
00096     
00097     
00098     
00099     
00100     def __init__(self, brf, tlb, resolution, rotation_z=math.radians(0.)):
00101         
00102         self.grid = np.zeros(np.round((tlb-brf)/resolution).astype('int')+1,dtype='int')
00103 
00104         self.tlb = tlb
00105         self.brf = brf
00106         self.grid_shape = np.matrix(self.grid.shape).T
00107         self.resolution = resolution
00108 
00109         n_cells = self.grid.shape[0]*self.grid.shape[1]*self.grid.shape[2]
00110         self.grid_points_list = [[] for i in range(n_cells)]
00111         self.rotation_z = rotation_z
00112 
00113     
00114     
00115     
00116     def grid_lines(self, rotation_angle=0.):
00117         grid_size = np.multiply(self.grid_shape,self.resolution)
00118         rot_mat = tr.rotZ(rotation_angle)
00119         p5 = self.tlb
00120         p6 = p5+np.matrix([0.,-grid_size[1,0],0.]).T
00121         p8 = p5+np.matrix([0.,0.,-grid_size[2,0]]).T
00122         p7 = p8+np.matrix([0.,-grid_size[1,0],0.]).T
00123 
00124         p3 = self.brf
00125         p4 = p3+np.matrix([0.,grid_size[1,0],0.]).T
00126         p2 = p3+np.matrix([0.,0.,grid_size[2,0]]).T
00127         p1 = p2+np.matrix([0.,grid_size[1,0],0.]).T
00128 
00129         p1 = rot_mat*p1
00130         p2 = rot_mat*p2
00131         p3 = rot_mat*p3
00132         p4 = rot_mat*p4
00133         p5 = rot_mat*p5
00134         p6 = rot_mat*p6
00135         p7 = rot_mat*p7
00136         p8 = rot_mat*p8
00137 
00138         l = [(p1,p2),(p1,p4),(p2,p3),(p3,p4),(p5,p6),(p6,p7),(p7,p8),(p8,p5),(p1,p5),(p2,p6),(p4,p8),(p3,p7)]
00139         
00140         return l
00141 
00142     
00143     
00144     
00145     
00146     
00147     def fill_grid(self,pts,ignore_z=False):
00148         if ignore_z:
00149             idx = np.where(np.min(np.multiply(pts[0:2,:]>self.brf[0:2,:],
00150                                               pts[0:2,:]<self.tlb[0:2,:]),0))[1]
00151         else:
00152             idx = np.where(np.min(np.multiply(pts[0:3,:]>self.brf,pts[0:3,:]<self.tlb),0))[1]
00153 
00154         if idx.shape[1] == 0:
00155             print 'aha!'
00156             return
00157         pts = pts[:,idx.A1.tolist()]
00158 
00159         
00160         p_all = np.round((pts[0:3,:]-self.brf)/self.resolution)
00161         
00162         pts[0:3,:] = tr.Rz(self.rotation_z).T*pts[0:3,:]
00163 
00164         for i,p in enumerate(p_all.astype('int').T):
00165             if ignore_z:
00166                 p[0,2] = 0
00167 
00168             if np.any(p<0) or np.any(p>=self.grid_shape.T):
00169                 continue
00170 
00171             tup = tuple(p.A1)
00172             self.grid_points_list[
00173                     tup[0] + self.grid_shape[0,0] * 
00174                     tup[1] + self.grid_shape[0,0] * 
00175                     self.grid_shape[1,0] * 
00176                     tup[2]].append(pts[:,i])
00177             self.grid[tuple(p.A1)] += 1
00178 
00179 
00180     def to_binary(self,thresh=1):
00181         ''' all cells with occupancy>=thresh set to 1, others set to 0.
00182         '''
00183         filled = (self.grid>=thresh)
00184         self.grid[np.where(filled==True)] = 1
00185         self.grid[np.where(filled==False)] = 0
00186 
00187     def argmax_z(self,index_min=-np.Inf,index_max=np.Inf,search_up=False,search_down=False):
00188         ''' searches in the z direction for maximum number of cells with occupancy==1
00189             call this function after calling to_binary()
00190             returns index.
00191         '''
00192         index_min = int(max(index_min,0))
00193         index_max = int(min(index_max,self.grid_shape[2,0]-1))
00194         z_count_mat = []
00195         
00196         for i in xrange(index_min,index_max+1):
00197             z_count_mat.append(np.where(self.grid[:,:,i]==1)[0].shape[0])
00198 
00199         if z_count_mat == []:
00200             return None
00201         z_count_mat = np.matrix(z_count_mat).T
00202         max_z = np.argmax(z_count_mat)
00203         max_count = z_count_mat[max_z,0]
00204         max_z += index_min
00205         print '#### max_count:', max_count
00206 
00207         if search_up:
00208             max_z_temp = max_z
00209             for i in range(1,5):
00210                 
00211                 
00212                 if (max_z+i)>index_max:
00213                     break
00214                 if (z_count_mat[max_z+i-index_min,0]*5.0)>max_count: 
00215                     max_z_temp = max_z+i
00216             max_z = max_z_temp
00217 
00218         if search_down:
00219             max_z_temp = max_z
00220             for i in range(1,5):
00221                 if (max_z-i)<index_min:
00222                     break
00223                 if (max_z-i)>index_max:
00224                     continue
00225 
00226                 if (z_count_mat[max_z-i-index_min,0]*5.0)>max_count:
00227                     max_z_temp = max_z-i
00228             max_z = max_z_temp
00229 
00230         return max_z,max_count
00231 
00232     def find_plane_indices(self,hmin=-np.Inf,hmax=np.Inf,assume_plane=False):
00233         ''' assume_plane - always return something.
00234             returns list of indices (z) corrresponding to horizontal plane points.
00235             returns [] if there is no plane
00236         '''
00237         index_min = int(max(round((hmin-self.brf[2,0])/self.resolution[2,0]),0))
00238         index_max = int(min(round((hmax-self.brf[2,0])/self.resolution[2,0]),self.grid_shape[2,0]-1))
00239         z_plane,max_count = self.argmax_z(index_min,index_max,search_up=True)
00240         if z_plane == None:
00241             print 'oink oink.'
00242             return []
00243 
00244 
00245 
00246 
00247 
00248 
00249 
00250 
00251         extra_remove_meters = 0.005
00252         n_more_to_remove = int(round(extra_remove_meters/self.resolution[2,0]))
00253         l = range(max(z_plane-10,0),
00254                   min(z_plane+n_more_to_remove+1,self.grid_shape[2,0]-1))
00255         
00256 
00257         if assume_plane == False:
00258             n_more = int(round(0.1/self.resolution[2,0]))
00259             l_confirm = l+ range(max(l),min(z_plane+n_more+1,self.grid_shape[2,0]-1))
00260 
00261             grid_2d = np.max(self.grid[:,:,l],2)
00262             n_plane_cells = grid_2d.sum()
00263             grid_2d = ni.binary_fill_holes(grid_2d) 
00264 
00265             n_plane_cells = grid_2d.sum()
00266             min_plane_pts_threshold = (self.grid_shape[0,0]*self.grid_shape[1,0])/4
00267             print '###n_plane_cells:', n_plane_cells
00268             print 'min_plane_pts_threshold:', min_plane_pts_threshold
00269             print 'find_plane_indices grid shape:',self.grid_shape.T
00270 
00271             if  n_plane_cells < min_plane_pts_threshold:
00272                 print 'occupancy_grid_3d.find_plane_indices: There is no plane.'
00273                 print 'n_plane_cells:', n_plane_cells
00274                 print 'min_plane_pts_threshold:', min_plane_pts_threshold
00275                 l = []
00276 
00277         return l
00278 
00279 
00280     
00281     
00282     
00283     def grid_to_centroids(self,occupancy_threshold=1):
00284         p = np.matrix(np.row_stack(np.where(self.grid>=occupancy_threshold))).astype('float')
00285         p[0,:] = p[0,:]*self.resolution[0,0]
00286         p[1,:] = p[1,:]*self.resolution[1,0]
00287         p[2,:] = p[2,:]*self.resolution[2,0]
00288         p += self.brf
00289         return p
00290 
00291 
00292     def grid_to_points(self,array=None,occupancy_threshold=1):
00293         ''' array - if not None then this will be used instead of self.grid
00294             returns 3xN matrix of 3d coord of the cells which have occupancy >= occupancy_threshold
00295         '''
00296 
00297         if array == None:
00298             array = self.grid
00299 
00300         idxs = np.where(array>=occupancy_threshold)
00301         list_idxs = (idxs[0]+idxs[1]*self.grid_shape[0,0]+idxs[2]*self.grid_shape[0,0]*self.grid_shape[1,0]).tolist()
00302 
00303         l = []
00304         for pts_idxs in list_idxs:
00305             l += self.grid_points_list[pts_idxs]
00306 
00307         if l == []:
00308             p = np.matrix([])
00309         else:
00310             p = np.column_stack(l)
00311 
00312         return p
00313 
00314     def labeled_array_to_points(self,array,label):
00315         ''' returns coordinates of centers of grid cells corresponding to
00316             label as a 3xN matrix.
00317         '''
00318         idxs = np.where(array==label)
00319         list_idxs = (idxs[0]+idxs[1]*self.grid_shape[0,0]+idxs[2]*self.grid_shape[0,0]*self.grid_shape[1,0]).tolist()
00320 
00321         l = []
00322         for pts_idxs in list_idxs:
00323             l += self.grid_points_list[pts_idxs]
00324 
00325         if l == []:
00326             p = np.matrix([])
00327         else:
00328             p = np.column_stack(l)
00329 
00330         return p
00331 
00332     def remove_vertical_plane(self):
00333         ''' removes plane parallel to the YZ plane.
00334             changes grid.
00335             returns plane_indices, slice corresponding to the vertical plane.
00336             points behind the plane are lost for ever!
00337         '''
00338         self.grid = self.grid.swapaxes(2,0)
00339         self.grid_shape = np.matrix(self.grid.shape).T
00340 
00341 
00342 
00343 
00344         z_max_first,max_count = self.argmax_z(search_down=False)
00345         z_max_second,max_count_second = self.argmax_z(index_min=z_max_first+int(round(0.035/self.resolution[0,0])) ,search_down=False)
00346         z_max_first,max_count = self.argmax_z(search_down=False)
00347 
00348         
00349         if (max_count_second*1./max_count) > 0.3:
00350             z_max = z_max_second
00351         else:
00352             z_max = z_max_first
00353         print 'z_max_first', z_max_first
00354         print 'z_max_second', z_max_second
00355         print 'z_max', z_max
00356 
00357         more = int(round(0.03/self.resolution[0,0]))
00358         plane_indices = range(max(0,z_max-more),min(z_max+more,self.grid_shape[2,0]))
00359 
00360         self.grid = self.grid.swapaxes(2,0)
00361         self.grid_shape = np.matrix(self.grid.shape).T
00362 
00363         ver_plane_slice = self.grid[plane_indices,:,:]
00364         self.grid[plane_indices,:,:] = 0
00365         max_x = max(plane_indices)
00366         behind_indices = range(max_x,self.grid_shape[0,0])
00367         self.grid[behind_indices,:,:] = 0
00368         return plane_indices,ver_plane_slice
00369 
00370     def remove_horizontal_plane(self, remove_below=True,hmin=-np.Inf,hmax=np.Inf,
00371                                 extra_layers=0):
00372         ''' call after to_binary()
00373             removes points corresponding to the horizontal plane from the grid.
00374             remove_below - remove points below the plane also.
00375             hmin,hmax - min and max possible height of the plane. (meters)
00376             This function changes grid.
00377 
00378             extra_layers - number of layers above the plane to remove. Sometimes
00379                            I want to be over zealous while removing plane points.
00380                            e.g. max_fwd_without_collision
00381 
00382             it returns the slice which has been set to zero, in case you want to
00383             leave the grid unchanged.
00384         '''
00385         l = self.find_plane_indices(hmin,hmax)
00386 
00387         if l == []:
00388             print 'occupancy_grid_3d.remove_horizontal_plane: No plane found.'
00389             return None,l
00390 
00391         add_num = min(10,self.grid_shape[2,0]-max(l)-1)
00392         max_l = max(l)+add_num
00393         l_edge = l+range(max(l),max_l+1)
00394 
00395         grid_2d = np.max(self.grid[:,:,l_edge],2)
00396 
00397         grid_2d = ni.binary_fill_holes(grid_2d) 
00398 
00399         connect_structure = np.empty((3,3),dtype='int')
00400         connect_structure[:,:] = 1
00401         eroded_2d = ni.binary_erosion(grid_2d,connect_structure,iterations=2)
00402         grid_2d = grid_2d-eroded_2d
00403         idxs = np.where(grid_2d!=0)
00404 
00405         if max_l>max(l):
00406             for i in range(min(5,add_num)):
00407                 self.grid[idxs[0],idxs[1],max(l)+i+1] = 0
00408 
00409         if remove_below:
00410             l = range(0,min(l)+1)+l
00411 
00412         max_z = max(l)
00413         for i in range(extra_layers):
00414             l.append(max_z+i+1)
00415 
00416         l_edge = l+range(max(l),max_l+1)
00417 
00418         plane_and_below_pts = self.grid[:,:,l_edge]
00419         self.grid[:,:,l] = 0 
00420         return plane_and_below_pts,l_edge
00421 
00422     def segment_objects(self, twod=False):
00423         ''' segments out objects after removing the plane.
00424             call after calling to_binary.
00425             returns labelled_array,n_labels
00426             labelled_array - same dimen as occupancy grid, each object has a different label.
00427         '''
00428         plane_and_below_pts,l = self.remove_horizontal_plane(extra_layers=0)
00429         if l == []:
00430             print 'occupancy_grid_3d.segment_objects: There is no plane.'
00431             return None,None
00432 
00433         if twod == False:
00434             labelled_arr,n_labels = self.find_objects()
00435         else:
00436             labelled_arr,n_labels = self.find_objects_2d()
00437         self.grid[:,:,l] = plane_and_below_pts
00438         return labelled_arr,n_labels
00439 
00440     def find_objects_2d(self):
00441         ''' projects all points into the xy plane and then performs
00442             segmentation by region growing.
00443         '''
00444         connect_structure = np.empty((3,3),dtype='int')
00445         connect_structure[:,:] = 1
00446         grid_2d = np.max(self.grid[:,:,:],2)
00447 
00448 
00449         labeled_arr,n_labels = ni.label(grid_2d,connect_structure)
00450         print 'found %d objects'%(n_labels)
00451 
00452         labeled_arr_3d = self.grid.swapaxes(2,0)
00453         labeled_arr_3d = labeled_arr_3d.swapaxes(1,2)
00454         print 'labeled_arr.shape:',labeled_arr.shape
00455         print 'labeled_arr_3d.shape:',labeled_arr_3d.shape
00456         labeled_arr_3d = labeled_arr_3d*labeled_arr
00457         labeled_arr_3d = labeled_arr_3d.swapaxes(2,0)
00458         labeled_arr_3d = labeled_arr_3d.swapaxes(1,0)
00459         labeled_arr = labeled_arr_3d
00460         
00461 
00462         if n_labels > 0:
00463             labels_list = range(1,n_labels+1)
00464             
00465             count_objects = ni.sum(self.grid,labeled_arr,labels_list)
00466             if n_labels == 1:
00467                 count_objects = [count_objects]
00468 
00469             t0 = time.time()
00470             new_labels_list = []
00471 
00472             for c,l in zip(count_objects,labels_list):
00473                 if c > 3:
00474                     new_labels_list.append(l)
00475                 else:
00476                     labeled_arr[np.where(labeled_arr == l)] = 0
00477 
00478             
00479             for nl,l in enumerate(new_labels_list):
00480                 labeled_arr[np.where(labeled_arr == l)] = nl+1
00481 
00482             n_labels = len(new_labels_list)
00483             t1 = time.time()
00484             print 'time:', t1-t0
00485 
00486         print 'found %d objects'%(n_labels)
00487 
00488         return labeled_arr_3d,n_labels
00489 
00490     def find_objects(self):
00491         ''' region growing kind of thing for segmentation. Useful if plane has been removed.
00492         '''
00493         connect_structure = np.empty((3,3,3),dtype='int')
00494         grid = copy.copy(self.grid)
00495 
00496         connect_structure[:,:,:] = 0
00497 
00498         connect_structure[1,1,:] = 1
00499         iterations = int(round(0.005/self.resolution[2,0]))
00500 
00501 
00502 
00503         
00504 
00505         connect_structure[:,:,:] = 1
00506         labeled_arr,n_labels = ni.label(grid,connect_structure)
00507         print 'ho!'
00508         print 'found %d objects'%(n_labels)
00509 
00510         if n_labels == 0:
00511             return labeled_arr,n_labels
00512 
00513         labels_list = range(1,n_labels+1)
00514         count_objects = ni.sum(grid,labeled_arr,labels_list)
00515         if n_labels == 1:
00516             count_objects = [count_objects]
00517 
00518 
00519 
00520 
00521 
00522 
00523 
00524 
00525 
00526         t0 = time.time()
00527         new_labels_list = []
00528 
00529         for c,l in zip(count_objects,labels_list):
00530             if c > 3:
00531                 new_labels_list.append(l)
00532             else:
00533                 labeled_arr[np.where(labeled_arr == l)] = 0
00534 
00535         
00536         for nl,l in enumerate(new_labels_list):
00537             labeled_arr[np.where(labeled_arr == l)] = nl+1
00538 
00539         n_labels = len(new_labels_list)
00540         t1 = time.time()
00541         print 'time:', t1-t0
00542 
00543 
00544         print 'found %d objects'%(n_labels)
00545         return labeled_arr,n_labels
00546 
00547 
00548 
00549 if __name__ == '__main__':
00550     import pygame_opengl_3d_display as po3d
00551     import hokuyo.pygame_utils as pu
00552     import processing_3d as p3d
00553 
00554     p = optparse.OptionParser()
00555 
00556     p.add_option('-f', action='store', type='string', dest='pkl_file_name',
00557                  help='file.pkl File with the scan,pos dict.',default=None)
00558     p.add_option('-c', action='store', type='string', dest='pts_pkl',
00559                  help='pkl file with 3D points',default=None)
00560 
00561     opt, args = p.parse_args()
00562     pts_pkl = opt.pts_pkl
00563     pkl_file_name = opt.pkl_file_name
00564 
00565 
00566     
00567 
00568 
00569 
00570 
00571 
00572 
00573     resolution = np.matrix([0.01,0.01,0.01]).T
00574     gr = occupancy_grid_3d(np.matrix([0.45,-0.5,-1.0]).T, np.matrix([0.65,0.05,-0.2]).T,
00575                            resolution)
00576 
00577     if pts_pkl != None:
00578         pts = ut.load_pickle(pts_pkl)
00579 
00580     elif pkl_file_name != None:
00581         dict = ut.load_pickle(pkl_file_name)
00582         pos_list = dict['pos_list']
00583         scan_list = dict['scan_list']
00584         min_angle = math.radians(-40)
00585         max_angle = math.radians(40)
00586 
00587         l1 = dict['l1']
00588         l2 = dict['l2']
00589         pts = p3d.generate_pointcloud(pos_list, scan_list, min_angle, max_angle, l1, l2)
00590 
00591     else:
00592         print 'specify a pkl file -c or -f'
00593         print 'Exiting...'
00594         sys.exit()
00595 
00596 
00597     print 'started filling the grid'
00598     t0 = time.time()
00599     gr.fill_grid(pts)
00600     t1 = time.time()
00601     print 'time to fill the grid:', t1-t0
00602 
00603     
00604     grid_pts = gr.grid_to_centroids()
00605 
00606     cloud = pu.CubeCloud(grid_pts,(0,0,0),(resolution/2).A1.tolist())
00607     pc = pu.PointCloud(pts,(100,100,100))
00608     lc = pu.LineCloud(gr.grid_lines(),(100,100,0))
00609     po3d.run([cloud,pc,lc])
00610 
00611 
00612 
00613