large refactor

auto_encoder.py

@@ -21,7 +21,7 @@ from stylegan2_ae.generator import Generator
 from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
 from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
 from dataset import DatasetFactory, StaticDataset
-from custom_utils.custom_utils import batch_to_dict
+from custom_utils.custom_utils import batch_to_dict, plot_3d
 from custom_utils.smpl import SMPL
 
 name='ae_body_1'
@@ -113,53 +113,6 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
 #     return fig
 
 
-def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
-    mesh_pred_unposed[:,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_unposed[:,0],
-            y=mesh_pred_unposed[:,1],
-            z=mesh_pred_unposed[:,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
 
 class Bottleneck(tf.keras.Model):