import json import numpy as np # import tensorflow as tf import torch import torch.nn as nn import scipy as sp import plotly.graph_objects as go from plotly.subplots import make_subplots from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode def m2uv_array(m2uv, is_smpl_orig=False): result = np.zeros((14475 if not is_smpl_orig else 3638, 2), dtype=np.int32) for k, v in m2uv.items(): if len(v) == 2: result[k] = list(v) elif len(v) == 1: x = v.pop() result[k] = [x, x] else: assert False return result class UV2Mesh(nn.Module): def __init__(self, res=256, is_skirt=False, is_smpl_orig=False): super(UV2Mesh, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + ('_orig' if is_smpl_orig else '') + '.npy') self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + ('_orig' if is_smpl_orig else '') + '.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv, is_smpl_orig=is_smpl_orig) self.uv_pixels = self.uv_map_discreet[m2uv] self.smpl_res = 14475 if not is_smpl_orig else 3638 def uv_to_mesh(self, uv): if uv.shape[1] == 3 or uv.shape[1] == 7: uv = torch.permute(uv, [0,2,3,1]) if uv.shape[1] != 512: uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR) # batch_size = uv.shape[0] # x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) # y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) # z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) # x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) # y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32)), axis=-1) # z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, self.smpl_res, 1), dtype=tf.int32) * 2), axis=-1) # x1 = tf.gather_nd(uv, x1_ind, batch_dims=1) # y1 = tf.gather_nd(uv, y1_ind, batch_dims=1) # z1 = tf.gather_nd(uv, z1_ind, batch_dims=1) # x2 = tf.gather_nd(uv, x2_ind, batch_dims=1) # y2 = tf.gather_nd(uv, y2_ind, batch_dims=1) # z2 = tf.gather_nd(uv, z2_ind, batch_dims=1) # p1 = tf.stack((x1, y1, z1), axis=-1) # p2 = tf.stack((x2, y2, z2), axis=-1) mesh = torch.mean(uv[:, self.uv_pixels[:,:,0], self.uv_pixels[:,:,1]].float(), dim=2) return mesh def __call__(self, uv): x = self.uv_to_mesh(uv) return x def _maskFaces(V, F, mask, return_unfixed=False): V_mask = mask F_mask = np.take(V_mask, F.reshape(-1), axis=0) F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 V_ind = np.zeros(V_mask.shape) V_ind[V_mask] = np.arange(V_mask.sum()) F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) V_new = V[V_mask, :] if return_unfixed: return V_new, F_new, F[F_mask, :] else: return V_new, F_new def maskFacesBatch(V, F, mask): Vs = [] Fs = [] for i in range(mask.shape[0]): v, f = _maskFaces(V[i], F, mask[i]) Vs.append(v) Fs.append(f) return Vs, Fs def flip_axis(x, ind=None): y = x.copy() if ind is not None: y[ind:,1] = x[ind:,2] y[ind:,2] = x[ind:,1] else: y[:,1] = x[:,2] y[:,2] = x[:,1] return y def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color, flip_xy=False): fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) if flip_xy: mesh_gt = flip_axis(mesh_gt) mesh_pred_unposed = flip_axis(mesh_pred_unposed) collision = flip_axis(collision, ind=14475) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=mesh_gt[:,0], y=mesh_gt[:,2], z=mesh_gt[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), row=1, col=1) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=mesh_pred_unposed[:,0], y=mesh_pred_unposed[:,2], z=mesh_pred_unposed[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] ), row=1, col=2) fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=collision[:,0], y=collision[:,2], z=collision[:,1], # i, j and k give the vertices of triangles i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ), row=1, col=3) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), ) helper = dict(xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) fig.update_layout(scene = helper, scene2 = helper, scene3 = helper, ) return fig def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False): fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) if flip_xy: collision_1 = flip_axis(collision_1, ind=14475) collision_2 = flip_axis(collision_2, ind=14475) gt = flip_axis(gt, ind=14475) fig.add_trace(go.Mesh3d( x=gt[:,0], y=gt[:,2], z=gt[:,1], i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ), row=1, col=1) fig.add_trace(go.Mesh3d( x=collision_1[:,0], y=collision_1[:,2], z=collision_1[:,1], i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ), row=1, col=2) fig.add_trace(go.Mesh3d( x=collision_2[:,0], y=collision_2[:,2], z=collision_2[:,1], i = collision_faces[:,0], j = collision_faces[:,1], k = collision_faces[:,2], vertexcolor=collision_color ), row=1, col=3) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1), scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1), ) helper = dict(xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) fig.update_layout(scene = helper, scene2 = helper, scene3 = helper, ) return fig def get_summary_str(model): # from torchview import draw_graph # draw_graph(model, input_size=((4,3,512,256), (4,512,256)), expand_nested=True) # lines = [] # model.summary(print_fn=lines.append, line_length=150,expand_nested=True,show_trainable=True) # Add initial spaces to avoid markdown formatting in TensorBoard trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) total_params = sum(param.numel() for param in model.parameters()) return [f'all_params: {total_params}, trainable: {trainable_params}'] def get_summary_html(model): lines = get_summary_str(model) html = '' for line in lines: html += "
" + line + "
" return html class Accumulator(): def __init__(self): self.dict = dict() def update(self, dict): for key in dict.keys(): if key in self.dict.keys(): self.dict[key].append(dict[key]) else: self.dict[key] = [dict[key]] def get(self): result = dict() for key in self.dict.keys(): result[key] = np.mean(self.dict[key]) return result def reset(self): self.dict = dict() from stl import mesh as mesh_stl def save_stl(file, mesh, faces): obj = mesh_stl.Mesh(np.zeros(faces.shape[0], dtype=mesh_stl.Mesh.dtype)) for i, f in enumerate(faces): for j in range(3): obj.vectors[i][j] = mesh[f[j],:] obj.save(file) def write_obj(file, vertices, faces, colors): with open(file, 'w') as file: for v, c in zip(vertices, colors): line = 'v ' + ' '.join([str(_) for _ in v]) + ' ' + ' '.join([str(_) for _ in c]) + '\n' file.write(line) for f in faces: line = 'f ' + ' '.join([str(_) for _ in f]) + '\n' file.write(line)