luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

direct training experiments

Browse Source
This commit is contained in:
luni
2023-06-01 12:55:42 +00:00
parent f3c1e3308e
commit e23d097398
4 changed files with 231 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/direct_train.py
+60 -38
View File
@@ -13,7 +13,7 @@ args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
if args['debug']:
from configs.direct_1 import config
from configs.direct_7 import config
# from configs.ae_static_orig_2 import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
@@ -52,7 +52,7 @@ from pytorch_msssim import MS_SSIM
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127_2/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/direct_train.py')
artifact.add_file('losses/losses.py')
@@ -83,21 +83,23 @@ def train_step(batch):
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
if config['use_skinning']:
mesh_orig = batch['mesh_unposed']
uv_orig = batch['uv_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
else:
mesh_orig = batch['mesh_posed']
uv_orig = batch['uv_posed']
mesh_pred_posed = mesh_pred
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
loss = l2_mesh_loss + (1 - ms_ssim) * 1e3
optimizer.zero_grad()
loss.backward()
@@ -109,9 +111,13 @@ def train_step(batch):
model.eval()
with torch.inference_mode():
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
# l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio,
'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss,
'train/normal_loss': n_loss }
@@ -127,14 +133,21 @@ def valid_step(batch):
input = torch.cat([batch['uv_normals'], batch['uv_body_posed'], batch['uv_static']], axis=1)
else:
input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
pred = model([input, batch['uv_mask']])
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
if config['use_skinning']:
mesh_orig = batch['mesh_unposed']
uv_orig = batch['uv_unposed']
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
else:
mesh_orig = batch['mesh_posed']
uv_orig = batch['uv_posed']
mesh_pred_posed = mesh_pred
e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
@@ -143,8 +156,8 @@ def valid_step(batch):
l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
loss = l2_mesh_loss + (1 - ms_ssim) * 1e3
metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
@@ -165,9 +178,14 @@ def plot_everything(batch: dict, type: str):
pred = model([input, batch['uv_mask']]).detach()
mesh_pred = uv2mesh(pred).cpu().numpy()
mesh_orig = batch['mesh_unposed'].cpu().numpy()
if config['use_skinning']:
mesh_orig = batch['mesh_unposed'].cpu().numpy()
uv_orig = batch['uv_unposed']
else:
mesh_orig = batch['mesh_posed'].cpu().numpy()
uv_orig = batch['uv_posed']
orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1])
orig_img = torch.permute(uv_orig, [0,2,3,1])
uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
uv_normals_img = (torch.permute(batch['uv_normals'], [0,2,3,1]) + 1.1) / 2.2
@@ -178,20 +196,24 @@ def plot_everything(batch: dict, type: str):
F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
mask = batch['mesh_mask'].cpu().numpy().astype(bool)
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
if config['use_skinning']:
shape = batch['pose_shape'][:, 72:]
gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
for i in range(mesh_orig.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
for i in range(mesh_orig.shape[0]):
mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
else:
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
mesh_posed = mesh_pred
gt = [mesh_orig[i][mask[i], :] for i in range(pred.shape[0])]
gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
body = batch['mesh_body_posed'].cpu().numpy()