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

pose autoencoder pipeline

Browse Source
This commit is contained in:
luni
2023-06-16 22:11:57 +00:00
parent 5f34b667fc
commit 7a31f0d04c
4 changed files with 405 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/pose_dataset.py
+39 -22
View File
@@ -17,10 +17,11 @@ from data.pose_renderer import PoseRenderer
import trimesh
class Dataset(Dataset):
def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
class PoseDataset(Dataset):
def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True, ae_mode=False):
self.folder = folder
self.device = device
self.ae_mode = ae_mode
self.frames = read_frames_file(os.path.join(folder, frames_file))
self.reader = DataReader(folder, 'superSMPL')
@@ -49,7 +50,7 @@ class Dataset(Dataset):
# and rotate it 90 degress along the z axis
body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.get_device()))
out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device))
return out
@@ -68,57 +69,73 @@ class Dataset(Dataset):
shape = torch.tensor(self.shapes[frame.id], device=self.device)
pose[:3] = 0
pose_out = pose.clone()
pose_out[1] = np.random.uniform(-np.pi, np.pi)
angle = torch.rand(1) * 2 * np.pi - np.pi
matrix = self.smpl.rodrigues(torch.tensor([0.0, angle.numpy()[0], 0.0], device=self.device)[None, None, :]).to(self.device)[0]
angle = angle.to(self.device)
pose_out[1] = angle
# pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
# generate random number between -np.pi, np.pi
# start = time()
body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
if self.ae_mode:
body = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
else:
body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False)
mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0]))
body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
body_back[0,:,0] = -body_back[0,:,0]
# mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False)
# mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0]))
# body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
# body_back[0,:,0] = -body_back[0,:,0]
# print(f'smpl took {time() - start} seconds')
# start = time()
render = self._render_body(body[0])[0,:,:,:3]
render_back = self._render_body(body_back)[0,:,:,:3]
render_out = self._render_body(body_out[0])[0,:,:,:3]
# render_back = self._render_body(body_back)[0,:,:,:3]
if not self.ae_mode:
render_out = self._render_body(body_out[0])[0,:,:,:3]
render = render.permute(2,0,1)
render_back = render_back.permute(2,0,1)
render_out = render_out.permute(2,0,1)
# render_back = render_back.permute(2,0,1)
if not self.ae_mode:
render_out = render_out.permute(2,0,1)
# import plotly.express as ex
# ex.imshow(render[0,:,:,:3].cpu().numpy()).show()
# ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show()
mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32)
mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)
# mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
if not self.ae_mode:
mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)
render = render * mask
render_back = render_back * mask_back
render_out = render_out * mask_out
# render_back = render_back * mask_back
if not self.ae_mode:
render_out = render_out * mask_out
# print(f'render took {time() - start} seconds')
data = {
'pose': render,
'shape': shape,
'pose_out': render_out,
'pose_back': render_back,
# 'pose_back': render_back,
# 'body': body,
# 'body_out': body_out,
# 'render': render,
# 'render_out': render_out,
'mask': mask,
'mask_out': mask_out,
'angle': angle,
'matrix': matrix,
}
if not self.ae_mode:
data['pose_out'] = render_out
data['mask_out'] = mask_out
if self.cache is not None:
self.cache[frame] = data
@@ -129,7 +146,7 @@ if __name__ == '__main__':
from torch.utils.data import DataLoader
device = "cuda" if torch.cuda.is_available() else "cpu"
dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
for batch in train_dataloader: