下载GRCNN项目:https://github.com/skumra/robotic-grasping.git
导出onnx模型:
import torch
net = torch.load("trained-models/jacquard-rgbd-grconvnet3-drop0-ch32/epoch_42_iou_0.93")
x = torch.rand(1, 4, 300, 300).cuda()
torch.onnx.export(net, x, "./grcnn.onnx", opset_version = 13)
onnx模型结构如下:
onnxruntime推理
import cv2
import onnxruntime
import numpy as np
from skimage.feature import peak_local_max
def process_data(rgb, depth, width, height, output_size):
left = (width - output_size) // 2
top = (height - output_size) // 2
right = (width + output_size) // 2
bottom = (height + output_size) // 2
depth_img = depth[top:bottom, left:right]
depth_img = np.clip((depth_img - depth_img.mean()), -1, 1)
depth_img = depth_img.transpose(2, 0, 1)
rgb_img = rgb[top:bottom, left:right]
rgb_img = rgb_img.astype(np.float32) / 255.0
rgb_img -= rgb_img.mean()
rgb_img = rgb_img.transpose(2, 0, 1)
ret = np.concatenate((np.expand_dims(depth_img, 0), np.expand_dims(rgb_img, 0)), axis=1)
return np.concatenate((np.expand_dims(depth_img, 0), np.expand_dims(rgb_img, 0)), axis=1)
if __name__ == '__main__':
rgb = cv2.imread(