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