name

display_utils.py

@@ -3,11 +3,13 @@ from matplotlib import pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 # plt.switch_backend('agg')
+from torch import Tensor
+from typing import Optional
 
 
 def display_model(
         model_info,
-        model_faces=None,
+        model_faces:Optional[Tensor]=None,
         with_joints=False,
         kintree_table=None,
         ax=None,
@@ -27,7 +29,7 @@ def display_model(
     if model_faces is None:
         ax.scatter(verts[:, 0], verts[:, 1], verts[:, 2], alpha=0.2)
     elif not only_joint:
-        mesh = Poly3DCollection(verts[model_faces], alpha=0.2)
+        mesh = Poly3DCollection(verts[model_faces.cpu()], alpha=0.2)
         face_color = (141 / 255, 184 / 255, 226 / 255)
         edge_color = (50 / 255, 50 / 255, 50 / 255)
         mesh.set_edgecolor(edge_color)
@@ -46,8 +48,8 @@ def display_model(
     if savepath:
         # print('Saving figure at {}.'.format(savepath))
         plt.savefig(savepath, bbox_inches='tight', pad_inches=0)
-    if show:
-        plt.show()
+    # if show:
+    #     plt.show()
     plt.close()
     return ax