save progress

data/cloth3d_io.py

@@ -1,4 +1,5 @@
 import os
+import scipy
 import numpy as np
 import scipy.io as sio
 from struct import pack, unpack
@@ -222,16 +223,20 @@ def _todict(matobj):
 
 
 # can i cache map?
-def read_posed_garment_frame(folder, id, garment, frame, mask=None, shape=None, shapedirs=None, unshape=False):
+def read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False):
     pc16_path = os.path.join(folder, id, garment + '.pc16')
-    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
     verts = readPC2Frame(pc16_path, frame, True)
-    verts = verts[:, map, :]
+    verts = verts[map, :]
 
     if unshape:
         shape_offsets = np.tensordot(shape, shapedirs, axes=[[1], [2]])
         verts[mask, :] = verts[mask, :] - np.squeeze(shape_offsets)[mask, :]
 
+    if deorient:
+        R = scipy.spatial.transform.Rotation.from_euler(pose[:3])
+        # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
     return verts
 
 
@@ -241,7 +246,6 @@ def read_garment_mask(folder, id, garment):
     return mask
 
 
-# can i cache mask?
 def read_garment_template(folder, id, garment, mask):
     offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
     body_static = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B']
@@ -249,6 +253,16 @@ def read_garment_template(folder, id, garment, mask):
     return verts_static
 
 
+def read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    if deorient:
+        R = scipy.spatial.transform.Rotation.from_euler(pose[:3])
+        # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+        
+    return verts
 
 # import os
 # import sys

data/dataset.py

@@ -11,11 +11,11 @@ import scipy.io as sio
 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
+from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame
 
 
 class StaticDataset():
-    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False) -> None:
+    def __init__(self, folder, signature, resolutions=[512], in_memory=True, is_cape=False, use_no_orient=True) -> None:
         self.store = dict()
         self.signature = signature
         self.folder = folder
@@ -30,6 +30,7 @@ class StaticDataset():
         self.resolutions = resolutions
         self.filled = ''
         self.is_cape = is_cape
+        self.use_no_orient = use_no_orient
 
         if is_cape:
             self.uv_masks = dict()
@@ -49,6 +50,9 @@ class StaticDataset():
         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)
@@ -91,6 +95,13 @@ class StaticDataset():
         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)
         return batch
 
@@ -186,31 +197,24 @@ class DatasetFactory():
             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()
 
+            # redunda ntly loading info again
+            pose = loadInfo(os.path.join(folder, id))['poses'][:,frame_nr]
+            # !!!!!!!!!!!!!!!!!!!
+            map = None
+
             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 'pose_shape' in signature.keys():
+
-                pose_shape = np.load(os.path.join(folder, 'pose_shape',  id + separator + frame_nr + '.npy'))
-                if use_no_orient:
-                    pose_shape[:3] = 0
-                item['pose_shape'] = pose_shape
             if 'mesh_body_posed' in signature.keys():
                 if use_no_orient:
-                    mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
+                    item['mesh_body_posed'] = read_posed_body_frame(folder, id, frame_nr, pose, deorient=use_no_orient)
-                    item['mesh_body_posed'] = mesh_posed_body_no_orient
-                else:
-                    mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy'))
-                    item['mesh_body_posed'] = mesh_body_posed
             if 'mesh_posed' in signature.keys():
                 if use_no_orient:
-                    mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32)
+                    item['mesh_posed'] = read_posed_garment_frame(folder, id, garment, frame, map, mask=None, shape=None, shapedirs=None, pose=None, unshape=False, deorient=False)
-                    item['mesh_posed'] = mesh_posed
-                else:
-                    mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32)
-                    item['mesh_posed'] = mesh_posed
             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