got root joint location
This commit is contained in:
+155
-155
@@ -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)
|
||||
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)
|
||||
+36
-25
@@ -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__':
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user