From 51b62190215b64a7f40a95276699af7ae72ae08f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Dec 2022 12:57:08 +0000 Subject: [PATCH] got root joint location --- cloth3d/read.py | 310 ++++++++++++++++++++++++------------------------ data/dataset.py | 61 ++++++---- data/render.py | 2 +- 3 files changed, 192 insertions(+), 181 deletions(-) diff --git a/cloth3d/read.py b/cloth3d/read.py index 2611ab1..59f75bd 100644 --- a/cloth3d/read.py +++ b/cloth3d/read.py @@ -11,163 +11,163 @@ from superSMPL.IO import readOBJ, readPC2Frame class DataReader: - def __init__(self, folder, smpl_path): - # Data paths - self.SRC = folder - # SMPL model - smpl_path = smpl_path - self.smpl = { - 'f': SMPLModel(os.path.join(smpl_path, 'super_model_f.pkl')), - 'm': SMPLModel(os.path.join(smpl_path, 'super_model_m.pkl')) - } - - """ - Read sample info - Input: - - sample: name of the sample e.g.:'01_01_s0' - """ - def read_info(self, sample): - info_path = os.path.join(self.SRC, sample, 'info') - return loadInfo(info_path) - - """ Human data """ - """ - Read SMPL parameters for the specified sample and frame - Inputs: - - sample: name of the sample - - frame: frame number - """ - def read_smpl_params(self, sample, frame): - # Read sample data - info = self.read_info(sample) - # SMPL parameters - gender = 'm' if info['gender'] else 'f' - pose = info['poses'][:, frame].reshape(self.smpl[gender].pose_shape) - shape = info['shape'] - trans = info['trans'][:, frame].reshape(self.smpl[gender].trans_shape) - return gender, pose, shape, trans - - """ - Computes human mesh for the specified sample and frame - Inputs: - - sample: name of the sample - - frame: frame number - - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint - Outputs: - - V: human mesh vertices - - F: mesh faces - """ - def read_human(self, sample, frame, absolute=True): - # Read sample data - info = self.read_info(sample) - # SMPL parameters - gender, pose, shape, trans = self.read_smpl_params(sample, frame) - # Compute SMPL - V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans if absolute else None) - V -= J[0:1] - # Apply rotation on z-axis - zRot = zRotMatrix(info['zrot']) - return zRot.dot(V.T).T, self.smpl[gender].faces - - """ Garment data """ - """ - Reads garment vertices location for the specified sample, garment and frame - Inputs: - - sample: name of the sample - - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) - - frame: frame number - - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint - Outputs: - - V: 3D vertex locations for the specified sample, garment and frame - """ - def read_garment_vertices(self, sample, garment, frame, absolute=True): - # Read garment vertices (relative to root joint) - pc16_path = os.path.join(self.SRC, sample, garment + '_registered.pc16') - V = readPC2Frame(pc16_path, frame, True) - if absolute: - # Read sample data - info = self.read_info(sample) - # Transform to absolute - V += info['trans'][:,frame][None] - # Apply rotation on z-axis - zRot = zRotMatrix(info['zrot']) - return zRot.dot(V.T).T + def __init__(self, folder, smpl_path): + # Data paths + self.SRC = folder + # SMPL model + smpl_path = smpl_path + self.smpl = { + 'f': SMPLModel(os.path.join(smpl_path, 'super_model_f.pkl')), + 'm': SMPLModel(os.path.join(smpl_path, 'super_model_m.pkl')) + } + + """ + Read sample info + Input: + - sample: name of the sample e.g.:'01_01_s0' + """ + def read_info(self, sample): + info_path = os.path.join(self.SRC, sample, 'info') + return loadInfo(info_path) + + """ Human data """ + """ + Read SMPL parameters for the specified sample and frame + Inputs: + - sample: name of the sample + - frame: frame number + """ + def read_smpl_params(self, sample, frame): + # Read sample data + info = self.read_info(sample) + # SMPL parameters + gender = 'm' if info['gender'] else 'f' + pose = info['poses'][:, frame].reshape(self.smpl[gender].pose_shape) + shape = info['shape'] + trans = info['trans'][:, frame].reshape(self.smpl[gender].trans_shape) + return gender, pose, shape, trans + + """ + Computes human mesh for the specified sample and frame + Inputs: + - sample: name of the sample + - frame: frame number + - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint + Outputs: + - V: human mesh vertices + - F: mesh faces + """ + def read_human(self, sample, frame, absolute=True): + # Read sample data + info = self.read_info(sample) + # SMPL parameters + gender, pose, shape, trans = self.read_smpl_params(sample, frame) + # Compute SMPL + V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans if absolute else None) + V -= J[0:1] + # Apply rotation on z-axis + zRot = zRotMatrix(info['zrot']) + return zRot.dot(V.T).T, self.smpl[gender].faces, zRot.dot(trans.T).T[None,:] + + """ Garment data """ + """ + Reads garment vertices location for the specified sample, garment and frame + Inputs: + - sample: name of the sample + - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) + - frame: frame number + - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint + Outputs: + - V: 3D vertex locations for the specified sample, garment and frame + """ + def read_garment_vertices(self, sample, garment, frame, absolute=True): + # Read garment vertices (relative to root joint) + pc16_path = os.path.join(self.SRC, sample, garment + '_registered.pc16') + V = readPC2Frame(pc16_path, frame, True) + if absolute: + # Read sample data + info = self.read_info(sample) + # Transform to absolute + V += info['trans'][:,frame][None] + # Apply rotation on z-axis + zRot = zRotMatrix(info['zrot']) + return zRot.dot(V.T).T - """ - Reads garment faces for the specified sample and garment - Inputs: - - sample: name of the sample - - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) - Outputs: - - F: mesh faces - """ - def read_garment_topology(self, sample, garment): - # Read OBJ file - obj_path = os.path.join(self.SRC, sample, garment + '.obj') - return readOBJ(obj_path)[1] + """ + Reads garment faces for the specified sample and garment + Inputs: + - sample: name of the sample + - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) + Outputs: + - F: mesh faces + """ + def read_garment_topology(self, sample, garment): + # Read OBJ file + obj_path = os.path.join(self.SRC, sample, garment + '.obj') + return readOBJ(obj_path)[1] - """ - Reads garment UV map for the specified sample and garment - Inputs: - - sample: name of the sample - - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) - Outputs: - - Vt: UV map vertices - - Ft: UV map faces - """ - def read_garment_UVMap(self, sample, garment): - # Read OBJ file - obj_path = os.path.join(self.SRC, sample, garment + '.obj') - return readOBJ(obj_path)[2:] + """ + Reads garment UV map for the specified sample and garment + Inputs: + - sample: name of the sample + - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) + Outputs: + - Vt: UV map vertices + - Ft: UV map faces + """ + def read_garment_UVMap(self, sample, garment): + # Read OBJ file + obj_path = os.path.join(self.SRC, sample, garment + '.obj') + return readOBJ(obj_path)[2:] - """ - Reads vertex colors of the specified sample and garment - Inputs: - - sample: name of the sample - - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) - - F: mesh faces - - Vt: UV map vertices - - Ft: UV map faces - Output - - C: RGB colors - """ - def read_garment_vertex_colors(self, sample, garment, F, Vt, Ft): - # Read sample texture info - texture = self.read_info(sample)['outfit'][garment]['texture'] - # Plain color - if texture['type'] == 'color': return (255*texture['data']).astype(np.int32) - # Image texture - img_path = os.path.join(self.SRC, sample, garment + '.png') - img = Image.open(img_path) - # Get color of each UV vertex - color = np.array([img.getpixel(uv_to_pixel(vt)) for vt in Vt]) - # Compute correspondece between V (mesh) and Vt (UV map) - m2uv = mesh2UV(F, Ft) - return np.array([color[list(m2uv[idx])].mean(0) for idx in sorted(list(m2uv.keys()))], np.uint8) - - """ Scene data """ - """ - Read camera location and compute projection matrix - Input: - - sample: name of the sample - Output: - - P: camera projection matrix (3 x 4) - """ - def read_camera(self, sample): - # Read sample data - info = self.read_info(sample) - # Camera location - camLoc = info['camLoc'] - # Camera projection matrix - return proj(camLoc) - + """ + Reads vertex colors of the specified sample and garment + Inputs: + - sample: name of the sample + - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...) + - F: mesh faces + - Vt: UV map vertices + - Ft: UV map faces + Output + - C: RGB colors + """ + def read_garment_vertex_colors(self, sample, garment, F, Vt, Ft): + # Read sample texture info + texture = self.read_info(sample)['outfit'][garment]['texture'] + # Plain color + if texture['type'] == 'color': return (255*texture['data']).astype(np.int32) + # Image texture + img_path = os.path.join(self.SRC, sample, garment + '.png') + img = Image.open(img_path) + # Get color of each UV vertex + color = np.array([img.getpixel(uv_to_pixel(vt)) for vt in Vt]) + # Compute correspondece between V (mesh) and Vt (UV map) + m2uv = mesh2UV(F, Ft) + return np.array([color[list(m2uv[idx])].mean(0) for idx in sorted(list(m2uv.keys()))], np.uint8) + + """ Scene data """ + """ + Read camera location and compute projection matrix + Input: + - sample: name of the sample + Output: + - P: camera projection matrix (3 x 4) + """ + def read_camera(self, sample): + # Read sample data + info = self.read_info(sample) + # Camera location + camLoc = info['camLoc'] + # Camera projection matrix + return proj(camLoc) + # TESTING if __name__ == '__main__': - sample = '135_02_s8' - frame = 0 - garment = 'Tshirt' - - reader = DataReader() - F = reader.read_garment_topology(sample, garment) - Vt, Ft = reader.read_garment_UVMap(sample, garment) - C = reader.read_garment_vertex_colors(sample, garment, F, Vt, Ft) \ No newline at end of file + sample = '135_02_s8' + frame = 0 + garment = 'Tshirt' + + reader = DataReader() + F = reader.read_garment_topology(sample, garment) + Vt, Ft = reader.read_garment_UVMap(sample, garment) + C = reader.read_garment_vertex_colors(sample, garment, F, Vt, Ft) \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index 0175c10..53011b2 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -46,21 +46,20 @@ class Dataset(Dataset): # 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) - mask = np.load(os.path.join(path, frame.garment + '_mask.npy')) + mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy')) # TODO: what's the difference between masking and reading .obj file - faces = mask_faces(self.body_faces, mask) + 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 = self.reader.read_human(frame.id, frame.frame_nr) + body, body_faces, root_joint = self.reader.read_human(frame.id, frame.frame_nr) body_faces = quads2tris(body_faces) info = self.reader.read_info(frame.id) - combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda') + combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to('cuda') combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda') - combined_faces_int = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0])).astype(np.int32)).to('cuda') - idx = 36000 - x = combined[idx, :].cpu().detach().numpy() + idx_frame = 36000 + x = combined[idx_frame, :].cpu().detach().numpy() y = np.array(info['camLoc'], dtype=np.float32) dist = np.sqrt(np.sum((x - y) ** 2)) @@ -69,35 +68,42 @@ class Dataset(Dataset): # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0] deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32)) # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined, faces=combined_faces.cpu().numpy().astype(np.int32)) - t = time.perf_counter() - start - print(t) + print(f'deodr time: {time.perf_counter() - start}') + start = time.perf_counter() - depth_image, depth_mask, v2i, K = self.render(combined, combined_faces, info['camLoc']) - t = time.perf_counter() - start - print(t) - - pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True) - pc = pc[0,:,depth_mask].cpu().numpy() + depth_image, depth_mask, v2i, K, camera = self.render(combined, combined_faces, info['camLoc']) + print(f'pytorch3d time: {time.perf_counter() - start}') + start = time.perf_counter() + # TODO: is it sure I need this K, and not some transformed one? normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K) + print(f'normals time: {time.perf_counter() - start}') + root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2] - depth_image = depth_image.cpu().numpy() - dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]] - depth_image[v2i[idx, 1], v2i[idx, 0]] = 100 + # Testing + # pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True) + # pc = pc[0,:,depth_mask].cpu().numpy() + + # depth_image = depth_image.cpu().numpy() + # dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]] + # depth_image[v2i[idx, 1], v2i[idx, 0]] = 100 combined = combined.cpu().numpy() - diff = np.abs(dist - dist_img) + # diff = np.abs(dist - dist_img) - import plotly.express as ex + # import plotly.express as ex import plotly.graph_objects as go # ex.imshow(deodr_depth[:,:,0]).show() # ex.imshow(depth_image).show() - ex.imshow(normal_image[0].cpu().numpy()).show() + # ex.imshow(normal_image[0].cpu().numpy()).show() - # go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'), - # go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), - # ]).show() + # trans = np.array(info['trans'], dtype=np.float32)[:, self.frames[idx].frame_nr] + trans = root_joint[0,:] + go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'), + # go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), + go.Scatter3d(x=[trans[0]], y=[trans[1]], z=[trans[2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), + ]).show() # go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), # ]).show() @@ -107,7 +113,12 @@ class Dataset(Dataset): # ex.imshow(v2i_coverage).show() # image = self.transforms(image) - return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask } + return { 'mesh_posed': mesh_posed, + 'depth_image': depth_image, + 'mesh_mask': mesh_mask, + 'depth_mask': depth_mask, + 'normal_image': normal_image + } if __name__ == '__main__': diff --git a/data/render.py b/data/render.py index 5711023..4176208 100644 --- a/data/render.py +++ b/data/render.py @@ -198,4 +198,4 @@ class DepthRenderer(nn.Module): # import plotly.express as ex # ex.imshow(depth.cpu().numpy()).show() - return depth, mask, v2i, K + return depth, mask, v2i, K, cameras