use original cloth3d functions

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)

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))

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)

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

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()

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: