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:

data/preprocess.py

@@ -7,10 +7,11 @@ import scipy.io as sio
 from tqdm import tqdm
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from utils.utils import mask_faces
+from utils.data_utils import mask_faces, Frame
 from superSMPL.IO import readPC2, writePC2
 
-def generate_faces(folder):
+
+def generate_faces(folder, frames):
     frames = os.listdir(os.path.join(folder, 'mesh_mask'))
     F = np.load('custom_data/smpl_faces.npy')
     os.makedirs(os.path.join(folder, 'faces'), exist_ok=True)
@@ -21,12 +22,6 @@ def generate_faces(folder):
         np.save(os.path.join(folder, 'faces', frame), faces)
 
 
-class Frame:
-    def __init__(self, id, frame_nr):
-        self.id = id
-        self.frame_nr = frame_nr
-
-
 def generate_frames_file(folder, file_name):
     ids = os.listdir(folder)#[:1]
     frames = []
@@ -40,11 +35,11 @@ def generate_frames_file(folder, file_name):
         garment = garments[0]
         verts = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V']
         for i in range(verts.shape[0]):
-            frames.append(Frame(id, i))
+            frames.append(Frame(id, i, garment))
 
     with open(os.path.join(folder, file_name), 'w') as f:
         for i, frame in enumerate(frames):
-            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + '\n')
+            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
 
 
 def genereated_registered_and_mask(folder, is_skirt = False):