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 o3dclass Camera:def __init__(self, width=1280, height=720,fps=15):self.im_width = widthself.im_height = heightself.fps = fpsself.intrinsic = Noneself.scale = None# 连接相机# self.connect()def connect(self):"""用于连接相机"""self.pipeline = Pipeline()config = Config()# color configprofile_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 configprofile_list = self.pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR)assert profile_list is not Nonedepth_profile = profile_list.get_default_video_stream_profile()assert depth_profile is not Noneprint("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 imgconfig.set_align_mode(OBAlignMode.SW_MODE)self.pipeline.enable_frame_sync()# start configself.pipeline.start(config)# get intrinsicself.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:continuecolor_frame = frames.get_color_frame()if color_frame is None:continuedepth_frame = frames.get_depth_frame()if depth_frame is None:continueif color_frame != None and depth_frame != None:breakreturn color_frame, depth_framedef 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) * scaleprint('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_datadef 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_datadef get_intrinsic(self):"""获取内参"""# get intrinsicitsc = self.pipeline.get_camera_param()raw_intrinsic = itsc.depth_intrinsicintrinsic = 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 testdef 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 imagedepth_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来进行过滤