diff --git a/gym_pybullet_drones/envs/BaseAviary.py b/gym_pybullet_drones/envs/BaseAviary.py index 4f4281250..8a40a3132 100755 --- a/gym_pybullet_drones/envs/BaseAviary.py +++ b/gym_pybullet_drones/envs/BaseAviary.py @@ -303,6 +303,15 @@ def step(self, # seg = ((seg-np.min(seg)) * 255 / (np.max(seg)-np.min(seg))).astype('uint8') # (Image.fromarray(np.reshape(seg, (h, w)))).save(self.IMG_PATH+"frame_"+str(self.FRAME_NUM)+".png") self.FRAME_NUM += 1 + if self.VISION_ATTR: + for i in range(self.NUM_DRONES): + self.rgb[i], self.dep[i], self.seg[i] = self._getDroneImages(i) + #### Printing observation to PNG frames example ############ + self._exportImage(img_type=ImageType.RGB, # ImageType.BW, ImageType.DEP, ImageType.SEG + img_input=self.rgb[i], + path=self.ONBOARD_IMG_PATH+"drone_"+str(i), + frame_num=int(self.step_counter/self.IMG_CAPTURE_FREQ) + ) #### Read the GUI's input parameters ####################### if self.GUI and self.USER_DEBUG: current_input_switch = p.readUserDebugParameter(self.INPUT_SWITCH, physicsClientId=self.CLIENT)