read rgb image as well

data/dataset.py

@@ -24,8 +24,9 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, frames_file='frames_all.txt'):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt'):
         self.folder = folder
+        self.folder_rgb = folder_rgb
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
 
@@ -49,7 +50,7 @@ class Dataset(Dataset):
         # mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False)
 
         mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
-        # TODO: what's the difference between masking and reading .obj file
+        # TODO: what's the difference between masking and reading .obj file. is the obj file the not registered version?
         faces = mask_faces(self.body_faces, mesh_mask)
         # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
         # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
@@ -61,11 +62,6 @@ class Dataset(Dataset):
         # combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda')
         combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda')
 
-        idx_frame = 35500
-        x = combined[idx_frame, :].cpu().detach().numpy()
-        y = np.array(info['camLoc'], dtype=np.float32)
-        dist = np.sqrt(np.sum((x - y) ** 2))
-
         start = time.perf_counter()
         # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
         # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
@@ -187,6 +183,10 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # ex.imshow(v2i_coverage).show()
 
+
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'))
+
+
         # image = self.transforms(image)
         return { 'mesh_posed': mesh_posed,
                  'depth_image': depth_image, 
@@ -194,7 +194,8 @@ class Dataset(Dataset):
                  'depth_mask': depth_mask, 
                  'normal_image': normal_image,
                  'depth_mask_garment': depth_mask_garment,
-                 'depth_mask_body': depth_mask_body
+                 'depth_mask_body': depth_mask_body,
+                 'rgb_image': rgb_image,
                 }
 
 
@@ -204,7 +205,7 @@ if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     # dataset = Dataset('/home/workspace/local-data/train')
-    dataset = Dataset('dataset/CLOTH3D/train')
+    dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader: