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

add new view definition to dataset

Browse Source
This commit is contained in:
luni
2023-01-14 18:25:01 +00:00
parent 35928ccbff
commit 581a5b2af2
1 changed files with 24 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+24 -5
View File
@@ -7,6 +7,7 @@ import math
import time
import torch
import kornia
import pytorch3d
import torchvision as tv
import numpy as np
import scipy.io as sio
@@ -53,9 +54,12 @@ class Dataset(Dataset):
path = os.path.join(self.folder, 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]
zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
trans = info['trans'][:,frame.frame_nr]
# 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], zrot)
mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
mesh_posed = torch.tensor(mesh_posed, device=self.device)
# 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'))
@@ -63,14 +67,28 @@ class Dataset(Dataset):
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)
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, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, trans, zrot)
# body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
body_faces = quads2tris(body_faces)
if self.random_view:
x_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
y_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
z_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi)
view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((x_angle, y_angle, z_angle)))
transform = pytorch3d.transforms.Rotate(view)
mesh_posed = transform.transform_points(mesh_posed)
body = transform.transform_points(body)
else:
view = torch.eye(3)
combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device)
# 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(self.device)
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]
@@ -194,10 +212,11 @@ class Dataset(Dataset):
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([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), trans.astype(np.float32)), axis=0), device=self.device)
# image = self.transforms(image)
data = { 'mesh_posed': torch.tensor(mesh_posed, device=self.device),
data = { 'mesh_posed': mesh_posed,
'depth_image': depth_image,
'mesh_mask': mesh_mask,
'depth_mask': depth_mask,