graspnet+Astra2相机实现部署

graspnet+Astra2相机实现部署

🚀 环境配置 🚀

• ubuntu 20.04
• Astra2相机
• cuda 11.0.1
• cudnn v8.9.7
• python 3.8.19
• pytorch 1.7.0
• numpy 1.23.5

1. graspnet的复现

具体的复现流程可以参考这篇文章：Ubuntu20.04下GraspNet复现流程

这里就不再详细介绍了

2. Astra2的Python API

以下内容都是参考官方文档：Orbbec SDK for Python 使用手册

我们首先确认输入到网络中的数据为一个点云数据，再一个我们需要一个rgb图像用来给点云上色，graspnetAPI帮我们写好了从深度图转换到点云的函数create_point_cloud_from_depth_image，所以我们只需要写好从相机的视频流获取深度图片和rgb图片的部分就好了，特别注意的是，大多数相机的rgb的fov是要大于深度图的fov所以我们要对两者进行对齐操作，对齐操作的本质就是在深度图中填充0，使得深度图和rgb图的大小一致。大多数的相机厂商已经提供了具体的示例来演示如何进行对齐，这里就不再赘述。

这里我直接给出我写的astra2.py，用于获取相机的深度图和rgb图的代码，大家可以参考一下思路

astra2.py

from pyorbbecsdk import *
import numpy as np
import cv2
import os
import open3d as o3d

class Camera:
    def __init__(self, width=1280, height=720,fps=15):
        self.im_width = width
        self.im_height = height
        self.fps = fps
        self.intrinsic = None
        self.scale = None

        # 连接相机
        # self.connect()
    
    def connect(self):
        """用于连接相机"""
        self.pipeline = Pipeline()
        config = Config()

        # color config
        profile_list = self.pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
        color_profile = profile_list.get_default_video_stream_profile()
        config.enable_stream(color_profile)
        # depth config
        profile_list = self.pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR)
        assert profile_list is not None
        depth_profile = profile_list.get_default_video_stream_profile()
        assert depth_profile is not None
        print("color profile : {}x{}@{}_{}".format(color_profile.get_width(),
                                                   color_profile.get_height(),
                                                   color_profile.get_fps(),
                                                   color_profile.get_format()))
        print("depth profile : {}x{}@{}_{}".format(depth_profile.get_width(),
                                                   depth_profile.get_height(),
                                                   depth_profile.get_fps(),
                                                   depth_profile.get_format()))
        config.enable_stream(depth_profile)
        # set synchronize for depth img and color img
        config.set_align_mode(OBAlignMode.SW_MODE)
        self.pipeline.enable_frame_sync()
        # start config
        self.pipeline.start(config)

        # get intrinsic
        self.intrinsic = self.get_intrinsic()
    
    def disconnect(self):
        """用于断开相机"""
        self.pipeline.stop()
    

    def get_frame(self):
        """通过流来获取color frame和depth frame"""
        while True:
            frames: FrameSet = self.pipeline.wait_for_frames(200)
            if frames is None:
                continue
            color_frame = frames.get_color_frame()
            if color_frame is None:
                continue
            depth_frame = frames.get_depth_frame()
            if depth_frame is None:
                continue
            if color_frame != None and depth_frame != None:
                break
        
        return color_frame, depth_frame

    def frame2data(self, color_frame, depth_frame):
        """暂时没用"""
        width = depth_frame.get_width()
        height = depth_frame.get_height()
        scale = depth_frame.get_depth_scale()
        depth_data = np.frombuffer(depth_frame.get_data(), dtype=np.uint16)
        depth_data = depth_data.reshape((height, width))

        width = color_frame.get_width()
        height = color_frame.get_height()
        color_data = np.asanyarray(color_frame.get_data(), dtype=np.uint16)
        # color_data = color_data.reshape((height, width, 3))
        return color_data.astype(np.float32), depth_data.astype(np.float32)
    
    def get_data(self):
        """通过流来获取color data和depth data"""

        # 连接相机
        self.connect()
        color_frame, depth_frame = self.get_frame()

        width = color_frame.get_width()
        height = color_frame.get_height()
        color_format = color_frame.get_format()
        data = np.asanyarray(color_frame.get_data())
        color_data = cv2.imdecode(data, cv2.IMREAD_COLOR)
        color_data.astype(np.float32)

        print('color_image.shape: ', color_data.shape)
        # print("===width: {}===".format(width))
        # print("===height: {}===".format(height))
        width = depth_frame.get_width()
        height = depth_frame.get_height()
        scale = depth_frame.get_depth_scale()
        print("===width: {}===".format(width))
        print("===height: {}===".format(height))
        print("===scale: {}===".format(scale))
        save_dir = os.path.join(os.getcwd(), "real/intrinsic")
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        filename = save_dir + "/camera_depth_scale.txt"
        save = np.array([scale])
        np.savetxt(filename, save, delimiter=' ')

        depth_data = np.frombuffer(depth_frame.get_data(), dtype=np.uint16)
        depth_data = depth_data.reshape((height, width))
        depth_data = depth_data.astype(np.float32) * scale
        print('depth_image.shape: ', depth_data.shape)

        depth_data = cv2.normalize(depth_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_16U)
        
        # 断开相机
        self.disconnect()

        return color_data, depth_data

    def get_data_saved(self):
        color_data, depth_data = self.get_data()
        
        # depth_image = depth_data_normalized.astype(np.uint8)

        save_image_dir = os.path.join(os.getcwd(), "real/images")
        if not os.path.exists(save_image_dir):
            os.mkdir(save_image_dir)
        depth_filename = save_image_dir + "/depth_{}x{}.png".format(depth_data.shape[0], depth_data.shape[1])
        color_filename = save_image_dir + "/color_{}x{}.png".format(color_data.shape[0], color_data.shape[1])
        cv2.imwrite(color_filename, color_data)
        # depth_data.tofile(depth_filename)
        cv2.imwrite(depth_filename, depth_data)

        return color_data, depth_data
        

    def get_intrinsic(self):
        """获取内参"""
        # get intrinsic
        itsc = self.pipeline.get_camera_param()
        raw_intrinsic = itsc.depth_intrinsic

        intrinsic = np.array([raw_intrinsic.fx, 0, raw_intrinsic.cx, 
                                0, raw_intrinsic.fy, raw_intrinsic.cy,
                                0, 0, 1]).reshape(3,3)
        print("intrinsic: ", itsc)
        print('depth intrinsic: ', raw_intrinsic)
        print("intrinsic matrix", intrinsic)
        save_dir = os.path.join(os.getcwd(), "real/intrinsic")
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        filename = save_dir + "/camera_itcs.txt"

        np.savetxt(filename, intrinsic, delimiter=' ')

        return intrinsic
    
    # for test
    def visualize(self):
        """显示rgbd图像"""
        color_data, depth_data = self.get_data()
        depth_image = cv2.normalize(depth_data, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        depth_image = cv2.applyColorMap(depth_image, cv2.COLORMAP_JET)
        # overlay color image on depth image
        depth_image = cv2.addWeighted(color_data, 0.5, depth_image, 0.5, 0)
        cv2.imshow("Depth with Color", depth_image)
        cv2.waitKey(500)
    
    

if __name__ == '__main__':
    camera = Camera()
    camera.visualize()
    color, depth = camera.get_data()
    print("depth.shape: ", depth.shape)

测试相机是否实现对齐,结果如下

表明了相机确实实现了对齐操作

3. 修改demo.py

我们使用get_data()函数就能够让相机进行一次拍摄，然后得到color_data和depth_data供我们进行后续的处理。然后可以修改一下demo.py，将从文件读取数据改为直接从相机进行读取

demo.py

...
from astra2 import Camera()
astra2 = Camera()

...

def get_and_process_data():

    # 使用相机获取一次数据
    color, depth = astra2.get_data()
    color = color / 255.0
    ...

然后别的部分可以保持不变或者以后再修改。这里一定要使用官方提供的checkpoint-rs.tar，如果使用checkpoint-kn.tar，会出现异常，暂时我也没有找到原因。

最后处理的效果如下

可以看到出现了许多我们不希望存在的抓取框,这个可以通过调整workspace_mask来进行过滤