小物体位姿估计

使用zebrapose

推理代码：

from zebrapose import zebra
from yolo import yolo
import cv2
import numpy as np
import pyrealsense2 as rs
import time
import json
import time
obj_dict = {
    'in':(143, 54)#(model id, yolo id)
    }

#symmetry_obj = ['gear']
symmetry_obj = [' ']
###
obj_name = 'in'

#K = [616.58,0,323.103, 0,616.778,238.464, 0.0, 0.0, 1.0]
K = [924.871,0,644.655, 0,925.167,357.695, 0.0, 0.0, 1.0]
###

def prj_marker2img(point,R,T,K):
    point_cam = np.matmul(R, point) + T
    point_pix = np.matmul(K, point_cam)
    point_pix = point_pix/point_pix[2]
    return point_pix[:2].astype(int)

def load_dims(model_id):
    with open('./assets/models_info.json', 'r') as f:
        model_info = json.load(f)

    info = model_info[str(model_id)]
    dims = [info['size_x'], info['size_y'], info['size_z']]
    return dims

def get_bbox_masks_from_yolo(results, target_id):
    bboxes = []
    masks = []
    for result in results:
        if result[0].item() == target_id:
            bboxes.append(result[1])
            masks.append(result[2])
    return bboxes, masks

if obj_name in symmetry_obj:
    symmetry = True
    point_list = []
    num_point = 12
    for i in range(num_point):
        x = np.cos(i*2*np.pi/num_point)
        y = np.sin(i*2*np.pi/num_point)
        point_list.append((x,y))
else:
    symmetry = False
    
dims = load_dims(obj_dict[obj_name][0])
zebra = zebra(obj_name, obj_dict[obj_name][0], K, 17)
yolo = yolo()
target_yolo_id = obj_dict[obj_name][1]

#color_image = cv2.imread('/home/fuzhao/poseEstimate/zebraPose_inference_gear2/data/rgb5/1702881914.jpg')
color_image = cv2.imread('/home/fuzhao/poseEstimate/zebraPose_inference_gear2/data/rgb5/32.jpg')

t0 = time.time()
yolo_results = yolo.predict(color_image, conf=0.5)

bboxes, masks = get_bbox_masks_from_yolo(yolo_results, target_yolo_id)
print("bboxes.shape",len(bboxes))
#bboxes, masks = bboxes[40:], masks[40:]
#bboxes, masks =[bboxes[1]], [masks[1]]

##for i in range(len(bboxes0)):
##    center_x, center_y, w, h = bboxes0[i]
##    if center_x > 188 and center_x < 215 and center_y >330 and center_y < 360:
##        bboxes = [bboxes0[i]]
##        masks = [masks0[i]]
##        break

##bbox = [(188.0+215.0)/2,(330.0+360.0)/2,27.0,30.0]
##fake_mask = cv2.imread('./images/masks/000019_000019.png',0)
##fake_mask = fake_mask.astype(np.float32)/255.0
##bboxes, masks = [bbox], [fake_mask]

#cv2.imwrite('yolo_mask.jpg', (masks[0]*255).astype(np.uint8))
##for bbox in bboxes:
##    bbox = bbox.astype(int)
##    x,y,w,h = bbox
##    cv2.rectangle(color_image, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)),(255,0,0),1)
##cv2.imshow('roi', color_image)
##cv2.waitKey(0)

r_list, t_list = zebra.predict(color_image, bboxes, masks, debug=True)

print("total infer time:",time.time() - t0 )
print("len(r_list):",len(r_list))
for i in range(len(r_list)):
    r_predict = r_list[i]
    t_predict = t_list[i]
    
    rot_mat = r_predict
    point_center = t_predict[:,0]
    K = zebra.cam_K

    if symmetry:
        upper_circle_point = []
        lower_circle_point = []
        radius = max(dims[0], dims[1])/2.0
        for point in point_list:
            point_object = np.array([point[0]*radius, point[1]*radius, dims[2]/2.0])
            point_img = prj_marker2img(point_object, rot_mat, point_center, K)
            upper_circle_point.append(point_img)

        for point in point_list:
            point_object = np.array([point[0]*radius, point[1]*radius, -dims[2]/2.0])
            point_img = prj_marker2img(point_object, rot_mat, point_center, K)
            lower_circle_point.append(point_img)              

        for i in range(num_point):
            cv2.line(color_image, upper_circle_point[i-1], upper_circle_point[i], (255,0,0),1)
            cv2.line(color_image, lower_circle_point[i-1], lower_circle_point[i], (0,255,0),1)
            #cv2.line(color_image, upper_circle_point[i], lower_circle_point[i], (0,0,255),1)
        
    else:
        point_object = np.array([dims[0]/2, dims[1]/2, dims[2]/2])
        point_img1 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([-dims[0]/2, dims[1]/2, dims[2]/2])
        point_img2 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([dims[0]/2, -dims[1]/2, dims[2]/2])
        point_img3 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([-dims[0]/2, -dims[1]/2, dims[2]/2])
        point_img4 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([dims[0]/2, dims[1]/2, -dims[2]/2])
        point_img5 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([-dims[0]/2, dims[1]/2, -dims[2]/2])
        point_img6 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([dims[0]/2, -dims[1]/2, -dims[2]/2])
        point_img7 = prj_marker2img(point_object,rot_mat,point_center,K)

        point_object = np.array([-dims[0]/2, -dims[1]/2, -dims[2]/2])
        point_img8 = prj_marker2img(point_object,rot_mat,point_center,K)  
        
        cv2.line(color_image, point_img5, point_img6, (255,0,0),1)
        cv2.line(color_image, point_img6, point_img8, (255,0,0),1)
        cv2.line(color_image, point_img8, point_img7, (255,0,0),1)
        cv2.line(color_image, point_img5, point_img7, (255,0,0),1)

        cv2.line(color_image, point_img1, point_img5, (0,255,0),1)
        cv2.line(color_image, point_img2, point_img6, (0,255,0),1)
        cv2.line(color_image, point_img3, point_img7, (0,255,255),1)
        cv2.line(color_image, point_img4, point_img8, (0,255,255),1)
        
        cv2.line(color_image, point_img1, point_img2, (0,0,255),1)
        cv2.line(color_image, point_img2, point_img4, (0,0,255),1)
        cv2.line(color_image, point_img4, point_img3, (0,0,255),1)
        cv2.line(color_image, point_img1, point_img3, (0,0,255),1)

cv2.imshow('image', color_image)
cv2.waitKey(0)