From cf4b67c5a7f39f6a8eb87d38147ee67c845092e5 Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Fri, 17 Jun 2022 21:42:45 +0200
Subject: [PATCH] add plotly plotting

---
 reconstruct.py | 55 ++++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/reconstruct.py b/reconstruct.py
index 14d2606..411ab15 100644
--- a/reconstruct.py
+++ b/reconstruct.py
@@ -4,6 +4,8 @@ import cv2
 import numpy as np
 import trimesh as tm
 import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+
 
 def m2uv_array(m2uv):
     result = np.zeros((14475, 2), dtype=np.int32)
@@ -83,17 +85,66 @@ def plot(uv, mesh, faces, body, F, template, uv_gt, mesh_gt):
     collision_color = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh, dtype=np.int32)*230), axis=0)
 
 
+    plot_3d(mesh_gt, mesh, faces, collision, collision_faces, collision_color)
     tm.Trimesh(mesh, faces).show()
     tm.Trimesh(mesh_gt, faces).show()
     # tm.Trimesh(collision, collision_faces, vertex_colors=collision_color).show()
 
 
+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],),)
+                        )
+
+    fig.show()
+
 
 if __name__ == '__main__':
-    file = sys.argv[1]
+    # file = sys.argv[1]
 
     for i in range(50):
-        file = r'Z:\Work\Huni\Barcelona\Output\stylegan\test-12-15000\\' + str(i) + '.npz'
+        file = r'Z:\Work\Huni\Barcelona\Output\stylegan\test-15-115000\\' + str(i) + '.npz'
 
         F = np.load('custom_utils\smpl_faces.npy')
         data = np.load(file)