use raw cloth3d input

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 
 import math
 import time
@@ -15,12 +15,16 @@ from torch.utils.data import Dataset
 
 from render import Render, DepthRenderer
 
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from superSMPL.IO import readPC2Frame, writePC2
+from utils.data_utils import Frame, read_frames_file, mask_faces
+
 
 class Dataset(Dataset):
-    def __init__(self, folder):
+    def __init__(self, folder, frames_file='frames_all.txt'):
         self.folder = folder
         self.transforms = tv.transforms.ToTensor()
-        self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
+        self.frames = read_frames_file(os.path.join(folder, frames_file))
 
         self.render = DepthRenderer()
         self.body_faces = np.load('custom_data/smpl_faces.npy')
@@ -31,16 +35,15 @@ class Dataset(Dataset):
 
 
     def __getitem__(self, idx):
-        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
-        mesh_posed = np.load(path)['mesh_posed']
+        frame : Frame = self.frames[idx]
+        path = os.path.join(self.folder, frame.id)
+        mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
 
-        id = self.frames[idx].split('_')[0]
-        faces = np.load(os.path.join(self.folder, 'faces', id + '.npy'))
-        mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy'))
-        body = np.load(os.path.join(self.folder, 'mesh_body_posed', self.frames[idx][:-1] + 'y'))
+        mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
+        faces = mask_faces(self.body_faces, mask)
+        body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
 
-        info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
-        info = sio.loadmat(info_folder)
+        info = sio.loadmat(os.path.join(path, 'info.mat'))
         combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda')
         combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda')
 
@@ -60,7 +63,8 @@ if __name__ == '__main__':
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    dataset = Dataset('/home/workspace/local-data/train')
+    # dataset = Dataset('/home/workspace/local-data/train')
+    dataset = Dataset('dataset/CLOTH3D/train')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True)
 
     for batch in train_dataloader: