microsoft / singleshotpose

This research project implements a real-time object detection and pose estimation method as described in the paper, Tekin et al. "Real-Time Seamless Single Shot 6D Object Pose Prediction", CVPR 2018. (https://arxiv.org/abs/1711.08848).
MIT License
727 stars 214 forks source link

Drawing 3D bounding boxes of all objects in an image #168

Open madhumilan opened 3 years ago

madhumilan commented 3 years ago

Can please someone tell if they are able to draw bounding boxes for all the test objects available in an image? I tried very much, but could not understand how exactly this can be done. The valid_multi.ipynb does not draw poses of multiple objects. Have someone tried this and succeeded? Your help is appreciated. Thanks.

ht2018zzk commented 1 year ago

@madhumilan Hello, I have also encountered the same problem, can not achieve the visualization of multiple goals, have you solved this problem now?

AristideLaignel commented 1 year ago

Hi, In the utils.py you can add this function :

def plot_3dboxes(img, corners, savename=None, class_names=None, color=(0, 204, 0)):

    pt2 = (corners[1][0], corners[1][1])
    pt3 = (corners[2][0], corners[2][1])
    pt4 = (corners[3][0], corners[3][1])
    pt5 = (corners[4][0], corners[4][1])
    pt6 = (corners[5][0], corners[5][1])
    pt7 = (corners[6][0], corners[6][1])
    pt8 = (corners[7][0], corners[7][1])
    pt9 = (corners[8][0], corners[8][1])

    rect1 = np.array([pt2, pt3, pt5, pt4], np.int32)
    rect2 = np.array([pt2, pt3, pt7, pt6], np.int32)
    rect3 = np.array([pt2, pt4, pt8, pt6], np.int32)
    rect4 = np.array([pt4, pt5, pt9, pt8], np.int32)
    rect5 = np.array([pt6, pt7, pt9, pt8], np.int32)

    image = np.asarray(img)      
    cv2.polylines(image, [rect1], True, color, 2)
    cv2.polylines(image, [rect2], True, color, 2)
    cv2.polylines(image, [rect3], True, color, 2)
    cv2.polylines(image, [rect4], True, color, 2)
    cv2.polylines(image, [rect5], True, color, 2)

    res_image = Image.fromarray(image)

    if savename:
        print("save plot results to %s" % savename)
        res_image.save(savename)
    return res_image

And after you can convert your image in your valid.py like that :

image_np = np.transpose(tensor_cpu.numpy(), (1, 2, 0))
image = Image.fromarray((image_np * 255).astype(np.uint8))

if image.mode != 'RGB':
     image= image.convert('RGB')

and plot your image in your loop like :

plot_3dboxes(img,corners2D_gt,"images/valid_image"+str(batch_idx)+".png")