open3d visualization

official/vision/beta/projects/mesh_rcnn/demo/visualize_mesh.py

@@ -0,0 +1,82 @@
+"""Mesh Visualization"""
+
+import matplotlib.pyplot as plt
+import tensorflow as tf
+import open3d as o3d
+from mpl_toolkits.mplot3d import art3d
+
+from official.vision.beta.projects.mesh_rcnn.ops.cubify import cubify
+
+
+def create_voxels(grid_dims, batch_size, occupancy_locs):
+    ones = tf.ones(shape=[len(occupancy_locs)])
+    voxels = tf.scatter_nd(
+        indices=tf.convert_to_tensor(occupancy_locs, tf.int32),
+        updates=ones,
+        shape=[batch_size, grid_dims, grid_dims, grid_dims])
+
+    return voxels
+
+
+def visualize_mesh_o3d(verts, faces, verts_mask, faces_mask, smoothing = True):
+
+    """Visualize meshes using Open3D
+
+    Args:
+        verts: A "Tensor" of shape of [V,3], where V is the number of vertices
+        faces: A "Tensor" of shape of [F,3], where each face is composed of 3 integers,
+        which map to the points in the vertex tensor
+        verts_mask: A "Tensor" of shape of [V,1], which consists of 0 or 1 to determine
+        which vertices are valid
+        face_mask: A "Tensor" of shape of [F,1], which consists of 0 or 1 to determine
+        which faces are valid
+        smoothing: A 'bool', whether to make meshes smooth
+    """
+
+
+    triangle_mesh = o3d.geometry.TriangleMesh()
+    verts_numpy = verts.numpy()
+    fm = faces_mask.numpy() == 1
+    faces_numpy = faces.numpy()[fm]
+
+    # Convert numpy array to Open3D format
+    triangle_mesh.vertices = o3d.utility.Vector3dVector(verts_numpy)
+    triangle_mesh.triangles = o3d.utility.Vector3iVector(faces_numpy)
+
+    # If smoothing is True, smooth the meshes
+    if smoothing:
+        mesh_out = triangle_mesh.filter_smooth_simple(number_of_iterations = 5)
+
+    # Compute triangle normals
+    mesh_out.compute_vertex_normals()
+
+    # Draw the meshes and visualize
+    o3d.visualization.draw_geometries([mesh_out])
+
+
+if __name__ == '__main__':
+    _grid_dims = 2
+    _batch_size = 5
+    _occupancy_locs = [
+        [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 1],
+
+        [1, 0, 0, 0], [1, 1, 1, 0], [1, 1, 0, 0], [1, 1, 0, 1],
+
+        [3, 0, 0, 0], [3, 0, 0, 1], [3, 0, 1, 0], [3, 0, 1, 1],
+        [3, 1, 0, 0], [3, 1, 0, 1], [3, 1, 1, 0], [3, 1, 1, 1],
+    ]
+    voxels = create_voxels(_grid_dims, _batch_size, _occupancy_locs)
+
+    mesh = cubify(voxels, 0.5)
+    _verts = mesh['verts']
+    _faces = mesh['faces']
+    _verts_mask = mesh['verts_mask']
+    _faces_mask = mesh['faces_mask']
+
+    batch_to_view = 0
+    visualize_mesh(_verts[batch_to_view, :],
+                   _faces[batch_to_view, :],
+                   _verts_mask[batch_to_view, :],
+                   _faces_mask[batch_to_view, :])
+
+    plt.show()

official/vision/beta/projects/mesh_rcnn/ops/visualize_mesh.py

@@ -2,61 +2,81 @@
 
 import matplotlib.pyplot as plt
 import tensorflow as tf
+import open3d as o3d
 from mpl_toolkits.mplot3d import art3d
 
 from official.vision.beta.projects.mesh_rcnn.ops.cubify import cubify
 
 
 def create_voxels(grid_dims, batch_size, occupancy_locs):
-  ones = tf.ones(shape=[len(occupancy_locs)])
-  voxels = tf.scatter_nd(
-      indices=tf.convert_to_tensor(occupancy_locs, tf.int32),
-      updates=ones,
-      shape=[batch_size, grid_dims, grid_dims, grid_dims])
+    ones = tf.ones(shape=[len(occupancy_locs)])
+    voxels = tf.scatter_nd(
+        indices=tf.convert_to_tensor(occupancy_locs, tf.int32),
+        updates=ones,
+        shape=[batch_size, grid_dims, grid_dims, grid_dims])
 
-  return voxels
+    return voxels
 
-def visualize_mesh(verts, faces, verts_mask, faces_mask):
-  v = verts.numpy()
-  f = faces.numpy()
-  vm = verts_mask.numpy() == 1
-  fm = faces_mask.numpy() == 1
 
-  new_f = f[fm]
+def visualize_mesh(verts, faces, verts_mask, faces_mask, smoothing = True):
 
-  fig = plt.figure()
-  ax = fig.add_subplot(projection="3d")
+    """Visualize meshes using Open3D
 
-  pc = art3d.Poly3DCollection(
-      v[new_f], facecolors=(1, 0.5, 1, 1), edgecolor="black")
+    Args:
+        verts: A "Tensor" of shape of [V,3], where V is the number of vertices
+        faces: A "Tensor" of shape of [F,3], where each face is composed of 3 integers,
+        which map to the points in the vertex tensor
+        verts_mask: A "Tensor" of shape of [V,1], which consists of 0 or 1 to determine
+        which vertices are valid
+        face_mask: A "Tensor" of shape of [F,1], which consists of 0 or 1 to determine
+        which faces are valid
+        smoothing: A 'bool', whether to make meshes smooth
+    """
 
-  ax.add_collection(pc)
 
-  plt.axis('off')
+    triangle_mesh = o3d.geometry.TriangleMesh()
+    verts_numpy = verts.numpy()
+    fm = faces_mask.numpy() == 1
+    faces_numpy = faces.numpy()[fm]
+
+    # Convert numpy array to Open3D format
+    triangle_mesh.vertices = o3d.utility.Vector3dVector(verts_numpy)
+    triangle_mesh.triangles = o3d.utility.Vector3iVector(faces_numpy)
+
+    # If smoothing is True, smooth the meshes
+    if smoothing:
+        mesh_out = triangle_mesh.filter_smooth_simple(number_of_iterations = 5)
+
+    # Compute triangle normals
+    mesh_out.compute_vertex_normals()
+
+    # Draw the meshes and visualize
+    o3d.visualization.draw_geometries([mesh_out])
+
 
 if __name__ == '__main__':
-  _grid_dims = 2
-  _batch_size = 5
-  _occupancy_locs = [
-      [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 1],
-
-      [1, 0, 0, 0], [1, 1, 1, 0], [1, 1, 0, 0], [1, 1, 0, 1],
-
-      [3, 0, 0, 0], [3, 0, 0, 1], [3, 0, 1, 0], [3, 0, 1, 1],
-      [3, 1, 0, 0], [3, 1, 0, 1], [3, 1, 1, 0], [3, 1, 1, 1],
-  ]
-  voxels = create_voxels(_grid_dims, _batch_size, _occupancy_locs)
-
-  mesh = cubify(voxels, 0.5)
-  _verts = mesh['verts']
-  _faces = mesh['faces']
-  _verts_mask = mesh['verts_mask']
-  _faces_mask = mesh['faces_mask']
-
-  batch_to_view = 0
-  visualize_mesh(_verts[batch_to_view, :],
-                 _faces[batch_to_view, :],
-                 _verts_mask[batch_to_view, :],
-                 _faces_mask[batch_to_view, :])
-
-  plt.show()
+    _grid_dims = 2
+    _batch_size = 5
+    _occupancy_locs = [
+        [0, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 1],
+
+        [1, 0, 0, 0], [1, 1, 1, 0], [1, 1, 0, 0], [1, 1, 0, 1],
+
+        [3, 0, 0, 0], [3, 0, 0, 1], [3, 0, 1, 0], [3, 0, 1, 1],
+        [3, 1, 0, 0], [3, 1, 0, 1], [3, 1, 1, 0], [3, 1, 1, 1],
+    ]
+    voxels = create_voxels(_grid_dims, _batch_size, _occupancy_locs)
+
+    mesh = cubify(voxels, 0.5)
+    _verts = mesh['verts']
+    _faces = mesh['faces']
+    _verts_mask = mesh['verts_mask']
+    _faces_mask = mesh['faces_mask']
+
+    batch_to_view = 0
+    visualize_mesh(_verts[batch_to_view, :],
+                   _faces[batch_to_view, :],
+                   _verts_mask[batch_to_view, :],
+                   _faces_mask[batch_to_view, :])
+
+    plt.show()

official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py

@@ -1,19 +1,17 @@
-
 import numpy as np
 import tensorflow as tf
+import open3d as o3d
 from matplotlib import pyplot as plt
 
 from official.vision.beta.projects.mesh_rcnn.modeling.heads.mesh_head import \
     MeshHead
-from official.vision.beta.projects.mesh_rcnn.modeling.heads.z_head import \
-    ZHead
 from official.vision.beta.projects.mesh_rcnn.ops.cubify import cubify
 from official.vision.beta.projects.mesh_rcnn.ops.mesh_ops import \
     compute_mesh_shape
 from official.vision.beta.projects.mesh_rcnn.ops.visualize_mesh import \
     visualize_mesh
 from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
-    load_weights_mesh_head, load_weights_zhead, pth_to_dict)
+    load_weights_mesh_head, pth_to_dict)
 
 PTH_PATH = r"C:\ML\Weights\meshrcnn_R50.pth"
 BACKBONE_FEATURES = [
@@ -37,26 +35,60 @@ def print_layer_names(layers_dict, offset=0):
       print(" " * offset + k)
       print_layer_names(layers_dict[k], offset+2)
 
-def test_load_zhead():
+def test_load_mesh_refinement_branch():
   weights_dict, n_read = pth_to_dict(PTH_PATH)
-  print(weights_dict.keys())
-  print(weights_dict['roi_heads'].keys())
-  print(weights_dict['roi_heads']['z_head'].keys())
-  print(weights_dict['roi_heads']['z_head']['z_pred'].keys())
-  print(weights_dict['roi_heads']['z_head']['z_pred']['weight'].shape)
-
-  input_specs = dict(
-    num_fc = 2,
-    fc_dim = 1024,
-    cls_agnostic = False,
-    num_classes = 9
-  )
-
-  zhead = ZHead.from_config(input_specs)
-
-  n_weights = load_weights_zhead(
-      zhead, weights_dict['roi_heads']['z_head'], 'pix3d')
-
+
+  grid_dims = 24
+  mesh_shapes = compute_mesh_shape(len(VOXEL_HEAD_OUTPUTS), grid_dims)
+  verts_shape, verts_mask_shape, faces_shape, faces_mask_shape = mesh_shapes
+  backbone_shape = [14, 14, 256]
+  input_layer = {
+      'feature_map': tf.keras.layers.Input(shape=backbone_shape),
+      'verts': tf.keras.layers.Input(shape=verts_shape[1:]),
+      'verts_mask': tf.keras.layers.Input(shape=verts_mask_shape[1:]),
+      'faces': tf.keras.layers.Input(shape=faces_shape[1:]),
+      'faces_mask': tf.keras.layers.Input(shape=faces_mask_shape[1:])
+  }
+  mesh_head = MeshHead()(input_layer)
+  model = tf.keras.Model(inputs=[input_layer], outputs=[mesh_head])
+
+  n_weights = load_weighimport numpy as np
+import tensorflow as tf
+import open3d as o3d
+from matplotlib import pyplot as plt
+
+from official.vision.beta.projects.mesh_rcnn.modeling.heads.mesh_head import \
+    MeshHead
+from official.vision.beta.projects.mesh_rcnn.ops.cubify import cubify
+from official.vision.beta.projects.mesh_rcnn.ops.mesh_ops import \
+    compute_mesh_shape
+from official.vision.beta.projects.mesh_rcnn.ops.visualize_mesh import \
+    visualize_mesh
+from official.vision.beta.projects.mesh_rcnn.utils.weight_utils.load_weights import (
+    load_weights_mesh_head, pth_to_dict)
+
+PTH_PATH = r"C:\ML\Weights\meshrcnn_R50.pth"
+BACKBONE_FEATURES = [
+  r"C:\ML\sofa_0134_mesh_features.npy",
+  r"C:\ML\bed_0003_mesh_features.npy",
+  r"C:\ML\bookcase_0002_mesh_features.npy",
+  r"C:\ML\chair_0093_mesh_features.npy",
+  r"C:\ML\table_0169_mesh_features.npy",
+]
+VOXEL_HEAD_OUTPUTS = [
+  r"C:\ML\sofa_0134_voxels.npy",
+  r"C:\ML\bed_0003_voxels.npy",
+  r"C:\ML\bookcase_0002_voxels.npy",
+  r"C:\ML\chair_0093_voxels.npy",
+  r"C:\ML\table_0169_voxels.npy",
+]
+
+def print_layer_names(layers_dict, offset=0):
+  if isinstance(layers_dict, dict):
+    for k in layers_dict.keys():
+      print(" " * offset + k)
+      print_layer_names(layers_dict[k], offset+2)
+
 def test_load_mesh_refinement_branch():
   weights_dict, n_read = pth_to_dict(PTH_PATH)
 
@@ -77,6 +109,60 @@ def test_load_mesh_refinement_branch():
   n_weights = load_weights_mesh_head(
       model, weights_dict['roi_heads']['mesh_head'], 'pix3d')
 
+  batched_backbone_features = []
+  print("backbone features shapes")
+  for f in BACKBONE_FEATURES:
+    backbone_features = np.load(f)
+    print(backbone_features.shape)
+    batched_backbone_features.append(backbone_features)
+
+  batched_backbone_features = np.concatenate(batched_backbone_features, axis=0)
+
+  batched_voxels = []
+  print("voxels shapes")
+  for f in VOXEL_HEAD_OUTPUTS:
+    voxels = np.load(f)
+    print(voxels.shape)
+    batched_voxels.append(voxels)
+
+  batched_voxels = np.concatenate(batched_voxels, axis=0)
+
+  backbone_features = tf.convert_to_tensor(batched_backbone_features, tf.float32)
+  backbone_features = tf.transpose(backbone_features, [0, 2, 3, 1])
+  voxels = tf.convert_to_tensor(batched_voxels, tf.float32)
+
+  mesh = cubify(voxels, 0.2)
+  verts = mesh['verts']
+  faces = mesh['faces']
+  verts_mask = mesh['verts_mask']
+  faces_mask = mesh['faces_mask']
+
+  inputs = {
+      'feature_map': backbone_features,
+      'verts': verts,
+      'verts_mask': verts_mask,
+      'faces': faces,
+      'faces_mask': faces_mask
+    }
+
+  outputs = model(inputs)[0]
+  new_verts_0 = outputs['verts']['stage_0']
+  new_verts_1 = outputs['verts']['stage_1']
+  new_verts_2 = outputs['verts']['stage_2']
+
+  batch_to_view = 1
+  for batch_to_view in range(len(VOXEL_HEAD_OUTPUTS)):
+    visualize_mesh(new_verts_2[batch_to_view, :],
+                  faces[batch_to_view, :],
+                  verts_mask[batch_to_view, :],
+                  faces_mask[batch_to_view, :]
+    )
+
+if __name__ == '__main__':
+  test_load_mesh_refinement_branch()
+ts_mesh_head(
+      model, weights_dict['roi_heads']['mesh_head'], 'pix3d')
+
   batched_backbone_features = []
   print("backbone features shapes")
   for f in BACKBONE_FEATURES:
@@ -126,8 +212,5 @@ def test_load_mesh_refinement_branch():
                   faces_mask[batch_to_view, :]
     )
 
-  plt.show()
-
 if __name__ == '__main__':
   test_load_mesh_refinement_branch()
-  test_load_zhead()