From 304a66e5b7a084b80738d3a47202821a56b597ce Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 6 Jan 2023 11:46:57 +0000 Subject: [PATCH] read rgb image as well --- data/dataset.py | 19 ++++++++++--------- test.py | 5 ++--- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index 0c3fe84..20308ed 100644 --- a/data/dataset.py +++ b/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: diff --git a/test.py b/test.py index 7cc64a1..aa40e95 100644 --- a/test.py +++ b/test.py @@ -57,9 +57,8 @@ def batch_to_device(batch, device): if __name__ == '__main__': device = "cuda" if torch.cuda.is_available() else "cpu" - - train_dataset = Dataset('/home/workspace/local-data/train') - valid_dataset = Dataset('/home/workspace/local-data/valid') + train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt') + valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt') train_dataloader = DataLoader(train_dataset, batch_size=4, drop_last=True, num_workers=2) valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2) train_accumulator = Accumulator()