allow auto simulate/collect

robotflow-initiative · Apr 1, 2024 · 4561b4a · 4561b4a
1 parent 9a08e13
commit 4561b4a
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 19 deletions.
diff --git a/Test/test_camera_image.py b/Test/test_camera_image.py
@@ -12,14 +12,14 @@
 env.step()
 camera.SetTransform(position=[0, 0.25, 0], rotation=[30, 0, 0])
 camera.LookAt(target=box.data["position"])
-camera.GetDepth(width=1024, height=512, zero_dis=0.2, one_dis=1)
+camera.GetDepth16Bit(width=1024, height=512, zero_dis=0.2, one_dis=1)
 camera.GetDepthEXR(width=1024, height=512)
 camera.GetRGB(width=1024, height=512)
 env.step()
 print(camera.data["depth"])
 print(camera.data["depth_exr"])
 print(camera.data["rgb"])
-image = np.frombuffer(camera.data["depth"], dtype=np.uint8)
+image = np.frombuffer(camera.data["depth"], dtype=np.uint16)
 image = cv2.imdecode(image, cv2.IMREAD_COLOR)
 print(image.shape)
 env.close()

diff --git a/pyrfuniverse/__init__.py b/pyrfuniverse/__init__.py
@@ -1,5 +1,5 @@
 # Version of the library that will be used to upload to pypi
-__version__ = "0.12.4"
+__version__ = "0.13.0"
 
 import os.path
 import json

diff --git a/pyrfuniverse/attributes/camera_attr.py b/pyrfuniverse/attributes/camera_attr.py
@@ -35,20 +35,20 @@ def parse_message(self, data: dict):
             self.data['3d_bounding_box']: The 3d bounding box of objects in world coordinate.
         """
         super().parse_message(data)
-        if "rgb" in self.data:
-            self.data["rgb"] = base64.b64decode(self.data["rgb"])
-        if "normal" in self.data:
-            self.data["normal"] = base64.b64decode(self.data["normal"])
-        if "id_map" in self.data:
-            self.data["id_map"] = base64.b64decode(self.data["id_map"])
-        if "depth" in self.data:
-            self.data["depth"] = base64.b64decode(self.data["depth"])
-        if "depth_exr" in self.data:
-            self.data["depth_exr"] = base64.b64decode(self.data["depth_exr"])
-        if "amodal_mask" in self.data:
-            self.data["amodal_mask"] = base64.b64decode(self.data["amodal_mask"])
-        if "heat_map" in self.data:
-            self.data["heat_map"] = base64.b64decode(self.data["heat_map"])
+        # if "rgb" in self.data:
+        #     self.data["rgb"] = base64.b64decode(self.data["rgb"])
+        # if "normal" in self.data:
+        #     self.data["normal"] = base64.b64decode(self.data["normal"])
+        # if "id_map" in self.data:
+        #     self.data["id_map"] = base64.b64decode(self.data["id_map"])
+        # if "depth" in self.data:
+        #     self.data["depth"] = base64.b64decode(self.data["depth"])
+        # if "depth_exr" in self.data:
+        #     self.data["depth_exr"] = base64.b64decode(self.data["depth_exr"])
+        # if "amodal_mask" in self.data:
+        #     self.data["amodal_mask"] = base64.b64decode(self.data["amodal_mask"])
+        # if "heat_map" in self.data:
+        #     self.data["heat_map"] = base64.b64decode(self.data["heat_map"])
 
     def AlignView(self):
         """
@@ -145,8 +145,8 @@ def GetID(
 
     def GetDepth(
         self,
-        zero_dis: float,
-        one_dis: float,
+        zero_dis: float = 0.,
+        one_dis: float = 1.,
         width: int = None,
         height: int = None,
         fov: float = 60.0,
@@ -185,6 +185,8 @@ def GetDepth(
 
     def GetDepth16Bit(
         self,
+        zero_dis: float = 0.,
+        one_dis: float = 1.,
         width: int = None,
         height: int = None,
         fov: float = 60.0,
@@ -211,6 +213,8 @@ def GetDepth16Bit(
                 height = int(intrinsic_matrix[1, 2] * 2)
         self._send_data(
             "GetDepth16Bit",
+            float(zero_dis),
+            float(one_dis),
             intrinsic_matrix,
             int(width),
             int(height),

diff --git a/pyrfuniverse/envs/base_env.py b/pyrfuniverse/envs/base_env.py
@@ -234,10 +234,32 @@ def close(self):
         if self.communicator is not None:
             self.communicator.close()
 
+    def SetAutoSimulate(self, auto: bool):
+        """
+        Enable or Disable Auto Simulate
+        """
+        self._send_env_data("SetAutoSimulate", auto)
+
+    def SetAutoCollect(self, auto: bool):
+        """
+        Enable or Disable Auto Collect
+        """
+        self._send_env_data("SetAutoCollect", auto)
+
     def Simulate(self, time_step: float = -1, count: int = 1):
+        """
+        Physics simulation
+
+        Args:
+            time_step: delta time of simulation pre step
+            count：count of simulation
+        """
         self._send_env_data("Simulate", float(time_step), int(count))
 
     def Collect(self):
+        """
+        Collect env data
+        """
         self._send_env_data("Collect")
 
     def GetAttr(self, id: int):