commit all

verify_data.py

@@ -1,4 +1,5 @@
 import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "6"
 import sys
 import time
 import cv2
@@ -12,18 +13,17 @@ import plotly.express as ex
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
 from custom_utils.smpl import SMPL
+from convert_to_uv_api import convert_to_uv_api
 
 
 if __name__ == '__main__':
-    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
 
     uv2mesh = UV2Mesh()
-    smpl = SMPL()
 
     F = np.load("custom_utils/smpl_faces.npy")
 
     factory = DatasetFactory()
-    signature_dynamic = { 'id': (),
+    signature_dynamic = { 'id': (), 'frame_nr': (),
                         #   'pose_shape': (72 + 11,),
                           'uv_posed': (512, 256, 3),
                           'uv_body_posed': (512, 256, 3),
@@ -42,10 +42,10 @@ if __name__ == '__main__':
                         #  'uv_shaped_body': (3, 512, 256,),
                           }
 
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100',
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
                                   batch_size=1,
                                   signature=signature_dynamic)
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100',
+    static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
                                    signature=signature_static)
 
 
@@ -60,10 +60,9 @@ if __name__ == '__main__':
         uv_posed = (np.squeeze(batch['uv_posed'].numpy()) + 1) / 2
         uv_body_posed = (np.squeeze(batch['uv_body_posed'].numpy()) + 1) / 2
         uv_unposed = (np.squeeze(batch['uv_unposed'].numpy()) + 1) / 2
-        uv_curvature = np.tile(np.squeeze(batch['uv_curvature'].numpy())[:,:,np.newaxis], [1,1,3]) / 6
+        uv_curvature = np.tile(np.argmax(np.squeeze(batch['uv_curvature'].numpy()), axis=-1)[:,:,np.newaxis], [1,1,3]) / 6
         uv_static = (np.squeeze(batch['uv_static'].numpy()) + 1) / 2
 
-
         img = np.concatenate([uv_posed, uv_body_posed, uv_unposed, uv_curvature, uv_static], axis=1)
 
         ex.imshow(img).show()