diff --git a/official/vision/beta/projects/mesh_rcnn/demo/visualize_mesh.py b/official/vision/beta/projects/mesh_rcnn/demo/visualize_mesh.py new file mode 100644 index 00000000000..1b6c5036ac6 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/official/vision/beta/projects/mesh_rcnn/ops/visualize_mesh.py b/official/vision/beta/projects/mesh_rcnn/ops/visualize_mesh.py index 4a9a49610a6..4ee3e65098c 100644 --- a/official/vision/beta/projects/mesh_rcnn/ops/visualize_mesh.py +++ b/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() \ No newline at end of file diff --git a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py b/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py index 79fa6ccac7c..f97648c3b42 100644 --- a/official/vision/beta/projects/mesh_rcnn/utils/weight_utils/test_load_weights.py +++ b/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()