luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

add random view option (definition changed later)

Browse Source
This commit is contained in:
luni
2023-01-14 17:12:37 +00:00
parent e1d7cafeab
commit 7d9610b1f8
1 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+9 -7
View File
@@ -24,10 +24,11 @@ from cloth3d.util import quads2tris
class Dataset(Dataset): class Dataset(Dataset):
def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True): def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True):
self.folder = folder self.folder = folder
self.folder_rgb = folder_rgb self.folder_rgb = folder_rgb
self.device = device self.device = device
self.random_view = random_view
self.transforms = tv.transforms.ToTensor() self.transforms = tv.transforms.ToTensor()
self.frames = read_frames_file(os.path.join(folder, frames_file)) self.frames = read_frames_file(os.path.join(folder, frames_file))
@@ -52,8 +53,9 @@ class Dataset(Dataset):
path = os.path.join(self.folder, frame.id) path = os.path.join(self.folder, frame.id)
info = self.reader.read_info(frame.id) info = self.reader.read_info(frame.id)
zrot = info['zrot'] if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
# mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True) # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot']) mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
# mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False) # 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')) mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
@@ -61,7 +63,7 @@ class Dataset(Dataset):
faces = mask_faces(self.body_faces, mesh_mask) faces = mask_faces(self.body_faces, mesh_mask)
# faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment)) # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
# body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True) # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot']) body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
# body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False) # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
body_faces = quads2tris(body_faces) body_faces = quads2tris(body_faces)
@@ -190,19 +192,19 @@ class Dataset(Dataset):
# v2i_coverage[v2i[:,0], v2i[:, 1]] = 1 # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
# ex.imshow(v2i_coverage).show() # 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'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32) rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32)
view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device) view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
# image = self.transforms(image) # image = self.transforms(image)
data = { 'mesh_posed': mesh_posed, data = { 'mesh_posed': torch.tensor(mesh_posed, device=self.device),
'depth_image': depth_image, 'depth_image': depth_image,
'mesh_mask': mesh_mask, 'mesh_mask': mesh_mask,
'depth_mask': depth_mask, 'depth_mask': depth_mask,
'normal_image': normal_image, 'normal_image': normal_image,
'depth_mask_garment': depth_mask_garment, 'depth_mask_garment': depth_mask_garment,
'depth_mask_body': depth_mask_body, 'depth_mask_body': depth_mask_body,
'rgb_image': rgb_image, 'rgb_image': rgb_image.to(self.device),
'view': view, 'view': view,
} }