SHARE
TWEET

Untitled

a guest Apr 25th, 2019 76 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python
  2.  
  3. import numpy as np
  4. import trimesh
  5.  
  6. import objslampp
  7.  
  8. models = objslampp.datasets.YCBVideoModels()
  9. dataset = objslampp.datasets.YCBVideoDataset('train')
  10. frame = dataset[0]
  11.  
  12. rgb = frame['color']
  13. depth = frame['depth']
  14. label = frame['label']
  15. K = frame['meta']['intrinsic_matrix']
  16. class_ids = frame['meta']['cls_indexes']
  17.  
  18. class_id = class_ids[0]
  19. mask = label == class_id
  20. cad_file = models.get_cad_model(class_id=class_id)
  21. dim = 32
  22. pitch = models.get_bbox_diagonal(cad_file) / dim
  23.  
  24. pcd = objslampp.geometry.pointcloud_from_depth(
  25.     depth, fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2]
  26. )
  27.  
  28. pcd_ins = pcd[mask]
  29. centroid = np.nanmean(pcd_ins, axis=0)
  30. aabb_min = centroid - pitch * dim / 2
  31. aabb_max = aabb_min + pitch * dim
  32.  
  33. camera = trimesh.scene.Camera(
  34.     resolution=(rgb.shape[1], rgb.shape[0]),
  35.     focal=(K[0, 0], K[1, 1]),
  36. )
  37.  
  38. centers = trimesh.voxel.matrix_to_points(
  39.     np.ones((dim, dim, dim)),
  40.     pitch=pitch,
  41.     origin=aabb_min + pitch,
  42. )
  43.  
  44. r, c = objslampp.geometry.project_to_camera(
  45.     centers,
  46.     fx=K[0, 0],
  47.     fy=K[1, 1],
  48.     cx=K[0, 2],
  49.     cy=K[1, 2],
  50.     image_shape=rgb.shape[:2],
  51. )
  52. r, c = r.round().astype(int), c.round().astype(int)
  53. is_target = mask[r, c]
  54. z_depth = depth[r, c]
  55. z_voxel = centers[:, 2]
  56.  
  57. nonnan = ~np.isnan(z_depth)
  58. # empty: 1, target: 2, other: 3, unknown: 4
  59. labels = np.full((centers.shape[0],), 4, dtype=np.int32)
  60. labels[np.greater((z_depth - z_voxel), pitch / 2, where=nonnan)] = 1
  61. is_occupied = np.less(np.abs(z_depth - z_voxel), pitch / 2, where=nonnan)
  62. labels[is_occupied & is_target] = 2
  63. labels[is_occupied & (~is_target)] = 3
  64.  
  65. # -----------------------------------------------------------------------------
  66.  
  67. scene = trimesh.Scene()
  68.  
  69. nonnan = ~np.isnan(pcd).any(axis=2)
  70. geom = trimesh.PointCloud(vertices=pcd[mask & nonnan], colors=rgb[mask & nonnan])
  71. scene.add_geometry(geom)
  72.  
  73. import imgviz
  74. # keep = labels == 4  # unknown
  75. # keep = np.isin(labels, [2, 3, 4])
  76. keep = np.isin(labels, [3])
  77. centers = centers[keep]
  78. labels = labels[keep]
  79.  
  80. colors = imgviz.label2rgb(labels[None])[0]
  81. colors = np.hstack((colors, [[127]] * len(colors)))
  82. boxes = trimesh.voxel.multibox(centers, pitch=pitch, colors=colors)
  83. scene.add_geometry(boxes)
  84.  
  85. geom = trimesh.creation.box(extents=aabb_max - aabb_min)
  86. geom.apply_translation(aabb_min + (aabb_max - aabb_min) / 2)
  87. # geom.visual.face_colors = (0, 1., 0, 0.5)
  88. # scene.add_geometry(geom)
  89. geom = objslampp.extra.trimesh.box_to_wired_box(geom)
  90. scene.add_geometry(geom)
  91.  
  92. scene.camera.resolution = camera.resolution
  93. scene.camera.focal = camera.focal
  94. scene.camera.transform = objslampp.extra.trimesh.camera_transform(np.eye(4))
  95.  
  96. scene.show()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top