updated dataset

data/dataset.py

@@ -21,6 +21,7 @@ from torchvision.transforms.functional import InterpolationMode
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
 from data.utils import Frame
+from data.cloth3d_io import readPC2, readPC2Frame
 
 
 def read_frames(folder, shuffle=True, is_cape=False, testing=False, frames_file=None):
@@ -60,7 +61,6 @@ class StaticCache():
         self.folder = folder
         assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
         self.resolutions = resolutions
-        self.filled = ''
         self.is_cape = is_cape
         self.device = device
         
@@ -74,50 +74,65 @@ class StaticCache():
                 self.uv_masks[res] = down.cpu().numpy()
 
 
-    def read(self, id):
+    def read(self, id, garment):
         data = dict()
 
+        mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', f'{id}_{garment}.npy')); mesh_mask = mesh_mask.astype(bool)
+        data['mesh_mask'] = mesh_mask
+
         for res in self.resolutions:
             res_string = ''
             if res != 512:
                 res_string = '_' + str(res)
 
             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
+                uv_static = np.load(os.path.join(self.folder, f'uv_static{res_string}', f'{id}_{garment}.npz')); uv_static = uv_static['uv_static']; data['uv_static' + res_string] = uv_static
             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
+                    uv_mask = np.load(os.path.join(self.folder, f'uv_mask{res_string}', f'{id}_{garment}.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:
-            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)
+            mesh_static = np.load(os.path.join(self.folder, 'mesh_static', f'{id}_{garment}.npy'))
-        if 'mesh_mask' in self.signature:
+            data['mesh_static'] = np.zeros((14475, 3), dtype=np.float32)
-            if self.is_cape:
+            data['mesh_static'][mesh_mask] = mesh_static.astype(np.float32)
-                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
 
 
-    def __getitem__(self, id) -> dict:
+    def __getitem__(self, frame) -> dict:
-        if id not in self.store:
+        if frame not in self.store:
-            item = self.read(id)
+            item = self.read(frame.id, frame.garment)
             stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()}
             for key, value in stored_data.items():
                 if key.startswith('uv_') and value.shape[-1] == 3:
                     stored_data[key] = torch.permute(value, [2,0,1])
-            self.store[id] = stored_data
+            self.store[frame] = stored_data
         
-        data = { key: value.to(self.device) for key, value in self.store[id].items()}
+        data = { key: value.to(self.device) for key, value in self.store[frame].items()}
         return data
 
 
+
+def read_frames_file(frames_file, shuffle=False):
+    frames = []
+    with open(frames_file, 'r') as f:
+        lines = f.readlines()
+        lines = [line.strip() for line in lines]
+    frames = [Frame(line.split()[0], int(line.split()[1]), line.split()[2]) for line in lines]
+
+    if shuffle:
+        random.shuffle(frames)
+
+    return frames
+
+
 class Dataset(Dataset):
-    def __init__(self, folder, signature, resolutions=[512], frames_file='frames_all.txt', device='cuda', use_cache=False, is_cape=False):
+    def __init__(self, folder, signature, resolutions=[512], frames_file='', device='cuda', use_cache=False, is_cape=False):
         self.folder = folder
         self.device = device
         self.is_cape = is_cape
@@ -125,13 +140,20 @@ class Dataset(Dataset):
         self.resolutions = resolutions
 
         self.transforms = tv.transforms.ToTensor()
-        self.frames = read_frames(folder, frames_file=frames_file, is_cape=is_cape, shuffle=True)
+        self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True)
 
-        self.body_faces = np.load('smpl/smpl_faces.npy')
         self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
 
         self.cache = dict() if use_cache else None
 
+        ids = set([frame.id for frame in self.frames])
+        self.pose_cache = dict()
+        for id in ids:
+            self.pose_cache[id] = np.load(os.path.join(self.folder, 'pose', f'{id}.npy'))
+        self.shape_cache = dict()
+        for id in ids:
+            self.shape_cache[id] = np.load(os.path.join(self.folder, 'shape', f'{id}.npy'))
+
 
     def __len__(self):
         return len(self.frames)
@@ -145,73 +167,56 @@ class Dataset(Dataset):
                 data = { **data, **self.static_cache[frame.id] }
                 return data
 
-        separator = '_'
-        if self.is_cape:
-            separator = '.'
-
-        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'
-
         item = dict()
+        mesh_mask = self.static_cache[frame]['mesh_mask'].cpu().numpy().astype(bool)
+
+        item['id'] = frame.id
+        item['frame_nr'] = frame.frame_nr
+        item['garment'] = frame.garment
 
-        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')).astype(np.float32)
+            pose = self.pose_cache[frame.id][frame.frame_nr]
-            pose_shape[:3] = 0
+            pose[:3] = 0
-            item['pose_shape'] = pose_shape
+            item['pose_shape'] = np.concatenate((pose, self.shape_cache[frame.id]), axis=0)
         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']
+            mesh_body_posed_no_orient = readPC2Frame(os.path.join(self.folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
-            item['mesh_body_posed'] = mesh_posed_body_no_orient
+            item['mesh_body_posed'] = mesh_body_posed_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)
+            mesh_posed = readPC2Frame(os.path.join(self.folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
-            item['mesh_posed'] = mesh_posed
+            item['mesh_posed'] = np.zeros((14475, 3), dtype=np.float32)
+            item['mesh_posed'][mesh_mask] = 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']
+            mesh_unposed_no_shape = readPC2Frame(os.path.join(self.folder, 'mesh_unposed_no_shape', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
-            item['mesh_unposed'] = mesh_unposed_no_shape
+            item['mesh_unposed'] = np.zeros((14475, 3), dtype=np.float32)
-        if 'mesh_curvature' in self.signature:
+            item['mesh_unposed'][mesh_mask] = mesh_unposed_no_shape
-            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)
         if 'mesh_normals' in self.signature:
-            mesh_normals = np.load(os.path.join(self.folder, 'mesh_normals', frame.id + separator + frame.frame_nr + '.npz')); mesh_normals = mesh_normals['mesh_normals']
+            mesh_normals = readPC2Frame(os.path.join(self.folder, 'mesh_normals', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True)
-            item['mesh_normals'] = mesh_normals
+            item['mesh_normals'] = np.zeros((14475, 3), dtype=np.float32)
+            item['mesh_normals'][mesh_mask] = mesh_normals
 
         for res in self.resolutions:
             res_string = ''
             if res != 512:
                 res_string = '_' + str(res)
 
-            if 'uv_curvature' + res_string in self.signature:
+            if 'uv_normals' + 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')); 
+                uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_normals = uv_normals['uv_normals']
-                if res_string != '':
+                item['uv_normals' + res_string] = uv_normals
-                    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).float().to(self.device)
-                item['uv_curvature' + res_string] = uv_curvature
 
             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']
+                uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed']
                 item['uv_unposed' + res_string] = uv_unposed_no_shape
 
             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_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed']
-                uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
+                item['uv_body_posed' + res_string] = uv_body_posed_no_orient
-                item['uv_body_posed' + res_string] = uv_posed_body_no_orient
 
             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']
+                uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz')); uv_posed = uv_posed['uv_posed']
                 item['uv_posed' + res_string] = uv_posed
 
-            if 'uv_normals' + res_string in self.signature:
-                uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
-                item['uv_normals' + res_string] = uv_normals
 
-
+        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor or type(value) == int) else value) for key, value in item.items()}
-        data = { key: (torch.from_numpy(value) if not (type(value) == str or type(value) == torch.Tensor) else value) for key, value in item.items()}
 
         for key, value in data.items():
             if key.startswith('uv_') and (value.shape[-1] == 3 or value.shape[-1] == 7):
@@ -221,38 +226,93 @@ class Dataset(Dataset):
             self.cache[frame] = data
         
         data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
-        data = { **data, **self.static_cache[frame.id] }
+        data = { **data, **self.static_cache[frame] }
         return data
 
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader
+    import plotly.graph_objects as go
+    import plotly.express as ex
+    from plotly.subplots import make_subplots
+    from custom_utils.custom_utils import _maskFaces
+
+    F = np.load('smpl/smpl_faces.npy')
+    F_skirt = np.load('smpl/smpl_faces_skirt.npy')
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    signature = { 'id', 
+    signature = { 'pose_shape',
-                  'frame_nr',
-                  'pose_shape',
                   'uv_posed',
                   'uv_body_posed',
                   'mesh_body_posed',
                   'mesh_posed',
                   'mesh_unposed',
                   'uv_unposed',
-                  'uv_curvature',
+                  'uv_normals',
-                  'mesh_curvature',
+                  'mesh_normals',
                   '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')
+    dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader:
         batch
+        b = 0
+        mesh_mask = batch['mesh_mask'][b].cpu().numpy()
+        faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
+        _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
+
+        mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask]
+        mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask]
+        mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask]
+        mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy()
+
+        fig = make_subplots(rows=2, cols=2,
+                            specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}],
+                                [{'type': 'mesh3d'}, {'type': 'mesh3d'}]])
+
+        fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1)
+        fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2)
+        fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=1)
+        fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], opacity=1), row=2, col=2)
+
+        fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
+                                    yaxis=dict(range=[-1, 1], autorange=False),
+                                    zaxis=dict(range=[-1, 1], autorange=False),
+                                    aspectmode='cube'),
+                        scene2=dict(xaxis=dict(range=[-1, 1], autorange=False),
+                                    yaxis=dict(range=[-1, 1], autorange=False),
+                                    zaxis=dict(range=[-1, 1], autorange=False),
+                                    aspectmode='cube'),
+                        scene3=dict(xaxis=dict(range=[-1, 1], autorange=False),
+                                    yaxis=dict(range=[-1, 1], autorange=False),
+                                    zaxis=dict(range=[-1, 1], autorange=False),
+                                    aspectmode='cube'),
+                        scene4=dict(xaxis=dict(range=[-1, 1], autorange=False),
+                                    yaxis=dict(range=[-1, 1], autorange=False),
+                                    zaxis=dict(range=[-1, 1], autorange=False),
+                                    aspectmode='cube'))
+
+        fig.show()
+
+
+        uv_mask = batch['uv_mask'][b].cpu().numpy()
+        uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0)
+        uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2
+        uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2
+        uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2
+        uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2
+        uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3
+
+        img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2)
+        img = np.transpose(img, (1, 2, 0))
+        ex.imshow(img).show()
+
         pass
 
     pass 
@@ -281,7 +341,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()
+#         ex.imshow(batch['uv_normals'][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)
 

data/preprocessing.py

@@ -510,6 +510,62 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
             # fig.show()
             pass
 
+
+def create_frames_file(folder, out_folder, split=0.9):
+    ids = os.listdir(folder)
+
+    frames_non_skirt = []
+    frames_skirt = []
+    ids_non_skirt = set()
+    ids_skirt = set()
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        frames = info['poses'].shape[1]
+
+        assert info['poses'].shape[0] == 72
+
+        for garment in garments:
+            for i in range(frames):
+                if garment == 'Skirt' or garment == 'Dress':
+                    frames_skirt.append(Frame(id, i, garment))
+                    ids_skirt.add(id)
+                else:
+                    frames_non_skirt.append(Frame(id, i, garment))
+                    ids_non_skirt.add(id)
+
+    nr_train_non_skirt = int(len(ids_non_skirt) * split)
+    nr_train_skirt = int(len(ids_skirt) * split)
+
+    ids_skirt = list(ids_skirt)
+    ids_non_skirt = list(ids_non_skirt)
+
+    with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt[:nr_train_non_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt[:nr_train_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt[nr_train_non_skirt:]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt[nr_train_skirt:]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+
 if __name__ == '__main__':
 
 
@@ -530,4 +586,6 @@ if __name__ == '__main__':
     # create_mesh_static(folder, out_folder)
     # create_mesh_body_posed_no_orient(folder, out_folder)
     # create_mesh_normals(folder, out_folder)
-    create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
+    # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
+
+    create_frames_file(folder, out_folder)

data/utils.py

@@ -1,7 +1,7 @@
 import tables as tb
 
 class Frame:
-    def __init__(self, id, frame_nr, garment=''):
+    def __init__(self, id: str, frame_nr: int, garment: str=''):
         self.id = id
         self.frame_nr = frame_nr
         self.garment = garment