From d1cdece896407ba44418f4ec42b94e1b30376599 Mon Sep 17 00:00:00 2001 From: Wlad <9556979+gosticks@users.noreply.github.com> Date: Mon, 8 Feb 2021 15:33:27 +0100 Subject: [PATCH] cleanup --- animate.py | 18 ++++++ camera_estimation.py | 2 - example_camera.py | 1 - example_fit.py | 53 ++-------------- example_fit_anim.py | 139 ++++++++++++------------------------------ modules/angle.py | 3 +- modules/angle_clip.py | 8 +-- renderer.py | 4 ++ train.py | 101 ++++++++++++++++++++++++++++++ train_pose.py | 5 +- utils/render.py | 22 +------ utils/video.py | 48 +++++++++++++++ 12 files changed, 223 insertions(+), 181 deletions(-) create mode 100644 animate.py create mode 100644 train.py create mode 100644 utils/video.py diff --git a/animate.py b/animate.py new file mode 100644 index 0000000..ecca5b9 --- /dev/null +++ b/animate.py @@ -0,0 +1,18 @@ +import pickle +import time +from tqdm import tqdm +from utils.render import make_video +import torch +from tqdm.auto import trange + +from dataset import SMPLyDataset +from model import * +from utils.general import * +from renderer import * +from camera_estimation import TorchCameraEstimate +from modules.camera import SimpleCamera +from train_pose import train_pose_with_conf +from utils.general import rename_files, get_new_filename + + +def animate_with_conf(config, start=0, end=None): diff --git a/camera_estimation.py b/camera_estimation.py index 22daca9..422ca39 100644 --- a/camera_estimation.py +++ b/camera_estimation.py @@ -19,7 +19,6 @@ class CameraEstimate: def __init__( self, model: smplx.SMPL, - dataset, keypoints, renderer, image_path=None, @@ -31,7 +30,6 @@ class CameraEstimate: self.use_progress_bar = use_progress_bar self.verbose = verbose self.model = model - self.dataset = dataset self.output_model = model(return_verts=True) self.renderer = renderer self.dtype = dtype diff --git a/example_camera.py b/example_camera.py index ccbac0c..863fe46 100644 --- a/example_camera.py +++ b/example_camera.py @@ -25,7 +25,6 @@ init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_tra camera = TorchCameraEstimate( model, - dataset=dataset, keypoints=keypoints, renderer=Renderer(), device=device, diff --git a/example_fit.py b/example_fit.py index 88b35ed..190e9fd 100644 --- a/example_fit.py +++ b/example_fit.py @@ -1,14 +1,9 @@ # library imports -from utils.render import make_video -import torch +from train import optimize_sample import matplotlib.pyplot as plt # local imports -from train_pose import train_pose_with_conf -from modules.camera import SimpleCamera -from model import SMPLyModel -from utils.general import getfilename_from_conf, load_config, setup_training -from camera_estimation import TorchCameraEstimate +from utils.general import load_config from dataset import SMPLyDataset # load and select sample @@ -16,50 +11,14 @@ config = load_config() dataset = SMPLyDataset.from_config(config=config) sample_index = 0 -# prepare data and SMPL model -model = SMPLyModel.model_from_conf(config) -init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training( - model=model, - renderer=True, - dataset=dataset, - sample_index=sample_index, - offscreen=True -) - -# configure PyTorch device and format -dtype = torch.float32 -device = torch.device('cpu') - - -camera = TorchCameraEstimate( - model, - dataset=dataset, - keypoints=keypoints, - renderer=r, - device=device, - dtype=dtype, - image_path=img_path, - est_scale=est_scale -) - -# render camera to the scene -camera.setup_visualization(r.init_keypoints, r.keypoints) - - # train for pose -result, best, train_loss, step_imgs = train_pose_with_conf( - config=config, - model=model, - keypoints=keypoints, - keypoint_conf=conf, - camera=camera, - renderer=r, - device=device, +pose, train_loss, step_imgs = optimize_sample( + sample_index, + dataset, + config ) -make_video(step_imgs, "test.avi") - # color = r.get_snapshot() # plt.imshow(color) # plt.show() diff --git a/example_fit_anim.py b/example_fit_anim.py index 9c053ba..bd82444 100644 --- a/example_fit_anim.py +++ b/example_fit_anim.py @@ -1,27 +1,31 @@ import pickle import time -from utils.render import make_video +from train import create_animation +from tqdm import tqdm +from utils.video import make_video, video_from_pkl import torch -from tqdm.auto import trange from dataset import SMPLyDataset from model import * from utils.general import * from renderer import * -from camera_estimation import TorchCameraEstimate -from modules.camera import SimpleCamera -from train_pose import train_pose_with_conf from utils.general import rename_files, get_new_filename START_IDX = 0 # starting index of the frame to optimize for -FINISH_IDX = 50 # choose a big number to optimize for all frames in samples directory +FINISH_IDX = 2 # choose a big number to optimize for all frames in samples directory # if False, only run already saved animation without optimization -RUN_OPTIMIZATION = True +RUN_OPTIMIZATION = False -final_poses = [] # optimized poses array that is saved for playing the animation result_image = [] idx = START_IDX +device = torch.device('cpu') +dtype = torch.float32 + +config = load_config() +dataset = SMPLyDataset.from_config(config) +model = SMPLyModel.model_from_conf(config) + def get_next_frame(idx): """ @@ -37,86 +41,41 @@ def get_next_frame(idx): return keypoints, keypoints_conf, image_path -device = torch.device('cpu') -dtype = torch.float32 - -config = load_config() -dataset = SMPLyDataset.from_config(config) -model = SMPLyModel.model_from_conf(config) - -samples_dir = config['data']['rootDir'] - # Rename files in samples directory to uniform format if config['data']['renameFiles']: - rename_files(samples_dir + "/") + rename_files(config['data']['rootDir'] + "/") -results_dir = config['output']['rootDir'] -result_prefix = config['output']['prefix'] - -model_out = model() -joints = model_out.joints.detach().cpu().numpy().squeeze() ''' Optimization part without visualization ''' if RUN_OPTIMIZATION: - for idx in trange(FINISH_IDX, desc='Optimizing'): - idx = START_IDX + idx - init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training( - model=model, - renderer=True, - offscreen=True, - dataset=dataset, - sample_index=idx - ) + final_poses, filename = create_animation( + dataset, + config, + START_IDX, + FINISH_IDX, + offscreen=True, + save_to_file=True + ) - r.start() - cam = TorchCameraEstimate( - model, - dataset=dataset, - keypoints=keypoints, - renderer=None, - device=torch.device('cpu'), - dtype=torch.float32, - image_path=img_path, - est_scale=est_scale, - use_progress_bar=False, - verbose=False - ) +def save_to_video(poses, video_name, config, fps=30): + r = DefaultRenderer( + offscreen=True + ) + r.start() - # print("\nCamera optimization of frame", idx, "is finished.") + model_anim = SMPLyModel.model_from_conf(config) - cur_pose, final_pose, loss, frames = train_pose_with_conf( - config=config, - model=model, - keypoints=keypoints, - keypoint_conf=conf, - camera=cam, - renderer=r, - device=device, - use_progress_bar=False - ) + frames = [] - camera_transformation, camera_int, camera_params = cam.get_results() + for body_pose, cam_trans in tqdm(poses): + r.render_model_with_tfs(model_anim, body_pose, keep_pose=True, + render_joints=False, transforms=cam_trans) + frames.append(r.get_snapshot()) - # print("\nPose optimization of frame", idx, "is finished.") - R = camera_transformation.numpy().squeeze() - idx += 1 - - # append optimized pose and camera transformation to the array - final_poses.append((final_pose, R)) - - print("Optimization of", idx, "frames finished") - - ''' - Save final_poses array into results folder as a pickle dump - ''' - filename = results_dir + get_new_filename() - print("Saving results to", filename) - with open(filename, "wb") as fp: - pickle.dump(final_poses, fp) - print("Results have been saved to", filename) + make_video(frames, video_name, fps) # TODO: use current body pose and camera transform for next optimization? @@ -150,37 +109,15 @@ def replay_animation(file, start_frame=0, end_frame=None, with_background=False, time.sleep(1 / fps) -def video_from_pkl(filename, video_name): - with open(filename, "rb") as fp: - final_poses = pickle.load(fp) - - save_to_video(final_poses, video_name) - - -def save_to_video(poses, video_name, fps=30): - r = DefaultRenderer( - offscreen=True - ) - r.start() - - model_anim = SMPLyModel.model_from_conf(config) - - frames = [] - - for body_pose, cam_trans in tqdm(poses): - r.render_model_with_tfs(model_anim, body_pose, keep_pose=True, - render_joints=False, transforms=cam_trans) - frames.append(r.get_snapshot()) - - make_video(frames, video_name, fps) - - ''' Play the animation. ''' -anim_file = results_dir + result_prefix + "0.pkl" if RUN_OPTIMIZATION: anim_file = filename +else: + results_dir = config['output']['rootDir'] + result_prefix = config['output']['prefix'] + anim_file = results_dir + result_prefix + "3.pkl" -video_from_pkl(anim_file, "test-anim.avi") +video_from_pkl(anim_file, "test-anim.avi", config) replay_animation(anim_file) diff --git a/modules/angle.py b/modules/angle.py index 1815694..8f82aba 100644 --- a/modules/angle.py +++ b/modules/angle.py @@ -40,5 +40,4 @@ class AnglePriorsLoss(nn.Module): angles = pose[:, self.angle_idx] # compute cost based not exponential of angle * direction - # then use MSE for cost - return torch.exp(angles * self.angle_directions).pow(2).sum() + return torch.exp(angles * self.angle_directions).pow(2).sum() diff --git a/modules/angle_clip.py b/modules/angle_clip.py index 476e529..2a1b4d2 100644 --- a/modules/angle_clip.py +++ b/modules/angle_clip.py @@ -10,7 +10,7 @@ class AngleClipper(nn.Module): dtype=torch.float32, angle_idx=[24, 10, 9], # directions=[-1, 1, 1, 1], - weights=[1.0, 1.0, 1.0] + weight=0.01 ): super(AngleClipper, self).__init__() @@ -28,8 +28,8 @@ class AngleClipper(nn.Module): # create buffer for weights self.register_buffer( - "weights", - torch.tensor(weights, dtype=dtype).to(device=device) + "weight", + torch.tensor(weight, dtype=dtype).to(device=device) ) def forward(self, pose): @@ -39,4 +39,4 @@ class AngleClipper(nn.Module): penalty = angles[torch.abs(angles) > self.limit] # get relevant angles - return penalty.pow(2).sum() * 0.01 + return penalty.pow(2).sum() * self.weight diff --git a/renderer.py b/renderer.py index 6d7cce9..0c1cf8b 100644 --- a/renderer.py +++ b/renderer.py @@ -316,6 +316,10 @@ class Renderer: return color + def wait_for_close(self): + while self.viewer.is_active: + pass + class DefaultRenderer(Renderer): """Utility class for easier default renderer setup diff --git a/train.py b/train.py new file mode 100644 index 0000000..b2fad4d --- /dev/null +++ b/train.py @@ -0,0 +1,101 @@ +# library imports +import pickle +import torch +from utils.video import make_video +from tqdm.auto import trange + +# local imports +from train_pose import train_pose_with_conf +from model import SMPLyModel +from utils.general import get_new_filename, setup_training +from camera_estimation import TorchCameraEstimate + + +def optimize_sample(sample_index, dataset, config, device=torch.device('cpu'), dtype=torch.float32, offscreen=False, verbose=False, display_result=False): + # prepare data and SMPL model + model = SMPLyModel.model_from_conf(config) + init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training( + model=model, + renderer=True, + dataset=dataset, + sample_index=sample_index, + offscreen=offscreen + ) + + camera = TorchCameraEstimate( + model, + keypoints=keypoints, + renderer=r, + device=device, + dtype=dtype, + image_path=img_path, + est_scale=est_scale, + verbose=verbose, + use_progress_bar=verbose + ) + + camera_transformation, camera_int, camera_params = camera.get_results() + + if not offscreen: + # render camera to the scene + camera.setup_visualization(r.init_keypoints, r.keypoints) + + # train for pose + pose, loss_history, step_imgs = train_pose_with_conf( + config=config, + model=model, + keypoints=keypoints, + keypoint_conf=conf, + camera=camera, + renderer=r, + device=device, + use_progress_bar=verbose + ) + + if display_result: + r.wait_for_close() + + return pose, camera_transformation, loss_history, step_imgs + + +def create_animation(dataset, config, start_idx=0, end_idx=None, device=torch.device('cpu'), dtype=torch.float32, offscreen=False, verbose=False, save_to_file=False): + final_poses = [] + + if end_idx is None: + end_idx = len(dataset) + + for idx in trange(end_idx - start_idx, desc='Optimizing'): + idx = start_idx + idx + + final_pose, cam_trans, train_loss, step_imgs = optimize_sample( + idx, + dataset, + config, + offscreen=True + ) + + if verbose: + print("Optimization of", idx, "frames finished") + + # print("\nPose optimization of frame", idx, "is finished.") + R = cam_trans.numpy().squeeze() + idx += 1 + + # append optimized pose and camera transformation to the array + final_poses.append((final_pose, R)) + + filename = None + + if save_to_file: + ''' + Save final_poses array into results folder as a pickle dump + ''' + results_dir = config['output']['rootDir'] + result_prefix = config['output']['prefix'] + filename = results_dir + get_new_filename() + print("Saving results to", filename) + with open(filename, "wb") as fp: + pickle.dump(final_poses, fp) + print("Results have been saved to", filename) + + return final_poses, filename diff --git a/train_pose.py b/train_pose.py index fccbd3b..477deaf 100644 --- a/train_pose.py +++ b/train_pose.py @@ -53,7 +53,6 @@ def train_pose( # renderer options renderer: Renderer = None, render_steps=True, - render_offscreen=True, vposer=None, @@ -209,14 +208,14 @@ def train_pose( renderer.render_model_with_tfs( model, pose_layer.cur_out, keep_pose=True, transforms=R) - if render_offscreen: + if renderer.use_offscreen: offscreen_step_output.append(renderer.get_snapshot()) # renderer.set_group_pose("body", R) if use_progress_bar: pbar.close() print("Final result:", loss.item()) - return pose_layer.cur_out, best_pose, loss_history, offscreen_step_output + return best_pose, loss_history, offscreen_step_output def train_pose_with_conf( diff --git a/utils/render.py b/utils/render.py index e54fa68..d45f8a7 100644 --- a/utils/render.py +++ b/utils/render.py @@ -1,28 +1,8 @@ + from typing import List, Set, Dict, Tuple, Optional import numpy as np import trimesh import pyrender -import cv2 - -from tqdm import tqdm - - -def make_video(images, video_name: str, fps=5): - - images = np.array(images) - width = images.shape[2] - height = images.shape[1] - video = cv2.VideoWriter( - video_name, 0, fps, (width, height), True) - - print("creating video with size", width, height) - - for idx in tqdm(range(len(images))): - img = images[idx] - im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - video.write(im_rgb) - - video.release() def render_model( diff --git a/utils/video.py b/utils/video.py new file mode 100644 index 0000000..a6f3025 --- /dev/null +++ b/utils/video.py @@ -0,0 +1,48 @@ +import pickle +from model import SMPLyModel +from renderer import DefaultRenderer +import cv2 +from tqdm import tqdm +import numpy as np + + +def make_video(images, video_name: str, fps=5): + images = np.array(images) + width = images.shape[2] + height = images.shape[1] + video = cv2.VideoWriter( + video_name, 0, fps, (width, height), True) + + print("creating video with size", width, height) + + for idx in tqdm(range(len(images))): + img = images[idx] + im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + video.write(im_rgb) + + video.release() + + +def video_from_pkl(filename, video_name, config): + with open(filename, "rb") as fp: + final_poses = pickle.load(fp) + + save_to_video(final_poses, video_name, config) + + +def save_to_video(poses, video_name, config, fps=30): + r = DefaultRenderer( + offscreen=True + ) + r.start() + + model_anim = SMPLyModel.model_from_conf(config) + + frames = [] + + for body_pose, cam_trans in tqdm(poses): + r.render_model_with_tfs(model_anim, body_pose, keep_pose=True, + render_joints=False, transforms=cam_trans) + frames.append(r.get_snapshot()) + + make_video(frames, video_name, fps)