def main(img_path, json_path=None): sess = tf.Session() model = RunModel(config, sess=sess) input_img, proc_param, img = preprocess_image(img_path, json_path) # Add batch dimension: 1 x D x D x 3 input_img = np.expand_dims(input_img, 0) joints, verts, cams, joints3d, theta = model.predict(input_img, get_theta=True) cams = theta[:, :model.num_cam] poses = theta[:, model.num_cam:(model.num_cam + model.num_theta)] shapes = theta[:, (model.num_cam + model.num_theta):] visualize(img, proc_param, joints[0], verts[0], cams[0]) ''' Start adjusting the shape ''' shape_adjuster = torch.load("./trained/model_save1.57873062595") smpl = SMPL("./models/neutral_smpl_with_cocoplus_reg.pkl") beta = torch.from_numpy(shapes).float().cuda() theta = torch.zeros((1, 72)).float().cuda() heights = torch.from_numpy(np.asarray([1.9])) volume = torch.from_numpy(np.asarray([90 / 1000.])) verts, joints3d, Rs = smpl.forward(beta, theta, True) flatten_joints3d = joints3d.view(1, -1) heights = torch.unsqueeze(heights, -1).float().cuda() volumes = torch.unsqueeze(volume, -1).float().cuda() input_to_net = torch.cat((flatten_joints3d, heights, volumes), 1) adjusted_betas = shape_adjuster.forward(input_to_net) adjusted_verts, adjusted_joints3d, Rs = smpl.forward( adjusted_betas, theta, True) adjusted_heights = measure.compute_height(adjusted_verts) adjusted_volumes = measure.compute_volume(adjusted_verts, smpl.f) print(adjusted_heights, adjusted_volumes) debug_display_cloud(verts[0], joints3d[0], adjusted_verts[0], adjusted_joints3d[0])
class Trainer(Trainable): def _setup(self): # 1) Initialize all needed parameters self.smpl = SMPL('../models/neutral_smpl_with_cocoplus_reg.pkl') self.camera = Camera() self.batch_size = int(self.config["batch_size"]) self.net = DenseNet(self.config) self.net.cuda() # optimizer setup self.optimizer = torch.optim.Adam( self.net.parameters(), lr=self.config["lr"], weight_decay=self.config["weight_decay"]) self.scheduler = ReduceLROnPlateau(self.optimizer, 'min', verbose=True) self.min_loss = inf self.num_iter = 5000 def _sample_random_theta(self): # The pose can be not natural, since the task is to predict shape. return np.asarray(np.random.uniform(-2 * np.pi, 2 * np.pi, 72)) # The inner function for batch generation def _get_batch(self): beta = 3 * torch.randn((self.batch_size, 10)).float().cuda() theta = torch.zeros((self.batch_size, 72)).float().cuda() # Without pose but with shape for measuring verts, joints3d, Rs = self.smpl.forward(beta, theta, True) heights = measure.compute_height(verts) volumes = measure.compute_volume(verts, self.smpl.f) # Must add artificial noise to the joints since joints detection algorithms are not perfect # -+ 2 cm noise = torch.normal(torch.zeros_like(joints3d), 0.04).float().cuda() noisy_joints3d = joints3d + noise # Visualize noisy joints #debug_display_joints(noisy_joints3d[0], joints3d[0]) scale = torch.rand((self.batch_size, 1, 1)) * 3.75 + 0.25 noisy_joints3d /= scale.float().cuda() # Must add artificial noise to the height and volume. # Volume must represent a weight of a person, so -+ 3 kg # Height is -+ 2 cm noise = torch.normal(torch.zeros_like(heights), 0.03).float().cuda() heights += noise noise = torch.normal(torch.zeros_like(volumes), volumes / 100.0).float().cuda() volumes += noise noisy_joints3d = noisy_joints3d.view(self.batch_size, -1) heights = torch.unsqueeze(heights, -1) volumes = torch.unsqueeze(volumes, -1) input_to_net = torch.cat((noisy_joints3d, heights, volumes), 1) return input_to_net.detach(), torch.squeeze(beta).detach() def _save(self, checkpoint_dir, postfix=None): file_path = checkpoint_dir + "/model_save_new" + postfix torch.save(self.net, file_path) return file_path def _restore(self, path): self.net = torch.load(path) def _stop(self): # If need, save your model when exit. saved_path = self.logdir + "/model_stop" torch.save(self.net, saved_path) print("save model at: ", saved_path) def _eval(self): total_loss = 0 for i in range(self.num_iter): inputs, beta = self._get_batch() # Evaluation =============================================================================================== self.net.eval() predicted_beta = self.net.forward(inputs) beta_loss = F.mse_loss(predicted_beta, beta) total_loss += float(beta_loss.cpu().detach().numpy()) total_loss = total_loss / float(self.num_iter) self.scheduler.step(total_loss) print '\n Evaluation finished: ', total_loss if self.min_loss > total_loss: self.min_loss = total_loss print 'Saving ...', ' the current loss is ', self.min_loss trainer._save('../trained/', str(self.min_loss)) theta = torch.zeros((self.batch_size, 72)).cuda() verts, joints3d, Rs = self.smpl.forward(beta, theta, True) predicted_verts, predicted_joints3d, Rs = self.smpl.forward( predicted_beta, theta, True) debug_display_cloud(verts[0], joints3d[0], predicted_verts[0], predicted_joints3d[0], self.min_loss, True) def _train(self): total_loss = 0 for i in range(self.num_iter): inputs, beta = self._get_batch() self.net.train() # clean the gradients self.net.zero_grad() self.camera.zero_grad() self.smpl.zero_grad() # Prediction =============================================================================================== # predicted_theta, predicted_beta, predicted_camera_parameters = net.forward(joints2d) predicted_beta = self.net.forward(inputs) beta_loss = F.mse_loss(predicted_beta, beta) beta_loss.backward() self.optimizer.step() total_loss += float(beta_loss.cpu().detach().numpy()) # debug_display_cloud if i % 100 == True: print i, total_loss = total_loss / float(self.num_iter) return TrainingResult(timesteps_this_iter=1, mean_loss=total_loss)
class Trainer(Trainable): def _setup(self): # 1) Initialize all needed parameters self.smpl = SMPL( '/home/king/Projects/GITHUB/hmr/models/neutral_smpl_with_cocoplus_reg.pkl' ) self.camera = Camera() self.batch_size = int(self.config["batch_size"]) self.net = DenseNet(self.config) self.net.cuda() # optimizer setup self.optimizer = torch.optim.Adam( self.net.parameters(), lr=self.config["lr"], weight_decay=self.config["weight_decay"]) self.posesV0 = glob.glob( '/home/king/Documents/measurement/dataset/up-3d/*.pkl') def _sample_random_theta(self): root = '/home/king/Documents/measurement/dataset/up-3d/' path = self.posesV0[np.random.random_integers(0, len(self.posesV0) - 1)] with open(path, 'rb') as f: # print path stored_parameters = pickle.load(f) orig_pose = np.array(stored_parameters['pose']).copy() orig_rt = np.array(stored_parameters['rt']).copy() orig_trans = np.array(stored_parameters['trans']).copy() orig_t = np.array(stored_parameters['t']).copy() thetas = stored_parameters['pose'] thetas[:3] = np.random.uniform(0, 2 * np.pi, 3) return np.asarray(thetas) # The inner function for batch generation def _get_batch(self, N): beta = 4 * torch.randn( (N, 10)).float().cuda() #torch.zeros((N, 10)).float().cuda() theta = np.ones((N, 72)) * np.expand_dims(self._sample_random_theta(), 0) #np.zeros((N, 72)) # theta = torch.from_numpy(theta).float().cuda() # Without pose but with shape for measuring verts, joints3d, Rs = self.smpl.forward(beta, theta, True) heights = measure.compute_height(verts) volumes = measure.compute_volume(verts, self.smpl.f) joints2d, camera_parameters = self.camera.forward(joints3d) return joints2d, verts, joints3d, camera_parameters, beta, theta, heights, volumes def _save(self, checkpoint_dir): file_path = checkpoint_dir + "/model_save" torch.save(self.net, file_path) return file_path def _restore(self, path): self.net = torch.load(path) def _stop(self): # If need, save your model when exit. saved_path = self.logdir + "/model_stop" torch.save(self.net, saved_path) print("save model at: ", saved_path) def _train(self): # 2) After the initialization has been done, we can start training the model # The training loop # Get the training samples joints2d, verts, joints3d, camera_parameters, beta, theta, heights, volumes = self._get_batch( self.batch_size) joints2d = joints2d.view((self.batch_size, -1)) heights = torch.unsqueeze(heights, -1) volumes = torch.unsqueeze(volumes, -1) input_to_net = torch.cat((joints2d, heights, volumes), 1) input_to_net = torch.unsqueeze(input_to_net, -1) eval_loss = inf while eval_loss > (1.3 / 32.): self.net.train() # clean the gradients self.net.zero_grad() self.camera.zero_grad() self.smpl.zero_grad() # Prediction =============================================================================================== # predicted_theta, predicted_beta, predicted_camera_parameters = net.forward(joints2d) predicted_theta, predicted_beta = self.net.forward(input_to_net) verts, predicted_joints3d, Rs = self.smpl.forward( predicted_beta, predicted_theta, True) predicted_joints2d, _ = self.camera.forward( predicted_joints3d, camera_parameters) # make sure they are references to the same object assert (_ is camera_parameters) verts, _, Rs = self.smpl.forward(predicted_beta, torch.zeros_like(predicted_theta), True) predicted_volumes = measure.compute_height(verts) predicted_heights = measure.compute_volume(verts, self.smpl.f) beta_loss = F.mse_loss(torch.squeeze(predicted_beta), torch.squeeze(beta)) theta_loss = F.mse_loss(torch.squeeze(predicted_theta), torch.squeeze(theta)) total_loss = theta_loss + beta_loss total_loss.backward() self.optimizer.step() # Evaluation =============================================================================================== self.net.eval() predicted_theta, predicted_beta = self.net.forward(input_to_net) verts, predicted_joints3d, Rs = self.smpl.forward( predicted_beta, predicted_theta, True) predicted_joints2d, _ = self.camera.forward( predicted_joints3d, camera_parameters) beta_loss = F.mse_loss(torch.squeeze(predicted_beta), torch.squeeze(beta)) theta_loss = F.mse_loss(torch.squeeze(predicted_theta), torch.squeeze(theta)) eval_loss = (theta_loss + beta_loss) / self.batch_size if float(np.random.random_sample()) > 0.99: true_verts, true_joints3d, Rs = self.smpl.forward( beta, theta, True) debug_display_cloud(verts[0], joints3d[0], true_verts[0], true_joints3d[0]) debug_display_joints(predicted_joints2d[0], joints2d[0]) print eval_loss return TrainingResult(timesteps_this_iter=1, mean_loss=eval_loss)
def predict(image, weight, height): global config, renderer config = flags.FLAGS config(sys.argv) # Using pre-trained model, change this to use your own. config.load_path = src.config.PRETRAINED_MODEL config.batch_size = 1 renderer = vis_util.SMPLRenderer(face_path=config.smpl_face_path) tf.reset_default_graph() sess = tf.Session() model = RunModel(config, sess=sess) input_img, proc_param, img = preprocess_image_V2(image) # Add batch dimension: 1 x D x D x 3 input_img = np.expand_dims(input_img, 0) joints, verts, cams, joints3d, theta = model.predict(input_img, get_theta=True) sess.close() cams = theta[:, :model.num_cam] poses = theta[:, model.num_cam:(model.num_cam + model.num_theta)] shapes = theta[:, (model.num_cam + model.num_theta):] viz_result = visualize(img, proc_param, joints[0], verts[0], cams[0]) ''' Start adjusting the shape ''' shape_adjuster = torch.load("./trained/model_release_1.5961363467") smpl = SMPL("./models/neutral_smpl_with_cocoplus_reg.pkl") beta = torch.from_numpy(shapes).float().cuda() theta = torch.zeros((1, 72)).float().cuda() heights = torch.from_numpy(np.asarray([height])) volume = torch.from_numpy(np.asarray([weight])) verts, joints3d, Rs = smpl.forward(beta, theta, True) flatten_joints3d = joints3d.view(1, -1) heights = torch.unsqueeze(heights, -1).float().cuda() volumes = torch.unsqueeze(volume, -1).float().cuda() input_to_net = torch.cat((flatten_joints3d, heights, volumes), 1) adjusted_betas = shape_adjuster.forward(input_to_net) adjusted_verts, adjusted_joints3d, Rs = smpl.forward( adjusted_betas, theta, True) adjusted_heights = measure.compute_height(adjusted_verts) adjusted_volumes = measure.compute_volume(adjusted_verts, smpl.f) print(adjusted_heights, adjusted_volumes) # debug_display_cloud(verts[0], joints3d[0], adjusted_verts[0], adjusted_joints3d[0]) # Change the posture for measurement from measurement import POSE1 theta = torch.from_numpy(np.expand_dims(POSE1, 0)).float().cuda() m_adjusted_verts, adjusted_joints3d, Rs = smpl.forward( adjusted_betas, theta, True) return viz_result, \ torch.squeeze(verts).detach().cpu().numpy(), \ torch.squeeze(adjusted_verts).detach().cpu().numpy(), \ torch.squeeze(m_adjusted_verts).detach().cpu().numpy(), \ torch.squeeze(adjusted_volumes).detach().cpu().numpy(),\ torch.squeeze(adjusted_heights).detach().cpu().numpy(),
class Trainer(Trainable): def _setup(self): # 1) Initialize all needed parameters self.smpl = SMPL( '/home/king/Projects/GITHUB/hmr/models/neutral_smpl_with_cocoplus_reg.pkl' ) self.camera = Camera() self.batch_size = int(self.config["batch_size"]) # optimizer setup self.posesV0 = glob.glob( '/home/king/Documents/measurement/dataset/up-3d/*.pkl') def _sample_random_theta(self): root = '/home/king/Documents/measurement/dataset/up-3d/' path = self.posesV0[np.random.random_integers(0, len(self.posesV0) - 1)] with open(path, 'rb') as f: # print path stored_parameters = pickle.load(f) orig_pose = np.array(stored_parameters['pose']).copy() orig_rt = np.array(stored_parameters['rt']).copy() orig_trans = np.array(stored_parameters['trans']).copy() orig_t = np.array(stored_parameters['t']).copy() thetas = stored_parameters['pose'] * 0 thetas[:1] = np.random.uniform(0, 2 * np.pi, 1) return np.asarray(thetas) # The inner function for batch generation def _get_batch(self, N): beta = 4 * torch.randn( (N, 10)).float().cuda() #torch.zeros((N, 10)).float().cuda() theta = np.ones((N, 72)) * np.expand_dims(self._sample_random_theta(), 0) #np.zeros((N, 72)) # theta = torch.from_numpy(theta).float().cuda() #theta.requires_grad=True # Without pose but with shape for measuring verts, joints3d, Rs = self.smpl.forward(beta, theta, True) heights = measure.compute_height(verts) volumes = measure.compute_volume(verts, self.smpl.f) self.camera._init_camera_randomly(N) joints2d = self.camera.forward(joints3d) return joints2d, verts, joints3d, beta, theta, heights, volumes def _save(self, checkpoint_dir): file_path = checkpoint_dir + "/model_save" torch.save(self.net, file_path) return file_path def _stop(self): # If need, save your model when exit. saved_path = self.logdir + "/model_stop" torch.save(self.net, saved_path) print("save model at: ", saved_path) def _train(self): # 2) After the initialization has been done, we can start training the model # The training loop # Get the training samples joints2d, verts, joints3d, beta, theta, heights, volumes = self._get_batch( 1) heights = torch.unsqueeze(heights, -1) volumes = torch.unsqueeze(volumes, -1) joints3d = joints3d.detach() joints2d = joints2d.detach() volumes = volumes.detach() heights = heights.detach() predicted_beta = Parameter(torch.zeros(1, 10).cuda()) predicted_theta = Parameter(torch.zeros(1, 72).cuda()) self.optimizer = torch.optim.SGD( [predicted_beta, predicted_theta], lr=self.config["lr"], weight_decay=self.config["weight_decay"], momentum=0.9, nesterov=True) for i in range(300000): self.camera.zero_grad() self.smpl.zero_grad() # Prediction =============================================================================================== # predicted_theta, predicted_beta, predicted_camera_parameters = net.forward(joints2d) predicted_verts, predicted_joints3d, Rs = self.smpl.forward( predicted_beta, predicted_theta, True) predicted_joints2d = self.camera.forward(predicted_joints3d) # verts, _, Rs = self.smpl.forward(predicted_beta, torch.zeros_like(predicted_theta), True) # predicted_volumes = measure.compute_height(verts) # predicted_heights = measure.compute_volume(verts, self.smpl.f) # torch.mean(predicted_joints2d).backward() # beta_loss = F.smooth_l1_loss(torch.squeeze(predicted_beta), torch.squeeze(beta)) # theta_loss = F.smooth_l1_loss(torch.squeeze(predicted_theta), torch.squeeze(theta)) joints2d_loss = torch.mean( torch.sum((torch.squeeze(predicted_joints3d) - torch.squeeze(joints3d))**2, 1)) total_loss = joints2d_loss # theta_loss + beta_loss total_loss.backward(retain_graph=True) self.optimizer.step() if float(np.random.random_sample()) > 0.95: debug_display_cloud(predicted_verts[0], predicted_joints3d[0], verts[0], joints3d[0]) debug_display_joints(predicted_joints2d[0], joints2d[0]) print total_loss