sbx_plot_mesh

# NOTE : run within DATASET folder
# project 2D mask, 3D bounding box, and vertices from the mesh into the dataset

import glob
import json
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import trimesh
from matplotlib.patches import Polygon
from PIL import Image
from pycocotools import mask
from skimage import measure

def project(xyz, K, RT):
  # xyz: [N, 3],  K: [3, 3],  RT: [3, 4]
  xyz = np.dot(xyz, RT[:, :3].T) + RT[:, 3:].T
  xyz = np.dot(xyz, K.T)
  xy = xyz[:, :2] / xyz[:, 2:]
  return xy

coco_json_search_path = "./*COCO*.json"
coco_fnames = glob.glob(coco_json_search_path)
coco_ds = json.load(open(coco_fnames[0]))

cat_id_2_name = {x['id']:x['name'] for x in coco_ds["categories"]}
image_id_2_image_data = {x["id"]:x for x in coco_ds['images']}
scene_id_2_depth_data = {i['scene_id']:i for i in coco_ds['images'] if 'Depth' in i['file_name']}

cat_id_2_mesh = {
  0 : trimesh.load('./meshes/001_pringles_can_C.obj'),
  1 : trimesh.load('./meshes/002_master_chef_can_google_16k.obj'),
  2 : trimesh.load('./meshes/003_cracker_box_google_16k.obj'),
  3 : trimesh.load('./meshes/008_pudding_box_16K_sbx_fix.obj'),
  4 : trimesh.load('./meshes/010_potted_meat_can_google_16k.obj'),
}

# Iterate over all annotations in the dataset
for annot_idx in range(len(coco_ds['annotations'])):
  annot_data = coco_ds['annotations'][annot_idx]

  # ignore annotations where the estimated visibility is less than X%
  if annot_data['visible_perc'] < 0.1:
    continue

  mesh1 = cat_id_2_mesh[annot_data["category_id"]]
  sample_rate = int(max(1,mesh1.vertices.shape[0] / 100.0))
  mesh_vertices = mesh1.vertices[::sample_rate,:]

  img_id = annot_data['image_id']
  img_data = image_id_2_image_data[img_id]
  rgb_img_fn = img_data['file_name']
  scene_id = img_data['scene_id']
  depth_img_fn = scene_id_2_depth_data[scene_id]['file_name']

  K = np.array(img_data['K'])
  RT = np.array(annot_data['RT'])

  bbox_2D = annot_data['bbox']
  bbox_3D = np.array(annot_data['bbox_3D'])
  bbox_xy = project(bbox_3D, K, RT)
  mesh_xy = project(mesh_vertices, K, RT)

  rle = annot_data['segmentation']
  compressed_rle = mask.frPyObjects(rle, rle.get('size')[0], rle.get('size')[1])
  mask_bin = mask.decode(compressed_rle)

  img_pil = Image.open(rgb_img_fn)
  depth_pil = Image.open(depth_img_fn)
  category_name = cat_id_2_name[annot_data["category_id"]]

  # Plot 4 views: 
  #   1, RGB image with the 2D bounding box around the item
  #   2, RGB image with 3D bounding box for the item projected on to the image
  #   3, RGB image with the sampled mesh vertices projected onto the item
  #   4, Depth image with the sampled mesh vertices projected onto the item

  img_width = img_data['width']
  img_height = img_data['height']
  DPI = 100
  fig = plt.figure(dpi=DPI, figsize=(img_width*2.0/float(DPI),img_height*2.0/float(DPI)))

  # 1, plot the rgb image with 2D bounding box
  ax1 = fig.add_subplot(2,2,1)
  ax1.set_title("RGB:" + category_name)
  np_img_rgb1 = np.asarray(img_pil)
  x,y,w,h = annot_data['bbox']
  cv2.rectangle(np_img_rgb1,(int(x),int(y)),(int(x+w),int(y+h)),(0,255,0),2)
  cv2.putText(np_img_rgb1,category_name,(int(x),int(y-10)),0,1.15,(0,255,0))
  ax1.imshow(np_img_rgb1)

  # 2, plot the mask and project 3D bounding box around the item
  ax2 = fig.add_subplot(2,2,2)
  ax2.set_title("RGB:" + "3D bounding box")
  ax2.imshow(img_pil)
  contours = measure.find_contours(mask_bin, 0.5)
  for i in range(len(contours)):
    conts = [np.flip(con, axis=0) for con in contours[i]]
    p2 = Polygon(conts, alpha=0.5, color='red')
    ax2.add_artist(p2)
  for xy2 in bbox_xy:
    ax2.plot(xy2[0],xy2[1], 'bo', markersize=5.0)

  # 3, project the mesh onto the rgb image
  ax3 = fig.add_subplot(2,2,3)
  ax3.set_title("RGB: Mesh")
  np_img_rgb3 = np.asarray(img_pil)
  x,y,w,h = annot_data['bbox']
  cv2.rectangle(np_img_rgb3,(int(x),int(y)),(int(x+w),int(y+h)),(0,255,0),2)
  cv2.putText(np_img_rgb3,category_name,(int(x),int(y-10)),0,1.15,(0,255,0))
  ax3.imshow(np_img_rgb3)
  for xy2 in mesh_xy:
    ax3.plot(xy2[0],xy2[1], 'yo', markersize=1.0)

  # 4, project the mesh onto the depth image
  ax4 = fig.add_subplot(2,2,4)
  ax4.set_title("Depth: Mesh")
  np_img_depth = np.asarray(depth_pil)
  x,y,w,h = annot_data['bbox']
  cv2.rectangle(np_img_depth,(int(x),int(y)),(int(x+w),int(y+h)),(0,255,0),2)
  cv2.putText(np_img_depth,category_name,(int(x),int(y-10)),0,1.15,(0,255,0))
  ax4.imshow(np_img_depth)
  for xy2 in mesh_xy:
    ax4.plot(xy2[0],xy2[1], 'yo', markersize=1.0)

  plt.show()