class DataGenerator: def __init__(self): self._loadWorld('./main.xml') self._loadSensor('./camera.rob') self.chessEngine = ChessEngine(self.world, self.world.terrain('tabletop')) self.chessEngine.loadPieces() self.chessEngine.loadBoard() self.chessEngine.arrangeBoard(90) self.chessEngine.arrangePieces() self._imageDir = 'image_dataset' self._colorDir = self._imageDir + '/color' self._depthDir = self._imageDir + '/depth' self._rectifiedDir = self._imageDir + '/rectified' self._metaDataFN = self._imageDir + '/metadata.csv' self._dirs = [ self._imageDir, self._colorDir, self._depthDir, self._rectifiedDir ] self._colorFNFormat = self._colorDir + '/%06d_%02d.png' self._depthFNFormat = self._depthDir + '/%06d_%02d.png' self._rectifiedFNFormat = self._rectifiedDir + '/%06d.png' self._rectifiedPictureCorners = np.float32( [[DataUtils.IMAGE_SIZE, 0], [0, 0], [0, DataUtils.IMAGE_SIZE], [DataUtils.IMAGE_SIZE, DataUtils.IMAGE_SIZE]]) self._boardWorldCorners = self.chessEngine.getBoardCorners() def _loadWorld(self, world_fn): world = WorldModel() res = world.readFile(world_fn) if not res: print("Unable to read file", world_fn) exit(0) for i in range(world.numRigidObjects()): obj = world.rigidObject(i) #this will perform a reasonable center of mass / inertia estimate m = obj.getMass() m.estimate(obj.geometry(), mass=0.454, surfaceFraction=0.2) obj.setMass(m) self.world = world def _loadSensor(self, sensor_fn): try: self.world.readFile(sensor_fn) except Exception as e: print(f'Error in loading filename: {e}') robot = self.world.robot(0) if robot is None: print("Robot is none: check sensor filename") exit(1) sensor = self.world.robot(0).sensor(0) if sensor is None: print("Sensor is none: check sensor filename") exit(1) self.robot = robot self.sensor = sensor def _createDatasetDirectory(self): for d in self._dirs: try: os.mkdir(d) except Exception: pass def deleteDataset(self): try: shutil.rmtree(self._imageDir) except OSError as e: print(f'Error deleting {self._imageDir}: {e}') def _randomlyRotateCamera(self, min_r=0, max_r=70, dist=DIST_FROM_BOARD): """ Randomly rotates camera about x-axis and zooms out from the center of the tabletop :param: min_r: the minimum random rotation in degrees :param: max_r: the maximum random rotation in degrees :param: dist: the distance to zoom out from center of the table :return: the angle of rotation sampled """ min_r = math.radians(min_r) max_r = math.radians(max_r) table_center = self.chessEngine.getTableCenter() table_R = so3.from_axis_angle(([1, 0, 0], -math.pi)) table_t = table_center rot_deg = random.uniform(min_r, max_r) zoom_out_R = so3.from_axis_angle(([1, 0, 0], rot_deg)) zoom_out_t = vectorops.mul( [0, math.sin(rot_deg), -math.cos(rot_deg)], dist) xform = se3.mul((table_R, table_t), (zoom_out_R, zoom_out_t)) sensing.set_sensor_xform(self.sensor, xform) return rot_deg def _rotateCamera(self, r, dist=DIST_FROM_BOARD): """ Rotates a camera and zooms out from center of the tabletop :param: r: rotation in degrees :param: dist: distance to zoom out from center of the table :return: angle of rotation """ return self._randomlyRotateCamera(r, r, dist) def generateImages(self, max_pics=100, data_distribution=None): assert max_pics > 0 if data_distribution is None: data_distribution = DataUtils.LIMITED_DISTRIBUTION assert len(data_distribution) == 13 self._createDatasetDirectory() for i in range(self.world.numRigidObjects()): self.world.rigidObject(i).appearance().setSilhouette(0) metadata_f = open(self._metaDataFN, mode='w+') metadata_f.write('color,depth,piece\n') DEPTH_SCALE = 8000 def loop_through_sensors(world=self.world, sensor=self.sensor, max_pics=max_pics): for counter in tqdm(range(max_pics)): if counter > 0: self.chessEngine.randomizePieces(32) # Arrange pieces and model world self.chessEngine.arrangePieces() self._randomlyRotateCamera(20, 40, 0.6) sensor.kinematicReset() sensor.kinematicSimulate(world, 0.01) rgb, depth = sensing.camera_to_images(self.sensor) # Project RGB and depth images to rectify them local_corner_coords = np.float32([ sensing.camera_project(self.sensor, self.robot, pt)[:2] for pt in self._boardWorldCorners ]) H = cv2.getPerspectiveTransform(local_corner_coords, self._rectifiedPictureCorners) color_rectified = cv2.warpPerspective( rgb, H, (DataUtils.IMAGE_SIZE, DataUtils.IMAGE_SIZE)) depth_rectified = cv2.warpPerspective( depth, H, (DataUtils.IMAGE_SIZE, DataUtils.IMAGE_SIZE)) depth_rectified = np.uint8( (depth_rectified - depth_rectified.min()) / (depth_rectified.max() - depth_rectified.min())) pieces_arrangement = self.chessEngine.getPieceArrangement() images, classes = DataUtils.split_image_by_classes( color_rectified, depth_rectified, data_distribution, pieces_arrangement) rectified_fname = self._rectifiedFNFormat % (counter) Image.fromarray(color_rectified).save(rectified_fname) for idx in range(sum(data_distribution)): color_fname = self._colorFNFormat % (counter, idx) depth_fname = self._depthFNFormat % (counter, idx) Image.fromarray(images[idx, :, :, :3]).save(color_fname) Image.fromarray(images[idx, :, :, 3]).save(depth_fname) metadata_f.write( f'{color_fname},{depth_fname},{classes[idx]}\n') vis.show(False) vis.loop(callback=loop_through_sensors) metadata_f.close()