diff --git a/.gitignore b/.gitignore index 4df31e5..97d70b4 100644 --- a/.gitignore +++ b/.gitignore @@ -159,4 +159,5 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +outputs/ diff --git a/examples/pose_detection.py b/examples/pose_detection.py new file mode 100644 index 0000000..a2b485a --- /dev/null +++ b/examples/pose_detection.py @@ -0,0 +1,33 @@ +"""Scan the scene using only the gripper camera, segment the scene with an +image scene segmentation network, project the segmentations into a 3d +pointcloud, and visualize the segmented pointcloud.""" + +from spot_utils.controllers.startup import DEFAULT_SPOT_IP, RobotClient +from spot_utils.perception.capture import capture_rgbd, save_rgbd +import bosdyn +from bosdyn.client.image import ImageClient +import os +import time + +if __name__ == "__main__": + + cameras = ["frontleft", "frontright"] + print("setting up robot") + bosdyn.client.util.setup_logging(False) + sdk = bosdyn.client.create_standard_sdk("RobotClient") + robot = sdk.create_robot(DEFAULT_SPOT_IP) + + bosdyn.client.util.authenticate(robot) + robot.time_sync.wait_for_sync() + + save_dir = "./outputs_ood" + os.makedirs(save_dir, exist_ok=True) + + while(True): + input("Capture image?") + time_str = str(time.time()) + for camera_name in cameras: + image_client = robot.ensure_client(ImageClient.default_service_name) + robot_client = RobotClient(robot=robot, sdk=sdk, image_client=image_client) + rgbd = capture_rgbd(robot_client, camera=camera_name) + save_rgbd(rgbd=rgbd, save_name=time_str, save_dir=os.path.join(save_dir, camera_name)) \ No newline at end of file diff --git a/spot_utils/controllers/startup.py b/spot_utils/controllers/startup.py index 603db0c..8e390b5 100644 --- a/spot_utils/controllers/startup.py +++ b/spot_utils/controllers/startup.py @@ -8,7 +8,6 @@ from bosdyn.client.manipulation_api_client import ManipulationApiClient from bosdyn.client.robot_command import RobotCommandClient from bosdyn.client.robot_state import RobotStateClient - from spot_utils.structures.robot import RobotClient DEFAULT_SPOT_IP = "192.168.80.3" diff --git a/spot_utils/perception/capture.py b/spot_utils/perception/capture.py index b008baa..a572eef 100644 --- a/spot_utils/perception/capture.py +++ b/spot_utils/perception/capture.py @@ -11,6 +11,10 @@ from spot_utils.structures.image import DepthImage, Intrinsics, RGBDImage, RGBImage from spot_utils.structures.robot import RobotClient +from PIL import Image as pImage +import os +import pickle +import time CAMERAS = ["frontleft", "frontright", "left", "right", "back", "hand"] @@ -113,3 +117,33 @@ def capture_rgbd( ) return rgbd + +def save_rgbd(rgbd:RGBDImage, save_name:str, save_dir:str): + + + rgb_path = os.path.join(save_dir, "rgb", f"{time_str}.png") + rgb_im = pImage.fromarray(rgbd.rgb) + os.makedirs(os.path.dirname(rgb_path), exist_ok = True) + rgb_im.save(rgb_path) + + depth_path = os.path.join(save_dir, "depth", f"{time_str}.png") + depth_im = pImage.fromarray(rgbd.depth) + os.makedirs( os.path.dirname(depth_path) , exist_ok = True) + depth_im.save(depth_path) + + other_data = { + "frame": rgbd.frame.to_matrix(), + "depth_scale": rgbd.depth_scale, + "intrinsics": rgbd.intrinsics + } + + data_path = os.path.join(save_dir, "data", f"{time_str}.pkl") + os.makedirs(os.path.dirname(data_path), exist_ok = True) + + with open(data_path, "wb") as f: + pickle.dump(other_data, f) + + + + + diff --git a/spot_utils/structures/robot.py b/spot_utils/structures/robot.py index c5e2341..9fbd960 100644 --- a/spot_utils/structures/robot.py +++ b/spot_utils/structures/robot.py @@ -18,11 +18,11 @@ class RobotClient: robot: Robot sdk: Sdk - state_client: RobotStateClient - command_client: RobotCommandClient - image_client: ImageClient - manipulation_client: ManipulationApiClient - lease_keepalive: LeaseKeepAlive + state_client: RobotStateClient = None + command_client: RobotCommandClient = None + image_client: ImageClient = None + manipulation_client: ManipulationApiClient = None + lease_keepalive: LeaseKeepAlive = None @dataclass