FPV frame saving

utiasDSL · Aug 2, 2023 · 3f6075f · 3f6075f
1 parent ecd18e3
commit 3f6075f
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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)