test-20

custom_utils/custom_utils.py

@@ -1,5 +1,6 @@
 import numpy as np
 import tensorflow as tf
+import plotly.graph_objects as go
 
 
 def m2uv_array(m2uv):
@@ -85,4 +86,53 @@ def maskFacesBatch(V, F, mask):
         v, f = _maskFaces(V[i], F, mask[i])
         Vs.append(v)
         Fs.append(f)
-    return Vs, Fs
+    return Vs, Fs
+
+
+def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_color):
+    mesh_pred[:,1] += 1
+    collision[:,1] += 2
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_gt[:,0],
+            y=mesh_gt[:,1],
+            z=mesh_gt[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2]
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=mesh_pred[:,0],
+            y=mesh_pred[:,1],
+            z=mesh_pred[:,2],
+            # i, j and k give the vertices of triangles
+            i = faces[:,0],
+            j = faces[:,1],
+            k = faces[:,2],
+        ),
+        go.Mesh3d(
+            # 8 vertices of a cube
+            x=collision[:,0],
+            y=collision[:,1],
+            z=collision[:,2],
+            # i, j and k give the vertices of triangles
+            i = collision_faces[:,0],
+            j = collision_faces[:,1],
+            k = collision_faces[:,2],
+            vertexcolor=collision_color
+        )
+    ])
+
+    fig.update_layout(scene_aspectmode='manual',
+                    scene_aspectratio=dict(x=1, y=2, z=1))
+    fig.update_layout(
+        scene = dict(
+            xaxis = dict(nticks=4, range=[-1,1],),
+                        yaxis = dict(nticks=4, range=[-1,3],),
+                        zaxis = dict(nticks=4, range=[-1,1],),)
+                        )
+
+    return fig