def on_changed(self, which): if 'model' in which: if not isinstance(self.model, dict): dd = pkl.load(open(self.model, "rb"), encoding="latin1") else: dd = self.model backwards_compatibility_replacements(dd) # for s in ['v_template', 'weights', 'posedirs', 'pose', 'trans', 'shapedirs', 'betas', 'J']: for s in ['posedirs', 'shapedirs']: if (s in dd) and not hasattr(dd[s], 'dterms'): dd[s] = ch.array(dd[s]) self.f = dd['f'] self.shapedirs = dd['shapedirs'] self.J_regressor = dd['J_regressor'] if 'J_regressor_prior' in dd: self.J_regressor_prior = dd['J_regressor_prior'] self.bs_type = dd['bs_type'] self.bs_style = dd['bs_style'] self.weights = ch.array(dd['weights']) if 'vert_sym_idxs' in dd: self.vert_sym_idxs = dd['vert_sym_idxs'] if 'weights_prior' in dd: self.weights_prior = dd['weights_prior'] self.kintree_table = dd['kintree_table'] self.posedirs = dd['posedirs'] if not hasattr(self, 'betas'): self.betas = ch.zeros(self.shapedirs.shape[-1]) if not hasattr(self, 'trans'): self.trans = ch.zeros(3) if not hasattr(self, 'pose'): self.pose = ch.zeros(72) if not hasattr(self, 'v_template'): self.v_template = ch.array(dd['v_template']) if not hasattr(self, 'v_personal'): self.v_personal = ch.zeros_like(self.v_template) self._set_up()
def optimize_smal(fposes, ftrans, fbetas, model, cams, segs, imgs, landmarks, landmarks_names, key_vids, symIdx=None, frameId=0, opt_model_dir=None, save_name=None, COMPUTE_OPT=True, img_paths=None, img_offset=None, img_scales=None): mesh_v_opt_save_path = join(opt_model_dir, 'mesh_v_opt_no_mc_' + str(frameId) + '.ply') mesh_v_opt_mc_save_path = join(opt_model_dir, 'mesh_v_opt_' + str(frameId) + '.ply') mesh_init_save_path = join(opt_model_dir, 'mesh_init_' + str(frameId) + '.ply') nViews = len(fposes) if not COMPUTE_OPT: if not exists(opt_model_dir): makedirs(opt_model_dir) dv = 0 compute_texture(nViews, opt_model_dir, dv, model, frameId, mesh_init_save_path, fposes, ftrans, fbetas, '_no_refine', cams, imgs, segs, img_paths, img_offset, img_scales) return # Write the initial mesh np_betas = np.zeros_like(model.betas) np_betas[:len(fbetas[0])] = fbetas[0] tmp = verts_decorated(v_template=model.v_template, pose=ch.zeros_like(model.pose.r), trans=ch.zeros_like(model.trans), J=model.J_regressor, kintree_table=model.kintree_table, betas=ch.array(np_betas), weights=model.weights, posedirs=model.posedirs, shapedirs=model.shapedirs, bs_type='lrotmin', bs_style='lbs', f=model.f) tmp_mesh = Mesh(v=tmp.r, f=tmp.f) tmp_path = join(opt_model_dir, 'mesh_init_' + str(frameId) + '.ply') tmp_mesh.write_ply(tmp_path) del tmp assert (nViews == len(cams)) assert (nViews == len(segs)) assert (nViews == len(imgs)) # Define a displacement vector. We set a small non zero displacement as initialization dv = ch.array(np.random.rand(model.r.shape[0], 3) / 1000.) # Cell structure for ARAP f = model.f _, A3, A = edgesIdx(nV=dv.shape[0], f=f, save_dir='.', name='smal') wedge = wedges(A3, dv) s = np.zeros_like(dv) arap = ARAP(reg_e=MatVecMult(A3.T, model.ravel() + dv.ravel()).reshape(-1, 3), model_e=MatVecMult(A3.T, model.ravel()).reshape(-1, 3), w=wedge, A=A) k_arap = settings['ref_k_arap_per_view'] * nViews for weight, part in zip(settings['ref_W_arap_values'], settings['ref_W_arap_parts']): k_arap, W_per_vertex = get_arap_part_weights( A, k_arap, [part], [weight]) #, animal_name) # was only Head W = np.zeros((W_per_vertex.shape[0], 3)) for i in range(3): W[:, i] = W_per_vertex k_lap = settings['ref_k_lap'] * nViews * W k_sym = settings['ref_k_sym'] * nViews k_keyp = settings['ref_k_keyp_weight'] * nViews # Load already computed mesh if not exists(opt_model_dir): makedirs(opt_model_dir) shape_model, compute = load_shape_models(nViews, opt_model_dir, dv, model, frameId, mesh_v_opt_save_path, fposes, ftrans, fbetas) mv = None # Remove inside mouth faces ''' if settings['ref_remove_inside_mouth']: # Giraffe faces_orig = shape_model[0].f.copy() im_v = im_up_v + im_down_v idx = [np.where(model.f == ix)[0] for ix in im_v] idx = np.concatenate(idx).ravel() for i in range(nViews): shape_model[i].f = np.delete(shape_model[i].f, idx, 0) ''' if compute: objs = {} FIX_CAM = True free_variables = [] kp_weights = k_keyp * np.ones((landmarks[0].shape[0], 1)) print('removing shoulders, often bad annotated') kp_weights[landmarks_names.index('leftShoulder'), :] *= 0 kp_weights[landmarks_names.index('rightShoulder'), :] *= 0 objs_pose = None j2d = None #k_silh_term = settings['ref_k_silh_term'] k_m2s = settings['ref_k_m2s'] k_s2m = settings['ref_k_s2m'] objs, params_, j2d = set_pose_objs(shape_model, cams, landmarks, key_vids, kp_weights=kp_weights, FIX_CAM=FIX_CAM, ONLY_KEYP=True, OPT_SHAPE=False) if np.any(k_arap) != 0: objs['arap'] = k_arap * arap if k_sym != 0: objs['sym_0'] = k_sym * (ch.abs(dv[:, 0] - dv[symIdx, 0])) objs['sym_1'] = k_sym * (ch.abs(dv[:, 1] + dv[symIdx, 1] - 0.00014954)) objs['sym_2'] = k_sym * (ch.abs(dv[:, 2] - dv[symIdx, 2])) if np.any(k_lap) != 0: lap_op = np.asarray( laplacian(Mesh(v=dv, f=shape_model[0].f)).todense()) objs['lap'] = k_lap * ch.dot(lap_op, dv) mv = None mv2 = MeshViewers(shape=(1, nViews)) #None vc = np.ones_like(dv) dv_r = fit_silhouettes_pyramid_opt(objs, shape_model, dv, segs, cams, j2d=j2d, weights=1., mv=mv, imgs=imgs, s2m_weights=k_s2m, m2s_weights=k_m2s, max_iter=100, free_variables=free_variables, vc=vc, symIdx=symIdx, mv2=mv2, objs_pose=objs_pose) # Save result image for i in range(nViews): img_res = render_mesh(Mesh(shape_model[i].r, shape_model[i].f), imgs[i].shape[1], imgs[i].shape[0], cams[i], img=imgs[i], world_frame=True) img_result = np.hstack((imgs[i], img_res * 255.)) save_img_path = save_name[i].replace('.pkl', '_v_opt.png') cv2.imwrite(save_img_path, img_result) shape_model[0].pose[:] = 0 shape_model[0].trans[:] = 0 V = shape_model[0].r.copy() vm = V[symIdx, :].copy() vm[:, 1] = -1 * vm[:, 1] V2 = (V + vm) / 2.0 mesh_out = Mesh(v=V2, f=shape_model[0].f) mesh_out.show() mesh_out.write_ply(mesh_v_opt_save_path) save_dv_data_path = mesh_v_opt_save_path.replace('.ply', '_dv.pkl') dv_data = {'betas': shape_model[0].betas.r, 'dv': dv_r} pkl.dump(dv_data, open(save_dv_data_path, 'wb')) compute_texture(nViews, opt_model_dir, dv, model, frameId, mesh_v_opt_save_path, fposes, ftrans, fbetas, '_non_opt', cams, imgs, segs, img_paths, img_offset, img_scales) return