update docker

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.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame
+from data.utils import Frame
 
 
-class StaticDataset():
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> None:
+def read_frames(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 = [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():
-            data['mesh_static'] = read_garment_template(self.folder, id, garment)
+        if 'mesh_static' in self.signature:
+            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[key] = tf.zeros(((0,) + self.signature[key]))
-        for i in range(batch['id'].shape[0]):
-            item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8'), batch['garment'][i].numpy().decode('UTF-8'))
-            for key in self.keys:
-                if key == 'pose':
-                    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)
+        batch_size = batch['id'].shape[0]
+        items = []
+        for i in range(batch_size):
+            items.append(self.__getitem__(batch['id'][i].numpy().decode('UTF-8')))
+        for key in items[0].keys():
+            batch[key] = torch.stack([item[key] for item in items], dim=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):
-        # if frames_file is None:
-        #     uv_static_folder = os.path.join(folder, 'pose_shape') 
-        #     frames = sorted(os.listdir(uv_static_folder))
+        self.body_faces = np.load('smpl/smpl_faces.npy')
+        self.static_cache = StaticCache(folder, signature, is_cape=is_cape)
 
-        #     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]
-
-        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)
+        self.cache = dict() if use_cache else None
 
 
-class DatasetFactory():
+    def __len__(self):
+        return len(self.frames)
 
-    def __init__(self):
-        super(DatasetFactory, self).__init__()
 
-    def _build_tf_signature(self, signature):
-        tf_signature = []
-        for key in signature.keys():
-            if key == 'id':
-                tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
-            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)
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                return self.cache[frame]
 
         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:
-            assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!'
+        item = dict()
 
-        to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
-        shapedirs = []
-        with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
-            female = pickle.load(f, encoding="latin1")
-            shapedirs.append(female['shapedirs'][to_keep, :])
-        with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
-            male = pickle.load(f, encoding="latin1")
-            shapedirs.append(male['shapedirs'][to_keep, :])
+        if 'id' in self.signature:
+            item['id'] = frame.id
+        if 'frame_nr' in self.signature:
+            item['frame_nr'] = frame.frame_nr
+        if 'pose_shape' in self.signature:
+            pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy'))
+            pose_shape[:3] = 0
+            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):
-            frame_split = tf.strings.split(frame, sep=separator)
-            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 = dict()
+            if 'uv_unposed' + res_string in self.signature:
+                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']
+                item['uv_unposed' + res_string] = uv_unposed_no_shape
 
-            # redunda ntly loading info again
-            pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr]
-            # !!!!!!!!!!!!!!!!!!!
-            map = None
+            if 'uv_body_posed' + res_string in self.signature:
+                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']
+                item['uv_body_posed' + res_string] = uv_posed_body_no_orient
 
-            if 'id' in signature.keys():
-                item['id'] = id
-            if 'frame_nr' in signature.keys():
-                item['frame_nr'] = frame_nr
-            if 'garment' in signature.keys():
-                item['garment'] = garment
+            if 'uv_posed' + res_string in self.signature:
+                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']
+                item['uv_posed' + res_string] = uv_posed
 
-            if 'mesh_body_posed' in signature.keys():
-                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
+        data = { **item, **self.static_cache[frame.id] }
 
-            if 'mesh_curvature' in signature.keys():
-                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
+        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()}
 
+        if self.cache is not None:
+            self.cache[frame] = data
+        return data
 
 
 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")
-
-
-    # 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))
-
+    from torch.utils.data import DataLoader
     
-    # 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
-    # from losses import collision_penalty as collision_penalty_2
+    signature = { 'id', 
+                  '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):

docker/Dockerfile

@@ -1,5 +1,5 @@
 # Use an existing image as the base
-FROM nvcr.io/nvidia/pytorch:23.03-py3
+FROM nvcr.io/nvidia/pytorch:21.10-py3
 
 # Add the current user to the image
 ARG USER_NAME=lhunor

docker/requirements.txt

@@ -5,4 +5,6 @@ wandb
 scipy
 plotly
 pandas
-cmapy
+cmapy
+pytorch3d
+tables

smpl/smpl_torch.py

@@ -0,0 +1,222 @@
+# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py
+
+import os
+import sys
+import pickle
+import torch
+import numpy as np
+import scipy.io as sio
+from time import time
+from torch.nn import Module
+from utils import plot_mesh
+
+class SMPL(Module):
+  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
+    
+    super(SMPL, self).__init__()
+    self.J_regressor = []
+    self.v_template = []
+    self.shapedirs = []
+    self.posedirs = []
+    self.weights = []
+    self.is_skirt = is_skirt
+    # Load SMPL models
+    # path = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(female['weights'], dtype=torch.float32))
+    with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(male['weights'], dtype=torch.float32))
+        kintree_table = male['kintree_table']
+        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
+        self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
+    self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device)
+    self.v_template = torch.stack(self.v_template, axis=0).to(device)
+    self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device)
+    self.posedirs = torch.stack(self.posedirs, axis=0).to(device)
+    self.weights = torch.stack(self.weights, axis=0).to(device)
+
+    # Topology no head/hands/feet
+    F = sio.loadmat(os.path.join(path, 'ssmpl_faces.mat'))
+    self.f = np.array(F['faces_default'])
+    self.f_skirt = np.array(F['faces_skirt'])
+    self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat')))
+    self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
+
+    self.weights_all = self.weights
+    # self.weights = self.weights[:,self.to_keep,:]
+    
+    self.shape = (1,27578,3)
+    self.device = device
+
+    if orig_smpl:
+        self.to_keep = self.to_keep[self.to_keep < 6890]
+        self.J_regressor = self.J_regressor[:,:,:6890]
+        self.v_template = self.v_template[:,:6890,:]
+        self.shapedirs = self.shapedirs[:,:6890,:,:]
+        self.posedirs = self.posedirs[:,:6890,:,:]
+        self.weights_all = self.weights_all[:,:6890,:]
+        self.weights = self.weights_all[:,self.to_keep,:]
+        self.shape = (1, 6890, 3)
+
+  @staticmethod
+  def rodrigues(r):
+    """
+    Rodrigues' rotation formula that turns axis-angle tensor into rotation
+    matrix in a batch-ed manner.
+    Parameter:
+    ----------
+    r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].
+    Return:
+    -------
+    Rotation matrix of shape [batch_size * angle_num, 3, 3].
+    """
+    # eps = r.clone().normal_(std=1e-8)
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
+    theta_dim = theta.shape[0]
+    r_hat = r / theta
+    # cos = torch.cos(theta)
+    z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
+    m = torch.stack(
+      (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+       -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
+    m = torch.reshape(m, (-1, 3, 3))
+    i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
+             + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
+    A = r_hat.permute(0, 2, 1)
+    # dot = torch.matmul(A, r_hat)
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
+    return R
+
+  @staticmethod
+  def with_zeros(x):
+    """
+    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.
+    Parameter:
+    ---------
+    x: Tensor to be appended.
+    Return:
+    ------
+    Tensor after appending of shape [4,4]
+    """
+    ones = torch.tensor(
+      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+    ).expand(x.shape[0],-1,-1).to(x.device)
+    ret = torch.cat((x, ones), dim=1)
+    return ret
+
+  @staticmethod
+  def pack(x):
+    """
+    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.
+    Parameter:
+    ----------
+    x: A tensor of shape [batch_size, 4, 1]
+    Return:
+    ------
+    A tensor of shape [batch_size, 4, 4] after appending.
+    """
+    zeros43 = torch.zeros(
+      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device)
+    ret = torch.cat((zeros43, x), dim=3)
+    return ret
+
+
+  def forward(self, betas, pose, remove=True):
+    
+    """
+          Construct a compute graph that takes in parameters and outputs a tensor as
+          model vertices. Face indices are also returned as a numpy ndarray.
+          
+          20190128: Add batch support.
+          Parameters:
+          ---------
+          pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation
+          relative to parent joint. For root joint it's global orientation.
+          Represented in a axis-angle format.
+          betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of
+          PCA components. Only 10 components were released by SMPL author.
+          trans: Global translation tensor of shape [N, 3].
+          Return:
+          ------
+          A 3-D tensor of [N * 6890 * 3] for vertices,
+          and the corresponding [N * 19 * 3] joint positions.
+    """
+    batch_num = betas.shape[0]
+
+    # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
+    v_shaped = []
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped = torch.stack(v_shaped, axis=-1)
+    
+    # J = torch.matmul(self.J_regressor, v_shaped)
+    J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    R_cube_big = self.rodrigues(pose.view(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
+
+    R_cube = R_cube_big[:, 1:, :, :]
+    I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
+      torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
+    lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
+    v_posed = torch.stack([v_shaped[i, :, :, g.to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[g.to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
+
+    results = []
+    results.append(
+      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+    )
+    for i in range(1, 24):
+      results.append(
+        torch.matmul(
+          results[self.parent[i]],
+          self.with_zeros(
+            torch.cat(
+              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              dim=2
+            )
+          )
+        )
+      )
+    
+    stacked = torch.stack(results, dim=1)
+    results = stacked - \
+      self.pack(
+        torch.matmul(
+          stacked,
+          torch.reshape(
+            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
+            (batch_num, 24, 4, 1)
+          )
+        )
+      )
+    
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+            
+    T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    rest_shape_h = torch.cat(
+      (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
+    )
+    v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
+    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    result = v
+
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
+
+    if remove:
+      result = result[:,self.to_keep,:]
+
+    return result, joints