diff --git a/cloth3d/read.py b/cloth3d/read.py new file mode 100644 index 0000000..2611ab1 --- /dev/null +++ b/cloth3d/read.py @@ -0,0 +1,173 @@ +import os +import sys +import numpy as np +from PIL import Image + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +from cloth3d.smpl_np import SMPLModel +from cloth3d.util import loadInfo, zRotMatrix, proj, mesh2UV, uv_to_pixel +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 + + """ + 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 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 diff --git a/cloth3d/smpl_np.py b/cloth3d/smpl_np.py new file mode 100644 index 0000000..e3880d7 --- /dev/null +++ b/cloth3d/smpl_np.py @@ -0,0 +1,187 @@ +import sys +import numpy as np +import pickle + +class SMPLModel(): + def __init__(self, model_path): + """ + SMPL model. + + Parameter: + --------- + model_path: Path to the SMPL model parameters, pre-processed by + `preprocess.py`. + + """ + with open(model_path, 'rb') as f: + if sys.version_info[0] == 2: + params = pickle.load(f) # Python 2.x + elif sys.version_info[0] == 3: + params = pickle.load(f, encoding='latin1') # Python 3.x + self.J_regressor = params['J_regressor'] + self.weights = params['weights'] + self.posedirs = params['posedirs'] + self.v_template = params['v_template'] + self.shapedirs = params['shapedirs'] + self.faces = params['f'] + self.kintree_table = params['kintree_table'] + + id_to_col = { + self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1]) + } + self.parent = { + i: id_to_col[self.kintree_table[0, i]] + for i in range(1, self.kintree_table.shape[1]) + } + + self.pose_shape = [24, 3] + self.beta_shape = [10] + self.trans_shape = [3] + + self.pose = np.zeros(self.pose_shape) + self.beta = np.zeros(self.beta_shape) + self.trans = np.zeros(self.trans_shape) + + self.verts = None + self.J = None + self.R = None + + self.update() + + def set_params(self, pose=None, beta=None, trans=None): + """ + Set pose, shape, and/or translation parameters of SMPL model. Verices of the + model will be updated and returned. + + Parameters: + --------- + pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation + relative to parent joint. For root joint it's global orientation. + Represented in a axis-angle format. + + beta: Parameter for model shape. A vector of shape [10]. Coefficients for + PCA component. Only 10 components were released by MPI. + + trans: Global translation of shape [3]. + + Return: + ------ + Updated vertices. + + """ + if pose is not None: + self.pose = pose + if beta is not None: + self.beta = beta + if trans is not None: + self.trans = trans + self.update() + return self.verts, self.J + + def update(self): + """ + Called automatically when parameters are updated. + + """ + # how beta affect body shape + v_shaped = self.shapedirs.dot(self.beta) + self.v_template + # joints location + self.J = self.J_regressor.dot(v_shaped) + pose_cube = self.pose.reshape((-1, 1, 3)) + # rotation matrix for each joint + self.R = self.rodrigues(pose_cube) + I_cube = np.broadcast_to( + np.expand_dims(np.eye(3), axis=0), + (self.R.shape[0]-1, 3, 3) + ) + lrotmin = (self.R[1:] - I_cube).ravel() + # how pose affect body shape in zero pose + v_posed = v_shaped + self.posedirs.dot(lrotmin) + # world transformation of each joint + G = np.empty((self.kintree_table.shape[1], 4, 4)) + G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1])))) + for i in range(1, self.kintree_table.shape[1]): + G[i] = G[self.parent[i]].dot( + self.with_zeros( + np.hstack( + [self.R[i],((self.J[i, :]-self.J[self.parent[i],:]).reshape([3,1]))] + ) + ) + ) + G = G - self.pack( + np.matmul( + G, + np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1]) + ) + ) + # transformation of each vertex + T = np.tensordot(self.weights, G, axes=[[1], [0]]) + rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1]))) + v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3] + self.verts = v + self.trans.reshape([1, 3]) + + def rodrigues(self, r): + """ + Rodrigues' rotation formula that turns axis-angle vector into rotation + matrix in a batch-ed manner. + + Parameter: + ---------- + r: Axis-angle rotation vector of shape [batch_size, 1, 3]. + + Return: + ------- + Rotation matrix of shape [batch_size, 3, 3]. + + """ + theta = np.linalg.norm(r, axis=(1, 2), keepdims=True) + # avoid zero divide + theta = np.maximum(theta, np.finfo(np.float64).tiny) + r_hat = r / theta + cos = np.cos(theta) + z_stick = np.zeros(theta.shape[0]) + m = np.dstack([ + z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], + r_hat[:, 0, 2], z_stick, -r_hat[:, 0, 0], + -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick] + ).reshape([-1, 3, 3]) + i_cube = np.broadcast_to( + np.expand_dims(np.eye(3), axis=0), + [theta.shape[0], 3, 3] + ) + A = np.transpose(r_hat, axes=[0, 2, 1]) + B = r_hat + dot = np.matmul(A, B) + R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m + return R + + def with_zeros(self, x): + """ + Append a [0, 0, 0, 1] vector to a [3, 4] matrix. + + Parameter: + --------- + x: Matrix to be appended. + + Return: + ------ + Matrix after appending of shape [4,4] + + """ + return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]]))) + + def pack(self, x): + """ + Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched + manner. + + Parameter: + ---------- + x: Matrices to be appended of shape [batch_size, 4, 1] + + Return: + ------ + Matrix of shape [batch_size, 4, 4] after appending. + + """ + return np.dstack((np.zeros((x.shape[0], 4, 3)), x)) \ No newline at end of file diff --git a/cloth3d/util.py b/cloth3d/util.py new file mode 100644 index 0000000..c63cc76 --- /dev/null +++ b/cloth3d/util.py @@ -0,0 +1,122 @@ +import numpy as np +import scipy.io as sio +from math import cos, sin + +def loadInfo(filename): + ''' + this function should be called instead of direct sio.loadmat + as it cures the problem of not properly recovering python dictionaries + from mat files. It calls the function check keys to cure all entries + which are still mat-objects + ''' + data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True) + del data['__globals__'] + del data['__header__'] + del data['__version__'] + return _check_keys(data) + +def _check_keys(dict): + ''' + checks if entries in dictionary are mat-objects. If yes + todict is called to change them to nested dictionaries + ''' + for key in dict: + if isinstance(dict[key], sio.matlab.mio5_params.mat_struct): + dict[key] = _todict(dict[key]) + return dict + +def _todict(matobj): + ''' + A recursive function which constructs from matobjects nested dictionaries + ''' + dict = {} + for strg in matobj._fieldnames: + elem = matobj.__dict__[strg] + if isinstance(elem, sio.matlab.mio5_params.mat_struct): + dict[strg] = _todict(elem) + elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]): + dict[strg] = [None] * len(elem) + for i,item in enumerate(elem): + if isinstance(item, sio.matlab.mio5_params.mat_struct): + dict[strg][i] = _todict(item) + else: + dict[strg][i] = item + else: + dict[strg] = elem + return dict + +# Computes matrix of rotation around z-axis for 'zrot' radians +def zRotMatrix(zrot): + c, s = cos(zrot), sin(zrot) + return np.array([[c, -s, 0], + [s, c, 0], + [0, 0, 1]], np.float32) + + +""" CAMERA """ +def intrinsic(): + RES_X = 640 + RES_Y = 480 + f_mm = 50 # blender default + sensor_w_mm = 36 # blender default + sensor_h_mm = sensor_w_mm * RES_Y / RES_X + + fx_px = f_mm * RES_X / sensor_w_mm; + fy_px = f_mm * RES_Y / sensor_h_mm; + + u = RES_X / 2; + v = RES_Y / 2; + + return np.array([[fx_px, 0, u], + [0, fy_px, v], + [0, 0, 1]], np.float32) + +def extrinsic(camLoc): + R_w2bc = np.array([[0, 1, 0], + [0, 0, 1], + [1, 0, 0]], np.float32) + + T_w2bc = -1 * R_w2bc.dot(camLoc) + + R_bc2cv = np.array([[1, 0, 0], + [0, -1, 0], + [0, 0, -1]], np.float32) + + R_w2cv = R_bc2cv.dot(R_w2bc) + T_w2cv = R_bc2cv.dot(T_w2bc) + + return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1) + +def proj(camLoc): + return intrinsic().dot(extrinsic(camLoc)) + +""" +Mesh to UV map +Computes correspondences between 3D mesh and UV map +NOTE: 3D mesh vertices can have multiple correspondences with UV vertices +""" +def mesh2UV(F, Ft): + m2uv = {v: set() for f in F for v in f} + for f, ft in zip(F, Ft): + for v, vt in zip(f, ft): + m2uv[v].add(vt) + # m2uv = {k:list(v) for k,v in m2uv.items()} + return m2uv + +# Maps UV coordinates to texture space (pixel) +IMG_SIZE = 2048 # all image textures have this squared size +def uv_to_pixel(vt): + px = vt * IMG_SIZE # scale to image plane + px %= IMG_SIZE # wrap to [0, IMG_SIZE] + # Note that Blender graphic engines invert vertical axis + return int(px[0]), int(IMG_SIZE - px[1]) # texel X, texel Y + + +def quads2tris(F): + out = [] + for f in F: + if len(f) == 3: out += [f] + elif len(f) == 4: out += [[f[0],f[1],f[2]], + [f[0],f[2],f[3]]] + else: print("This should not happen...") + return np.array(out, np.int32) \ No newline at end of file diff --git a/data/dataset.py b/data/dataset.py index 44ba593..c51923d 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -18,6 +18,8 @@ from render import Render, DepthRenderer sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from superSMPL.IO import readPC2Frame, writePC2 from utils.data_utils import Frame, read_frames_file, mask_faces +from cloth3d.read import DataReader +from cloth3d.util import quads2tris class Dataset(Dataset): @@ -27,30 +29,38 @@ class Dataset(Dataset): self.frames = read_frames_file(os.path.join(folder, frames_file)) self.render = DepthRenderer() + self.deodr_render : Render = Render(max_depth=10) self.body_faces = np.load('custom_data/smpl_faces.npy') + self.reader = DataReader(folder, 'superSMPL') def __len__(self): return len(self.frames) + # TODO: make it work kwith two part outfits def __getitem__(self, idx): frame : Frame = self.frames[idx] path = os.path.join(self.folder, frame.id) - 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) mask = np.load(os.path.join(path, frame.garment + '_mask.npy')) faces = mask_faces(self.body_faces, mask) - body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True) + # 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_faces = quads2tris(body_faces) - info = sio.loadmat(os.path.join(path, 'info.mat')) + info = self.reader.read_info(frame.id) combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda') - combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda') + combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to('cuda') 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] - depth_image = self.render(combined, combined_faces, info['camLoc'][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)) + depth_image = self.render(combined, combined_faces, info['camLoc']) t = time.perf_counter() - start print(t) @@ -65,7 +75,7 @@ if __name__ == '__main__': # dataset = Dataset('/home/workspace/local-data/train') dataset = Dataset('dataset/CLOTH3D/train') - train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True) + train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: batch diff --git a/data/render.py b/data/render.py index 516d7d1..649c397 100644 --- a/data/render.py +++ b/data/render.py @@ -150,9 +150,9 @@ class DepthRenderer(nn.Module): def forward(self, verts, faces, cam): R, T = look_at_view_transform(2.7, 0, 0) - cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T) + # cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T) ext = extrinsic(cam) - cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=50) + cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3]) raster_settings = RasterizationSettings( image_size=512, @@ -165,7 +165,8 @@ class DepthRenderer(nn.Module): cameras=cameras, raster_settings=raster_settings ) - depth = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])).zbuf[0,:,:,0] + output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])) + depth = output.zbuf[0,:,:,0] mask = depth != -1 # dmin = depth.min() diff --git a/utils/data_utils.py b/utils/data_utils.py index f754154..3316c2e 100644 --- a/utils/data_utils.py +++ b/utils/data_utils.py @@ -1,8 +1,11 @@ import os import sys import numpy as np +import scipy.io as sio +from math import cos, sin sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from superSMPL.IO import readPC2Frame, readOBJ class Frame: @@ -30,4 +33,4 @@ def read_frames_file(file): for line in lines: id, frame_nr, garment = line.split() frames.append(Frame(id, int(frame_nr), garment)) - return frames \ No newline at end of file + return frames