commit all

dataset.py

@@ -135,14 +135,27 @@ class StaticDataset():
 
 class DatasetIdx(tf.data.Dataset):
 
-    def __new__(cls, folder, shuffle=True, is_cape=False):
-        uv_static_folder = os.path.join(folder, 'pose_shape') 
-        frames = sorted(os.listdir(uv_static_folder))
+    def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
+        if frames_file is None:
+            uv_static_folder = os.path.join(folder, 'pose_shape') 
+            frames = sorted(os.listdir(uv_static_folder))
+
+            if is_cape:
+                frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
+            else:
+                frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+
+            if testing:
+                frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') ))
+                frames = [frame.split('.')[0] + '_4' for frame in frames]
+
+        if frames_file is not None:
+            frames = []
+            with open(os.path.join(folder, frames_file), 'r') as f:
+                lines = f.readlines()
+                lines = [line.strip() for line in lines]
+            frames = lines
 
-        if is_cape:
-            frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0]
-        else:
-            frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
         if shuffle:
             random.shuffle(frames)
 
@@ -179,10 +192,12 @@ class DatasetFactory():
                     do_repeat=True, 
                     is_cape=False, 
                     use_one_hot=False,
-                    use_no_orient=True):
+                    use_no_orient=True,
+                    testing=False,
+                    frames_file=None):
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
 
-        dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape)
+        dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file)
 
         separator = '_'
         if is_cape:
@@ -333,11 +348,11 @@ if __name__ == '__main__':
     os.environ["CUDA_VISIBLE_DEVICES"] = "6"
     from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
     import plotly.graph_objects as go
-    # import plotly.express as ex
+    import plotly.express as ex
     # import cmapy as cm
     uv2mesh = UV2Mesh(is_skirt=True)
 
-    F = np.load("custom_utils/smpl_faces_skirt.npy")
+    F = np.load("custom_utils/smpl_faces.npy")
 
 
     # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
@@ -355,7 +370,7 @@ if __name__ == '__main__':
                         #   'mesh_unposed': (14475, 3,),
                           'uv_unposed': (512, 256, 3),
                         #   'uv_unposed_32': (3, 32, 16,),
-                        #   'uv_curvature': (512, 256, 7),
+                          'uv_curvature': (512, 256, 7),
                         #   'mesh_curvature':  (14475)
                         }
     signature_static = { #'uv_static': (3, 512, 256,),
@@ -367,14 +382,17 @@ if __name__ == '__main__':
                           }
 
     dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
-                                  batch_size=300,
+                                  batch_size=1,
                                   signature=signature_dynamic,
                                   resolutions=[512],
                                   do_repeat=False,
-                                  use_one_hot=False)
+                                  use_one_hot=False,
+                                  is_cape=False,
+                                  frames_file='frames_dress_30k.txt')
     static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt',
                                    signature=signature_static,
-                                   resolutions=[512])
+                                   resolutions=[512],
+                                  is_cape=False)
 
 
     # for i, batch in enumerate(dataset):
@@ -386,6 +404,7 @@ if __name__ == '__main__':
     for i, batch in enumerate(dataset):
         print(i)
         batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+        ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show()
         sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
         sum_mask += batch['uv_mask'].numpy().sum(axis=0)
 
@@ -406,9 +425,9 @@ if __name__ == '__main__':
                 k = F[:,2]
             )]).show()
 
-    np.save('custom_data/train_no_shape_skirt_avg_uv.npy', avg_garment)
+    np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
 
-    # print()
+    print()
 
     # pass