new dataloader

data/dataset.py

@@ -1,79 +1,101 @@
 import os
 import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+
+import math
 import time
-import cv2
+import torch
 import random
-import pickle
+import pytorch3d
+import torchvision as tv
 import numpy as np
-import tensorflow as tf
 import scipy.io as sio
+import plotly.express as ex
+import torchvision
+
+from torch.utils.data import Dataset
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from custom_utils.custom_utils import nest_conditions, batch_to_dict
+from data.utils import Frame
-from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame
 
 
-class StaticDataset():
+def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> 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 = [Frame(line.split('_')[0], line.split('_')[1]) for line in lines]
+
+    if shuffle:
+        random.shuffle(frames)
+
+    return frames
+
+
+class StaticCache():
+    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False) -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
         self.in_memory = in_memory
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
-        self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256',
-                              'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256',
-                              'mesh_static', 'mesh_mask', 'uv_shaped_body']
-        self.keys = [value for value in signature.keys() if value in self.possible_keys]
-        if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys:
-            self.keys.append('uv_static_128')
         self.resolutions = resolutions
         self.filled = ''
         self.is_cape = is_cape
-        self.use_no_orient = use_no_orient
+        
-
         if is_cape:
             self.uv_masks = dict()
             uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32)
             self.uv_masks[512] = uv_mask.astype(np.float32)
             for res in [32, 64, 128, 256]:
-                down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+                down = Resize(uv_mask[:,:,None], size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST)
                 down = np.squeeze(down.numpy().astype(np.float32))
                 self.uv_masks[res] = down
 
 
-    def read(self, id, garment):
+    def read(self, id):
         data = dict()
 
-        info = loadInfo(os.path.join(self.folder, id))
-
-        data['shape'] = info['shape']
-        data['poses'] = info['poses']
-
-        data['map'] = sio.loadmat(os.path.join(self.folder, id, garment + '_map.mat'))['map'].reshape((-1))
-
-
-        if 'mesh_mask' in self.signature.keys():
-            if self.is_cape:
-                data['mesh_mask'] = np.ones((3638), np.float32)
-            else:
-                data['mesh_mask'] = read_garment_mask(self.folder, id, garment).astype(np.float32)
-
         for res in self.resolutions:
             res_string = ''
             if res != 512:
                 res_string = '_' + str(res)
 
-            if 'uv_static' + res_string in self.signature.keys():
+            if 'uv_static' + res_string in self.signature:
                 uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id +  '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
-            if 'uv_mask' + res_string in self.signature.keys():
+            if 'uv_mask' + res_string in self.signature:
                 if self.is_cape:
                     data['uv_mask' + res_string] = self.uv_masks[res]
                 else:
                     uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask
 
-        if 'mesh_static' in self.signature.keys():
+        if 'mesh_static' in self.signature:
-            data['mesh_static'] = read_garment_template(self.folder, id, garment)
+            mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id +  '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static.astype(np.float32)
+        if 'mesh_mask' in self.signature:
+            if self.is_cape:
+                data['mesh_mask'] = np.ones((3638), np.float32)
+            else:
+                mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask
+        # if 'uv_shaped_body' in self.signature:
+        #     uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body
 
         return data
 
@@ -90,462 +112,186 @@ class StaticDataset():
 
 
     def expand_batch(self, batch):
-        for key in self.keys:
+        batch_size = batch['id'].shape[0]
-                batch[key] = tf.zeros(((0,) + self.signature[key]))
+        items = []
-        for i in range(batch['id'].shape[0]):
+        for i in range(batch_size):
-            item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8'))
+            items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8')))
-            for key in self.keys:
+        for key in items[0].keys():
-                if key == 'pose':
+            batch[key] = torch.stack([item[key] for item in items], dim=0)
-                    frame_nr = int(batch['frame_nr'][i].numpy().decode('UTF-8'))
-                    pose = item[key][frame_nr, :]
-                    if self.use_no_orient:
-                        pose[:3] = 0
-                    batch[key] = tf.concat((batch[key], tf.expand_dims(pose, axis=0)), axis=0)
-                else:
-                    batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0)
         return batch
 
 
+class Dataset(Dataset):
+    def __init__(self, folder, signature, resolutions, frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+        self.folder = folder
+        self.device = device
+        self.is_cape = is_cape
+        self.signature = signature
+        self.resolutions = resolutions
 
-class DatasetIdx(tf.data.Dataset):
+        self.transforms = tv.transforms.ToTensor()
+        self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
 
-    def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
+        self.body_faces = np.load('smpl/smpl_faces.npy')
-        # if frames_file is None:
+        self.static_cache = StaticCache(folder, signature, is_cape=is_cape)
-        #     uv_static_folder = os.path.join(folder, 'pose_shape') 
-        #     frames = sorted(os.listdir(uv_static_folder))
 
-        #     if is_cape:
+        self.cache = dict() if use_cache else None
-        #         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]
-
-        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 shuffle:
-            random.shuffle(frames)
-
-        return tf.data.Dataset.from_tensor_slices(frames)
 
 
-class DatasetFactory():
+    def __len__(self):
+        return len(self.frames)
 
-    def __init__(self):
-        super(DatasetFactory, self).__init__()
 
-    def _build_tf_signature(self, signature):
+    def __getitem__(self, idx):
-        tf_signature = []
+        frame : Frame = self.frames[idx]
-        for key in signature.keys():
+        if self.cache is not None:
-            if key == 'id':
+            if frame in self.cache.keys():
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
+                return self.cache[frame]
-            elif key == 'frame_nr':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            elif key == 'garment':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            elif key == 'mesh_curvature':
-                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
-            else:
-                tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
-        return tf_signature
-
-    
-    def get_dataset(self,
-                    folder, 
-                    batch_size, 
-                    signature, 
-                    shuffle=True, 
-                    resolutions=[512], 
-                    in_memory=False, 
-                    do_repeat=True, 
-                    is_cape=False,
-                    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, testing=testing, frames_file=frames_file)
 
         separator = '_'
-        if is_cape:
+        if self.is_cape:
             separator = '.'
 
-        uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
+        uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2'
+        if self.is_cape:
+            uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient'
 
-        if use_no_orient == False:
+        item = dict()
-            assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
 
-        to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
+        if 'id' in self.signature:
-        shapedirs = []
+            item['id'] = frame.id
-        with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
+        if 'frame_nr' in self.signature:
-            female = pickle.load(f, encoding="latin1")
+            item['frame_nr'] = frame.frame_nr
-            shapedirs.append(female['shapedirs'][to_keep, :])
+        if 'pose_shape' in self.signature:
-        with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
+            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy'))
-            male = pickle.load(f, encoding="latin1")
+            pose_shape[:3] = 0
-            shapedirs.append(male['shapedirs'][to_keep, :])
+            item['pose_shape'] = pose_shape
+        if 'mesh_body_posed' in self.signature:
+            mesh_posed_body_no_orient = np.load(os.path.join(self.folder, 'mesh_posed_body_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
+            item['mesh_body_posed'] = mesh_posed_body_no_orient
+        if 'mesh_posed' in self.signature:
+            mesh_posed = np.load(os.path.join(self.folder, 'mesh_posed_no_orient', frame.id + separator + frame.frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
+            item['mesh_posed'] = mesh_posed
+        if 'mesh_unposed' in self.signature:
+            mesh_unposed_no_shape = np.load(os.path.join(self.folder, 'mesh_unposed_no_shape', frame.id + separator + frame.frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
+            item['mesh_unposed'] = mesh_unposed_no_shape
+        if 'mesh_curvature' in self.signature:
+            mesh_curvature = np.load(os.path.join(self.folder, 'mesh_curvature', frame.id + separator + frame.frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
+            item['mesh_curvature'] = mesh_curvature.astype(np.int32)
 
+        for res in self.resolutions:
+            res_string = ''
+            if res != 512:
+                res_string = '_' + str(res)
 
+            if 'uv_curvature' + res_string in self.signature:
+                uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz')); 
+                if res_string != '':
+                    uv_curvature = uv_curvature['uv_curvature'] 
+                else:
+                    uv_curvature = uv_curvature['curvature_uv']
+                uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).to(self.device)
+                item['uv_curvature' + res_string] = uv_curvature
 
-        def read_item(frame):
+            if 'uv_unposed' + res_string in self.signature:
-            frame_split = tf.strings.split(frame, sep=separator)
+                uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
-            id, frame_nr, garment = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8"), frame_split[2].numpy().decode("utf-8")
+                item['uv_unposed' + res_string] = uv_unposed_no_shape
-            item = dict()
 
-            # redunda ntly loading info again
+            if 'uv_body_posed' + res_string in self.signature:
-            pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr]
+                uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
-            # !!!!!!!!!!!!!!!!!!!
+                uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
-            map = None
+                item['uv_body_posed' + res_string] = uv_posed_body_no_orient
 
-            if 'id' in signature.keys():
+            if 'uv_posed' + res_string in self.signature:
-                item['id'] = id
+                uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
-            if 'frame_nr' in signature.keys():
+                item['uv_posed' + res_string] = uv_posed
-                item['frame_nr'] = frame_nr
-            if 'garment' in signature.keys():
-                item['garment'] = garment
 
-            if 'mesh_body_posed' in signature.keys():
+        data = { **item, **self.static_cache[frame.id] }
-                if use_no_orient:
-                    item['mesh_body_posed'] = read_posed_body_frame(folder, id, frame_nr, pose, deorient=use_no_orient)
-            if 'mesh_posed' in signature.keys():
-                if use_no_orient:
-                    item['mesh_posed'] = read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False)
-            # if 'mesh_unposed' in signature.keys():
-            #     mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
-            #     item['mesh_unposed'] = mesh_unposed_no_shape
 
-            if 'mesh_curvature' in signature.keys():
+        data = { key: (torch.from_numpy(value).to(self.device) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in data.items()}
-                mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
-                item['mesh_curvature'] = mesh_curvature.astype(np.int32)
-
-
-            for res in resolutions:
-                res_string = ''
-                if res != 512:
-                    res_string = '_' + str(res)
-
-                if 'uv_curvature' + res_string in signature.keys():
-                    uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + res_string, id + separator + frame_nr + '.npz')); 
-                    if res_string != '':
-                        uv_curvature = uv_curvature['uv_curvature'] 
-                    else:
-                        uv_curvature = uv_curvature['curvature_uv']
-                    uv_curvature = tf.one_hot(uv_curvature, 7)
-                    item['uv_curvature' + res_string] = uv_curvature
-
-                if 'uv_unposed' + res_string in signature.keys():
-                    uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
-                    item['uv_unposed' + res_string] = uv_unposed_no_shape
-
-                if 'uv_body_posed' + res_string in signature.keys():
-                    if use_no_orient:
-                        uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz'))
-                        uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
-                        item['uv_body_posed' + res_string] = uv_posed_body_no_orient
-                    else:
-                        uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']
-                        item['uv_body_posed' + res_string] = uv_body_posed
-
-                if 'uv_posed' + res_string in signature.keys():
-                    if use_no_orient:
-                        uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
-                        item['uv_posed' + res_string] = uv_posed
-                    else:
-                        uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
-                        item['uv_posed' + res_string] = uv_posed
-
-            as_list = []
-            for key in signature.keys():
-                if key in item:
-                    as_list.append(item[key])
-
-            return as_list
-
-
-        if in_memory:
-            dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, 
-                                                inp=[i], 
-                                                Tout=self._build_tf_signature(signature)
-                                                ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True)
-        else:
-            dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, 
-                                                inp=[i], 
-                                                Tout=self._build_tf_signature(signature)
-                                                ), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True)
-
-        if do_repeat:
-            dataset = dataset.repeat()
-        
-        return dataset
 
+        if self.cache is not None:
+            self.cache[frame] = data
+        return data
 
 
 if __name__ == '__main__':
-    os.environ["CUDA_VISIBLE_DEVICES"] = "6"
+    from torch.utils.data import DataLoader
-    from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces
-    import plotly.graph_objects as go
-    import plotly.express as ex
-    # import cmapy as cm
-    uv2mesh = UV2Mesh(is_skirt=True)
-
-    F = np.load("custom_utils/smpl_faces.npy")
-
-
-    # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train')
-    # for d in dataset:
-    #     print(d)
-
-    factory = DatasetFactory()
-    signature_dynamic = { 'id': (), 'frame_nr': (), 'garment': (),
-                        #   'pose_shape': (72 + 11,),
-                        #   'uv_posed': (3, 512, 256,),
-                        #   'uv_body_posed': (512, 256, 3),
-                        #   'uv_body_posed_32': (32, 16, 3),
-                        #   'mesh_body_posed': (14475, 3,),
-                        #   'mesh_posed': (14475, 3,),
-                        #   'mesh_unposed': (14475, 3,),
-                          'uv_unposed': (512, 256, 3),
-                        #   'uv_unposed_32': (3, 32, 16,),
-                          'uv_curvature': (512, 256, 7),
-                        #   'mesh_curvature':  (14475)
-                        }
-    signature_static = { #'uv_static': (3, 512, 256,),
-                        #  'uv_static_32': (3, 32, 16,),
-                        #  'mesh_static': (14475, 3,),
-                        #  'mesh_mask': (14475,),
-                         'uv_mask': (512, 256,),
-                        #  'uv_shaped_body': (3, 512, 256,),
-                          }
-
-    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt',
-                                  batch_size=1,
-                                  signature=signature_dynamic,
-                                  resolutions=[512],
-                                  do_repeat=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],
-                                  is_cape=False)
-
-
-    # for i, batch in enumerate(dataset):
-    #     print(i)
-
-
-    sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
-    sum_mask = np.zeros((512, 256), dtype=np.float32)
-    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)
-
-    sum_mask[sum_mask == 0] = 1
-    avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
-
-    # x = np.transpose(avg_garment, [1,2,0])
-    x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
-    go.Figure(data=[
-            go.Mesh3d(
-                # 8 vertices of a cube
-                x=x[:,0],
-                y=x[:,1],
-                z=x[:,2],
-                # i, j and k give the vertices of triangles
-                i = F[:,0],
-                j = F[:,1],
-                k = F[:,2]
-            )]).show()
-
-    np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
-
-    print()
-
-    # pass
-
-
-    # uv_res = 256
-
-    # def m2uv_array(m2uv):
-    #     result = np.zeros((14475, 2), dtype=np.int32)
-    #     for k, v in m2uv.items():
-    #         if len(v) == 2:
-    #             result[k] = list(v)
-    #         elif len(v) == 1:
-    #             x = v.pop()
-    #             result[k] = [x, x]
-    #         else:
-    #             assert False
-
-    #     return result
-
-    # def unique_uv_map_mask(uv_map_discreet):
-    #     mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
-    #     temp = np.zeros((uv_res * 2, uv_res))
-    #     for p in uv_map_discreet:
-    #         temp[p[0], p[1]] += 1
-    #     # temp = temp == 1
-    #     for i, p in enumerate(uv_map_discreet):
-    #         if temp[p[0], p[1]] == 1:
-    #             mask[i] = True
-    #         elif temp[p[0], p[1]] == 2:
-    #             mask[i] = True
-    #             temp[p[0], p[1]] = 0
-    #     return mask
-
-    # def uv_to_mesh(uv, m2uv):
-    #     res = uv.shape[0]
-    #     mesh = np.zeros((14475, 3))
-    #     uv_pixels = uv_map_discreet[m2uv]
-    #     # counter = np.ones((14475))
-    #     # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
-    #     mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
-    #     return mesh
-
-
-    # uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
-    # m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
-    # m2uv = m2uv_array(m2uv)
-    # uv_mask = np.load('custom_utils/mask_256x512.npy')
-    # curvature_label_type = '6_25'
-
-    # uv_res = 256
-    # uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
-    # uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
-
-    # uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
-    # uv_map_mask = unique_uv_map_mask(uv_map_discreet)
-    # uv_counter = np.zeros((uv_map.shape[0]))
-    # uv_counter[m2uv[:, 0]] = 1
-    # uv_counter[m2uv[:, 1]] += 1
-    # uv_counter = uv_counter[uv_map_mask]
-
-
-
-    # for i, batch in enumerate(dataset):
-    #     print(i)
-    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-
-    #     uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
-    #     mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
-    #     mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
-
-
-    #     values = mesh_unposed
-    #     uv_values = np.zeros((uv_map.shape[0], 3))
-    #     uv_values[m2uv[:, 0]] = values
-    #     uv_values[m2uv[:, 1]] += values
-    #     uv_counter = np.zeros((uv_map.shape[0]))
-    #     uv_counter[m2uv[:, 0]] = 1
-    #     uv_counter[m2uv[:, 1]] += 1
-    #     uv_values /= uv_counter[:, np.newaxis]
-    #     x = uv_values
-    #     # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
-    #     im = ((x - x.min()) / (x.max() - x.min())) * 255
-
-
-    #     x = np.transpose(uv_unposed, [1,2,0])
-    #     x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
-    #     _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
-    #     go.Figure(data=[
-    #             go.Mesh3d(
-    #                 # 8 vertices of a cube
-    #                 x=x[:,0],
-    #                 y=x[:,1],
-    #                 z=x[:,2],
-    #                 # i, j and k give the vertices of triangles
-    #                 i = faces[:,0],
-    #                 j = faces[:,1],
-    #                 k = faces[:,2]
-    #             )]).show()
-
-
-    #     print()
-
-
-
-
-    # # @tf.function
-    # def euclid_error(mask, gt, pred):
-    #     # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1)
-    #     return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
-
-    # all_error = []
-    # for batch in dataset:
-    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-    #     mesh_reconstruct = uv2mesh(batch['uv_posed'])
-    #     err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
-
-    #     i = 296
-    #     err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
-    #     colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
-
-    #     _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
-    #     x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
-    #     fig = go.Figure(data=[
-    #         go.Mesh3d(
-    #             # 8 vertices of a cube
-    #             x=x[:,0],
-    #             y=x[:,1],
-    #             z=x[:,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=y[:,0],
-    #             y=y[:,1],
-    #             z=y[:,2],
-    #             # i, j and k give the vertices of triangles
-    #             i = faces[:,0],
-    #             j = faces[:,1],
-    #             k = faces[:,2],
-    #         )])
-    #     fig.show()
-
-    #     fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
-    #     fig.update_traces(marker={'color': colors, 'size': 5})
-    #     fig.show()
-
-    #     all_error.append(err.numpy().sum())
-    #     # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
-    #     print(len(all_error))
-
     
-    # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # from losses import collision_loss as collision_penalty
+    signature = { 'id', 
-    # from losses import collision_penalty as collision_penalty_2
+                  'frame_nr',
+                  'pose_shape',
+                  'uv_posed',
+                  'uv_body_posed',
+                  'mesh_body_posed',
+                  'mesh_posed',
+                  'mesh_unposed',
+                  'uv_unposed',
+                  'uv_curvature',
+                  'mesh_curvature',
+                  'uv_static',
+                  'mesh_static',
+                  'mesh_mask',
+                  'uv_mask',
+                  'uv_shaped_body',
+                  }
+    
+    dataset = Dataset('/data-local/127/train', signature, resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
+
+    for batch in train_dataloader:
+        batch
+        pass
+
+    pass 
 
 
-    # for d in dataset:
-    #     d
 
-    #     v = uv2mesh(d[0])
-    #     b = d[3]
-    #     m = d[4]
 
-    #     b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
 
-    #     c = collision_penalty(v, b, m)
-    #     c2 = collision_penalty(v, b2, m)
 
-    #     c_ = collision_penalty_2(v, b, m)
-    #     c2_ = collision_penalty_2(v, b2, m)
 
-    #     print()
+
-    #     pass
+
-    
+
+# 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 cmapy as cm
+#     uv2mesh = UV2Mesh(is_skirt=True)
+
+#     F = np.load("custom_utils/smpl_faces.npy")
+
+#     sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
+#     sum_mask = np.zeros((512, 256), dtype=np.float32)
+#     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)
+
+#     sum_mask[sum_mask == 0] = 1
+#     avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
+
+#     # x = np.transpose(avg_garment, [1,2,0])
+#     x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
+#     go.Figure(data=[
+#             go.Mesh3d(
+#                 # 8 vertices of a cube
+#                 x=x[:,0],
+#                 y=x[:,1],
+#                 z=x[:,2],
+#                 # i, j and k give the vertices of triangles
+#                 i = F[:,0],
+#                 j = F[:,1],
+#                 k = F[:,2]
+#             )]).show()
+
+#     np.save('custom_data/train_dress_30k_no_shape_avg_uv.npy', avg_garment)
+
+#     print()

data/utils.py

@@ -1,16 +1,16 @@
 import tables as tb
 
 class Frame:
-    def __init__(self, id, frame_nr, garment):
+    def __init__(self, id, frame_nr, garment=''):
         self.id = id
         self.frame_nr = frame_nr
         self.garment = garment
 
     def __hash__(self):
-        return hash((self.id, self.frame_nr))
+        return hash((self.id, self.frame_nr, self.garment))
 
     def __eq__(self, other):
-        return (self.id, self.frame_nr) == (other.id, other.frame_nr)
+        return (self.id, self.frame_nr, self.garment) == (other.id, other.frame_nr, self.garment)
 
 
 def save_to_h5file(h5file, id, frame_nr, garment, data, name, type):