vhad/body-pose-animation
1
0
Fork 0
mirror of https://github.com/gosticks/body-pose-animation.git synced 2025-10-16 11:45:42 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: minor formatting issues

Browse Source
This commit is contained in:
Wlad 2021-02-08 13:30:42 +01:00
parent ec7261621d
commit dbcab9cffc
9 changed files with 112 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@ -98,4 +98,5 @@ reference
vposer_v1_0
*.avi
results/
results/
output/

66
camera_estimation.py
View File

@ -25,7 +25,11 @@ class CameraEstimate:
image_path=None,
dtype=torch.float32,
device=torch.device("cpu"),
verbose=True,
use_progress_bar=True,
est_scale=1):
self.use_progress_bar = use_progress_bar
self.verbose = verbose
self.model = model
self.dataset = dataset
self.output_model = model(return_verts=True)
@ -45,7 +49,8 @@ class CameraEstimate:
return torso_keypoints_2d, torso_keypoints_3d
def visualize_mesh(self, keypoints, smpl_points):
if self.renderer is None:
return
# hardcoded scaling factor
# scaling_factor = 1
# smpl_points /= scaling_factor
@ -82,8 +87,9 @@ class CameraEstimate:
current_pose = self.params_to_pose(params)
# TODO: use renderer.py methods
self.renderer.scene.set_group_pose("body", current_pose)
# self.renderer.scene.set_pose(self.verts, current_pose)
if self.renderer is not None:
self.renderer.scene.set_group_pose("body", current_pose)
# self.renderer.scene.set_pose(self.verts, current_pose)
def params_to_pose(self, params):
pose = np.eye(4)
@ -96,7 +102,7 @@ class CameraEstimate:
translation = np.zeros(3)
rotation = np.random.rand(3) * 2 * np.pi
params = np.concatenate((translation, rotation))
print(params)
# print(params)
init_points_2d, init_points_3d = self.get_torso_keypoints()
@ -104,7 +110,7 @@ class CameraEstimate:
res = minimize(self.sum_of_squares, x0=params, args=(init_points_3d, init_points_2d),
callback=self.iteration_callback, tol=1e-4, method="BFGS")
print(res)
# print(res)
transform_matrix = self.params_to_pose(res.x)
return transform_matrix
@ -141,8 +147,10 @@ class TorchCameraEstimate(CameraEstimate):
stop = True
tol = 3e-4
print("Estimating Initial transform...")
pbar = tqdm(total=100)
if self.verbose:
print("Estimating Initial transform...")
if self.use_progress_bar:
pbar = tqdm(total=100)
current = 0
while stop:
y_pred = self.C(params, init_points_3d_prepared)
@ -159,18 +167,22 @@ class TorchCameraEstimate(CameraEstimate):
if self.renderer is not None:
self.renderer.set_group_pose("body", current_pose)
per = int((tol/loss*100).item())
if per > 100:
pbar.update(abs(100 - current))
current = 100
else:
pbar.update(per - current)
current = per
if self.use_progress_bar:
if per > 100:
pbar.update(abs(100 - current))
current = 100
else:
pbar.update(per - current)
current = per
stop = loss > tol
if stop == True:
stop = self.patience_module(loss, 5)
pbar.update(abs(100 - current))
pbar.close()
if self.use_progress_bar:
pbar.update(abs(100 - current))
pbar.close()
self.memory = None
transform_matrix = self.torch_params_to_pose(params)
current_pose = transform_matrix.detach().numpy()
@ -180,7 +192,7 @@ class TorchCameraEstimate(CameraEstimate):
# camera_translation[0,2] = 5 * torch.ones(1)
camera_rotation = torch.tensor(
[[0,0,0]], requires_grad=False, dtype=self.dtype, device=self.device)
[[0, 0, 0]], requires_grad=False, dtype=self.dtype, device=self.device)
camera_intrinsics = torch.zeros(
4, 4, dtype=self.dtype, device=self.device)
camera_intrinsics[0, 0] = 5
@ -205,8 +217,9 @@ class TorchCameraEstimate(CameraEstimate):
stop = True
first = True
cam_tol = 6e-3
print("Estimating Camera transformations...")
pbar = tqdm(total=100)
# print("Estimating Camera transformations...")
if self.use_progress_bar:
pbar = tqdm(total=100)
current = 0
while stop:
@ -221,15 +234,17 @@ class TorchCameraEstimate(CameraEstimate):
else:
loss.backward()
opt2.step()
if visualize:
if visualize and self.renderer is not None:
self.renderer.scene.set_pose(
self.camera_renderer, self.torch_params_to_pose(params).detach().numpy())
per = int((cam_tol/loss*100).item())
if per > 100:
pbar.update(100 - current)
else:
pbar.update(per - current)
if self.use_progress_bar:
if per > 100:
pbar.update(100 - current)
else:
pbar.update(per - current)
current = per
stop = loss > cam_tol
@ -237,8 +252,9 @@ class TorchCameraEstimate(CameraEstimate):
if stop == True:
stop = self.patience_module(loss, 5)
pbar.update(100 - current)
pbar.close()
if self.use_progress_bar:
pbar.update(100 - current)
pbar.close()
camera_transform_matrix = self.torch_params_to_pose(
params)
return camera_intrinsics, transform_matrix, camera_transform_matrix

6
config.yaml
View File

@ -9,10 +9,10 @@ smpl:
useVposerInit: false
data:
renameFiles: false
rootDir: ./samples
rootDir: ./output
personId: 0
sampleImageFormat: "%%%.png"
sampleNameFormat: "%%%.json"
sampleImageFormat: "input_%%%%%%%%%%%%_rendered.png"
sampleNameFormat: "input_%%%%%%%%%%%%_keypoints.json"
sampleCoords: !!python/tuple [1080, 1080]
camera:
lr: 0.001

16
dataset.py
View File

@ -16,6 +16,7 @@ class SMPLyDataset(Dataset):
model_type="smplx",
person_id=0,
sample_format="%%%.json",
image_format="%%%.png",
start_index=1,
sample_id_pad=None
):
@ -23,6 +24,8 @@ class SMPLyDataset(Dataset):
self.model_type = model_type
self.size = size
self.person_id = person_id
self.image_format = image_format
self.sample_format = sample_format
if sample_id_pad:
self.sample_id_pad = sample_format.count('%')
else:
@ -43,6 +46,8 @@ class SMPLyDataset(Dataset):
size=config['data']['sampleCoords'],
person_id=config['data']['personId'],
model_type=config['smpl']['type'],
image_format=config['data']['sampleImageFormat'],
sample_format=config['data']['sampleNameFormat'],
# img_format=config['data']['sampleImageFormat'],
sample_id_pad=config['data']['sampleImageFormat'].count('%')
)
@ -55,8 +60,12 @@ class SMPLyDataset(Dataset):
index + self.start_index).zfill(self.sample_id_pad)
return name
def get_keypoint_name(self, index):
id = self.get_item_name(index)
return self.sample_format.replace("%" * len(id), id)
def __getitem__(self, index):
name = self.get_item_name(index) + ".json"
name = self.get_keypoint_name(index)
path = os.path.join(
self.root_dir, name)
if os.path.exists(path):
@ -69,7 +78,7 @@ class SMPLyDataset(Dataset):
return None
def transform(self, data, origin_format="body_25"):
"""
"""
transform: transforms the order of an origin array to the target format
"""
@ -85,6 +94,7 @@ class SMPLyDataset(Dataset):
return sample_count
def get_image_path(self, index):
name = self.get_item_name(index) + ".png"
id = self.get_item_name(index)
name = self.image_format.replace("%" * len(id), id)
path = os.path.join(self.root_dir, name)
return path

28
example_fit_anim.py
View File

@ -2,7 +2,7 @@ import pickle
import time
from utils.render import make_video
import torch
from tqdm.std import tqdm
from tqdm.auto import trange
from dataset import SMPLyDataset
from model import *
@ -16,7 +16,7 @@ from utils.general import rename_files, get_new_filename
START_IDX = 150 # starting index of the frame to optimize for
FINISH_IDX = 200 # choose a big number to optimize for all frames in samples directory
# if False, only run already saved animation without optimization
RUN_OPTIMIZATION = False
RUN_OPTIMIZATION = True
final_poses = [] # optimized poses array that is saved for playing the animation
result_image = []
@ -60,7 +60,7 @@ joints = model_out.joints.detach().cpu().numpy().squeeze()
Optimization part without visualization
'''
if RUN_OPTIMIZATION:
for idx in dataset:
for idx in trange(100, desc='Optimizing'):
init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training(
model=model,
@ -70,7 +70,9 @@ if RUN_OPTIMIZATION:
sample_index=idx
)
camera = TorchCameraEstimate(
r.start()
cam = TorchCameraEstimate(
model,
dataset=dataset,
keypoints=keypoints,
@ -78,24 +80,28 @@ if RUN_OPTIMIZATION:
device=torch.device('cpu'),
dtype=torch.float32,
image_path=img_path,
est_scale=est_scale
est_scale=est_scale,
use_progress_bar=False,
verbose=False
)
print("\nCamera optimization of frame", idx, "is finished.")
camera = SimpleCamera.from_estimation_cam(camera)
# print("\nCamera optimization of frame", idx, "is finished.")
cur_pose, final_pose, loss, frames = train_pose_with_conf(
config=config,
model=model,
keypoints=keypoints,
keypoint_conf=conf,
camera=camera,
camera=cam,
renderer=r,
device=device
device=device,
use_progress_bar=False
)
print("\nPose optimization of frame", idx, "is finished.")
R = camera.trans.numpy().squeeze()
camera_transformation, camera_int, camera_params = cam.get_results()
# print("\nPose optimization of frame", idx, "is finished.")
R = camera_transformation.numpy().squeeze()
idx += 1
# append optimized pose and camera transformation to the array

19
model.py
View File

@ -8,7 +8,7 @@ from human_body_prior.tools.model_loader import load_vposer
def get_model_path(type, gender, dir):
return os.path.joint(
return os.path.join(
dir,
type,
type.upper() + "_" +
@ -24,6 +24,8 @@ def get_model_path_from_conf(config):
class VPoserModel():
global_vposer = None
def __init__(
self,
model_type='smpl',
@ -70,14 +72,17 @@ class VPoserModel():
def get_pose(self):
return self.model.pose_body
def from_conf(config):
def from_conf(config, use_global=True):
model_path = get_model_path_from_conf(config)
return VPoserModel(
model_type=config['smpl']['type'],
gender=config['smpl']['gender'],
vposer_model_path=config['pose']['vposerPath'],
body_model_path=model_path)
if VPoserModel.global_vposer is None:
VPoserModel.global_vposer = VPoserModel(
model_type=config['smpl']['type'],
gender=config['smpl']['gender'],
vposer_model_path=config['pose']['vposerPath'],
body_model_path=model_path)
return VPoserModel.global_vposer
class SMPLyModel():

34
train_pose.py
View File

@ -53,11 +53,14 @@ def train_pose(
# renderer options
renderer: Renderer = None,
render_steps=True,
render_offscreen=True
):
render_offscreen=True,
print("[pose] starting training")
print("[pose] dtype=", dtype)
vposer=None,
use_progress_bar=True
):
# print("[pose] starting training")
# print("[pose] dtype=", dtype)
offscreen_step_output = []
@ -85,7 +88,6 @@ def train_pose(
# loss layers
if useBodyPrior:
vposer = VPoserModel()
# TODO: handle this in vposer model
vposer.model.to(device=device, dtype=dtype)
latent_body = vposer.get_pose()
@ -105,7 +107,8 @@ def train_pose(
optimizer = optimizer(parameters, learning_rate)
pbar = tqdm(total=iterations)
if use_progress_bar:
pbar = tqdm(total=iterations)
def predict():
# pose_extra = None
@ -197,8 +200,9 @@ def train_pose(
if patience == 0:
print("[train] aborted due to patience limit reached")
pbar.set_description("Error %f" % cur_loss)
pbar.update(1)
if use_progress_bar:
pbar.set_description("Error %f" % cur_loss)
pbar.update(1)
if renderer is not None and render_steps:
R = camera.trans.detach().cpu().numpy().squeeze()
@ -209,8 +213,9 @@ def train_pose(
offscreen_step_output.append(renderer.get_snapshot())
# renderer.set_group_pose("body", R)
pbar.close()
print("Final result:", loss.item())
if use_progress_bar:
pbar.close()
print("Final result:", loss.item())
return pose_layer.cur_out, best_pose, loss_history, offscreen_step_output
@ -223,7 +228,8 @@ def train_pose_with_conf(
device=torch.device('cpu'),
dtype=torch.float32,
renderer: Renderer = None,
render_steps=True
render_steps=True,
use_progress_bar=True
):
# configure PyTorch device and format
@ -245,6 +251,8 @@ def train_pose_with_conf(
if renderer is not None:
renderer.set_group_pose("body", cam_trans.cpu().numpy())
vposer = VPoserModel.from_conf(config)
return train_pose(
model=model.to(dtype=dtype),
keypoints=keypoints,
@ -259,11 +267,13 @@ def train_pose_with_conf(
learning_rate=config['pose']['lr'],
optimizer_type=config['pose']['optimizer'],
iterations=config['pose']['iterations'],
vposer=vposer,
body_prior_weight=config['pose']['bodyPrior']['weight'],
angle_prior_weight=config['pose']['anglePrior']['weight'],
body_mean_loss=config['pose']['bodyMeanLoss']['enabled'],
body_mean_weight=config['pose']['bodyMeanLoss']['weight'],
use_angle_sum_loss=config['pose']['angleSumLoss']['enabled'],
angle_sum_weight=config['pose']['angleSumLoss']['weight'],
render_steps=render_steps
render_steps=render_steps,
use_progress_bar=use_progress_bar
)

4
utils/general.py
View File

@ -140,8 +140,8 @@ def rename_files(dir):
def get_new_filename():
conf = load_config()
results_dir = conf['resultsPath']
result_prefix = conf['resultPrefix']
results_dir = conf['output']['rootDir']
result_prefix = conf['output']['prefix']
results = glob.glob(results_dir + "*.pkl")
if len(results) == 0:

1
utils/render.py
View File

@ -10,7 +10,6 @@ from tqdm import tqdm
def make_video(images, video_name: str, fps=5):
images = np.array(images)
print(images.shape)
width = images.shape[2]
height = images.shape[1]
video = cv2.VideoWriter(