scan.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 
00003 import sys
00004 import os
00005 import fnmatch
00006 from tvtk.api import tvtk
00007 from numpy.random import randint, rand, permutation
00008 import numpy as np
00009 import time
00010 
00011 from optparse import OptionParser
00012 
00013 num_objects_in_scene = 8
00014 model_center_square = 1.5
00015 
00016 camera_height = 1.5
00017 camera_height_delta = 0.5
00018 
00019 camera_rotation_delta = 5
00020 
00021 baseline = 0.075
00022 focal_length = 585.0
00023 
00024 width = 640
00025 height = 480
00026 
00027 
00028 pcd_header = """
00029 # .PCD v0.7 - Point Cloud Data file format
00030 VERSION 0.7
00031 FIELDS x y z
00032 SIZE 4 4 4
00033 TYPE F F F
00034 COUNT 1 1 1
00035 WIDTH %d
00036 HEIGHT 1
00037 VIEWPOINT %f %f %f 0 0 0 1
00038 POINTS %d
00039 DATA ascii
00040 """
00041 
00042 y_grid, x_grid = np.mgrid[-240:240,-320:320].astype(np.float32)
00043 
00044 def compute_points(depth_image):
00045         points = np.empty((depth_image.shape[0], depth_image.shape[1], 4), dtype=np.float32)
00046         points[:,:,2] = depth_image
00047         points[:,:,1]= y_grid * points[:,:,2] / focal_length
00048         points[:,:,0]= x_grid * points[:,:,2] / focal_length
00049         points[:,:,3]=1
00050         return points
00051 
00052 def get_models_from_dir(models_dir):
00053         classes = os.listdir(models_dir)
00054         models = {}
00055         for c in classes:
00056                 model_files = [ f for f in os.listdir(models_dir + '/' + c) if fnmatch.fnmatch(f,'*.vtk')]
00057                 models[c] = {}
00058                 for m in model_files:
00059                         r = tvtk.PolyDataReader()
00060                         r.file_name = models_dir + '/' + c + '/' + m
00061                         r.update()
00062                         models[c][m] = tvtk.PolyDataMapper(input=r.output)
00063         return models
00064 
00065 
00066 def get_render_window(options):
00067         ren = tvtk.Renderer(background=(0.0, 0.0, 0.0))
00068         rw = tvtk.RenderWindow(size=(width,height))
00069         rw.off_screen_rendering = 0
00070         rw.add_renderer(ren)
00071         rw.multi_samples = 0
00072         ac = ren.active_camera
00073         ac.view_angle = 44.61
00074         
00075         ac.position = [ options.shift, -options.distance , options.height]
00076         ac.view_up = [ 0, 0 , 1]
00077         ac.focal_point = [ options.shift, 0, options.height]
00078         ac.pitch(options.tilt)
00079         
00080         return rw
00081 
00082         
00083 
00084 def save_pointcloud(rw, options, classname, modelname, i, actor):
00085         narr = np.empty((height,width), dtype=np.float32)
00086         rw.get_zbuffer_data(0,0,width-1,height-1,narr)
00087 
00088         narr = np.flipud(narr)
00089 
00090         cr = rw.renderers[0].active_camera.clipping_range
00091         mask = (narr == 1.0)
00092 
00093         # Z buffer stores z values in log scale. Getting z coordinate in world coordinates.
00094         narr = (narr - 0.5) * 2.0
00095         narr = 2*cr[1]*cr[0]/(narr*(cr[1]-cr[0])-(cr[1]+cr[0]))
00096         narr = -narr
00097 
00098         narr += np.random.normal(0, options.noise_std, (480,640))
00099         narr[mask] = 0
00100         
00101         points = compute_points(narr)
00102         points = points[np.nonzero(narr > 0)]
00103 
00104         ac = rw.renderers[0].active_camera
00105         m = ac.view_transform_matrix.to_array()
00106         transform_matrix = np.linalg.inv(np.dot(np.array([[1,0,0,0],[0,-1, 0,0],[0,0,-1,0],[0,0,0,1]]), m))
00107         transform_matrix = np.dot(np.array([[0,1,0,0],[-1, 0, 0,0],[0,0,1,0],[0,0,0,1]]), transform_matrix)
00108         points = np.dot(transform_matrix, points.T)
00109         points = points.T
00110         points[:, 2] -= actor.center[2]
00111         
00112 
00113         filename = options.output_dir + classname + '/' +  modelname[0:-4] + '/'
00114         filename += 'rotation' + str(i*360/options.num_views) + '_distance' + str(options.distance) 
00115         filename += '_tilt' + str(options.tilt) + '_shift' + str(options.shift) + '.pcd'
00116         
00117         f = open(filename,'w')
00118         f.write(pcd_header % (points.shape[0], transform_matrix[0,3], transform_matrix[1,3], transform_matrix[2,3], points.shape[0]))
00119         for i in range(points.shape[0]):
00120                 f.write("%f %f %f\n" % (points[i,0], points[i,1], points[i,2]))
00121         f.close()
00122         
00123 
00124 
00125 if __name__ == '__main__':
00126         parser = OptionParser()
00127         parser.add_option("--input_dir", type="string", dest="input_dir")
00128         parser.add_option("--output_dir", type="string", dest="output_dir")
00129         parser.add_option("--height", type="float", dest="height", default=1.5)
00130         parser.add_option("--noise_std", type="float", dest="noise_std", default=0.0001)
00131         parser.add_option("--distance", type="float", dest="distance", default=4.0)
00132         parser.add_option("--tilt", type="float", dest="tilt", default=-30.0)
00133         parser.add_option("--shift", type="float", dest="shift", default=0)
00134         parser.add_option("--num_views", type="int", dest="num_views", default=6)
00135 
00136         options, args = parser.parse_args()
00137 
00138         
00139         if options.input_dir == None or options.output_dir == None:
00140                 print 'Please specify input and output directories'
00141                 sys.exit(-1)
00142 
00143         models = get_models_from_dir(options.input_dir)
00144         rw = get_render_window(options)
00145 
00146         if not os.path.exists(options.output_dir):
00147                 os.makedirs(options.output_dir)
00148         
00149         for classname, models_map in models.iteritems():
00150                 class_dir = options.output_dir + classname
00151                 if not os.path.exists(class_dir):
00152                         os.makedirs(class_dir)
00153                 for modelname, mapper in models_map.iteritems():
00154                         model_dir = options.output_dir + classname + '/' + modelname[0:-4]
00155                         if not os.path.exists(model_dir):
00156                                 os.makedirs(model_dir)
00157                         actor = tvtk.Actor(mapper=mapper)
00158                         angle = 360.0/options.num_views
00159                         rw.renderers[0].add_actor(actor)
00160                         for i in range(options.num_views):
00161                                 actor.rotate_z(angle)
00162                                 rw.render()
00163                                 save_pointcloud(rw, options, classname, modelname, i, actor)
00164                         rw.renderers[0].remove_actor(actor)
00165                                 
00166         
00167         
00168 
00169  
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Properties Friends Defines


furniture_classification
Author(s): Vladyslav Usenko
autogenerated on Thu May 23 2013 18:32:29