# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
import os
from lib.common.seg3d_lossless import Seg3dLossless
from lib.dataset.Evaluator import Evaluator
from lib.net import HGPIFuNet
from lib.common.train_util import *
from lib.common.render import Render
from lib.dataset.mesh_util import SMPLX, update_mesh_shape_prior_losses, get_visibility
import warnings
import logging
import torch
import lib.smplx as smplx
import numpy as np
from torch import nn
import os.path as osp
from skimage.transform import resize
import pytorch_lightning as pl
from huggingface_hub import cached_download
torch.backends.cudnn.benchmark = True
logging.getLogger("lightning").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")
class ICON(pl.LightningModule):
def __init__(self, cfg):
super(ICON, self).__init__()
self.cfg = cfg
self.batch_size = self.cfg.batch_size
self.lr_G = self.cfg.lr_G
self.use_sdf = cfg.sdf
self.prior_type = cfg.net.prior_type
self.mcube_res = cfg.mcube_res
self.clean_mesh_flag = cfg.clean_mesh
self.netG = HGPIFuNet(
self.cfg,
self.cfg.projection_mode,
error_term=nn.SmoothL1Loss() if self.use_sdf else nn.MSELoss(),
)
# TODO: replace the renderer from opengl to pytorch3d
self.evaluator = Evaluator(
device=torch.device(f"cuda:{self.cfg.gpus[0]}"))
self.resolutions = (
np.logspace(
start=5,
stop=np.log2(self.mcube_res),
base=2,
num=int(np.log2(self.mcube_res) - 4),
endpoint=True,
)
+ 1.0
)
self.resolutions = self.resolutions.astype(np.int16).tolist()
self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"]
self.pamir_keys = ["voxel_verts",
"voxel_faces", "pad_v_num", "pad_f_num"]
self.reconEngine = Seg3dLossless(
query_func=query_func,
b_min=[[-1.0, 1.0, -1.0]],
b_max=[[1.0, -1.0, 1.0]],
resolutions=self.resolutions,
align_corners=True,
balance_value=0.50,
device=torch.device(f"cuda:{self.cfg.test_gpus[0]}"),
visualize=False,
debug=False,
use_cuda_impl=False,
faster=True,
)
self.render = Render(
size=512, device=torch.device(f"cuda:{self.cfg.test_gpus[0]}")
)
self.smpl_data = SMPLX()
self.get_smpl_model = lambda smpl_type, gender, age, v_template: smplx.create(
self.smpl_data.model_dir,
kid_template_path=cached_download(osp.join(self.smpl_data.model_dir,
f"{smpl_type}/{smpl_type}_kid_template.npy"), use_auth_token=os.environ['ICON']),
model_type=smpl_type,
gender=gender,
age=age,
v_template=v_template,
use_face_contour=False,
ext="pkl",
)
self.in_geo = [item[0] for item in cfg.net.in_geo]
self.in_nml = [item[0] for item in cfg.net.in_nml]
self.in_geo_dim = [item[1] for item in cfg.net.in_geo]
self.in_total = self.in_geo + self.in_nml
self.smpl_dim = cfg.net.smpl_dim
self.export_dir = None
self.result_eval = {}
def get_progress_bar_dict(self):
tqdm_dict = super().get_progress_bar_dict()
if "v_num" in tqdm_dict:
del tqdm_dict["v_num"]
return tqdm_dict
# Training related
def configure_optimizers(self):
# set optimizer
weight_decay = self.cfg.weight_decay
momentum = self.cfg.momentum
optim_params_G = [
{"params": self.netG.if_regressor.parameters(), "lr": self.lr_G}
]
if self.cfg.net.use_filter:
optim_params_G.append(
{"params": self.netG.F_filter.parameters(), "lr": self.lr_G}
)
if self.cfg.net.prior_type == "pamir":
optim_params_G.append(
{"params": self.netG.ve.parameters(), "lr": self.lr_G}
)
if self.cfg.optim == "Adadelta":
optimizer_G = torch.optim.Adadelta(
optim_params_G, lr=self.lr_G, weight_decay=weight_decay
)
elif self.cfg.optim == "Adam":
optimizer_G = torch.optim.Adam(
optim_params_G, lr=self.lr_G, weight_decay=weight_decay
)
elif self.cfg.optim == "RMSprop":
optimizer_G = torch.optim.RMSprop(
optim_params_G,
lr=self.lr_G,
weight_decay=weight_decay,
momentum=momentum,
)
else:
raise NotImplementedError
# set scheduler
scheduler_G = torch.optim.lr_scheduler.MultiStepLR(
optimizer_G, milestones=self.cfg.schedule, gamma=self.cfg.gamma
)
return [optimizer_G], [scheduler_G]
def training_step(self, batch, batch_idx):
if not self.cfg.fast_dev:
export_cfg(self.logger, self.cfg)
self.netG.train()
in_tensor_dict = {
"sample": batch["samples_geo"].permute(0, 2, 1),
"calib": batch["calib"],
"label": batch["labels_geo"].unsqueeze(1),
}
for name in self.in_total:
in_tensor_dict.update({name: batch[name]})
if self.prior_type == "icon":
for key in self.icon_keys:
in_tensor_dict.update({key: batch[key]})
elif self.prior_type == "pamir":
for key in self.pamir_keys:
in_tensor_dict.update({key: batch[key]})
else:
pass
preds_G, error_G = self.netG(in_tensor_dict)
acc, iou, prec, recall = self.evaluator.calc_acc(
preds_G.flatten(),
in_tensor_dict["label"].flatten(),
0.5,
use_sdf=self.cfg.sdf,
)
# metrics processing
metrics_log = {
"train_loss": error_G.item(),
"train_acc": acc.item(),
"train_iou": iou.item(),
"train_prec": prec.item(),
"train_recall": recall.item(),
}
tf_log = tf_log_convert(metrics_log)
bar_log = bar_log_convert(metrics_log)
if batch_idx % int(self.cfg.freq_show_train) == 0:
with torch.no_grad():
self.render_func(in_tensor_dict, dataset="train")
metrics_return = {
k.replace("train_", ""): torch.tensor(v) for k, v in metrics_log.items()
}
metrics_return.update(
{"loss": error_G, "log": tf_log, "progress_bar": bar_log})
return metrics_return
def training_epoch_end(self, outputs):
if [] in outputs:
outputs = outputs[0]
# metrics processing
metrics_log = {
"train_avgloss": batch_mean(outputs, "loss"),
"train_avgiou": batch_mean(outputs, "iou"),
"train_avgprec": batch_mean(outputs, "prec"),
"train_avgrecall": batch_mean(outputs, "recall"),
"train_avgacc": batch_mean(outputs, "acc"),
}
tf_log = tf_log_convert(metrics_log)
return {"log": tf_log}
def validation_step(self, batch, batch_idx):
self.netG.eval()
self.netG.training = False
in_tensor_dict = {
"sample": batch["samples_geo"].permute(0, 2, 1),
"calib": batch["calib"],
"label": batch["labels_geo"].unsqueeze(1),
}
for name in self.in_total:
in_tensor_dict.update({name: batch[name]})
if self.prior_type == "icon":
for key in self.icon_keys:
in_tensor_dict.update({key: batch[key]})
elif self.prior_type == "pamir":
for key in self.pamir_keys:
in_tensor_dict.update({key: batch[key]})
else:
pass
preds_G, error_G = self.netG(in_tensor_dict)
acc, iou, prec, recall = self.evaluator.calc_acc(
preds_G.flatten(),
in_tensor_dict["label"].flatten(),
0.5,
use_sdf=self.cfg.sdf,
)
if batch_idx % int(self.cfg.freq_show_val) == 0:
with torch.no_grad():
self.render_func(in_tensor_dict, dataset="val", idx=batch_idx)
metrics_return = {
"val_loss": error_G,
"val_acc": acc,
"val_iou": iou,
"val_prec": prec,
"val_recall": recall,
}
return metrics_return
def validation_epoch_end(self, outputs):
# metrics processing
metrics_log = {
"val_avgloss": batch_mean(outputs, "val_loss"),
"val_avgacc": batch_mean(outputs, "val_acc"),
"val_avgiou": batch_mean(outputs, "val_iou"),
"val_avgprec": batch_mean(outputs, "val_prec"),
"val_avgrecall": batch_mean(outputs, "val_recall"),
}
tf_log = tf_log_convert(metrics_log)
return {"log": tf_log}
def compute_vis_cmap(self, smpl_type, smpl_verts, smpl_faces):
(xy, z) = torch.as_tensor(smpl_verts).split([2, 1], dim=1)
smpl_vis = get_visibility(xy, -z, torch.as_tensor(smpl_faces).long())
if smpl_type == "smpl":
smplx_ind = self.smpl_data.smpl2smplx(np.arange(smpl_vis.shape[0]))
else:
smplx_ind = np.arange(smpl_vis.shape[0])
smpl_cmap = self.smpl_data.get_smpl_mat(smplx_ind)
return {
"smpl_vis": smpl_vis.unsqueeze(0).to(self.device),
"smpl_cmap": smpl_cmap.unsqueeze(0).to(self.device),
"smpl_verts": smpl_verts.unsqueeze(0),
}
@torch.enable_grad()
def optim_body(self, in_tensor_dict, batch):
smpl_model = self.get_smpl_model(
batch["type"][0], batch["gender"][0], batch["age"][0], None
).to(self.device)
in_tensor_dict["smpl_faces"] = (
torch.tensor(smpl_model.faces.astype(np.int))
.long()
.unsqueeze(0)
.to(self.device)
)
# The optimizer and variables
optimed_pose = torch.tensor(
batch["body_pose"][0], device=self.device, requires_grad=True
) # [1,23,3,3]
optimed_trans = torch.tensor(
batch["transl"][0], device=self.device, requires_grad=True
) # [3]
optimed_betas = torch.tensor(
batch["betas"][0], device=self.device, requires_grad=True
) # [1,10]
optimed_orient = torch.tensor(
batch["global_orient"][0], device=self.device, requires_grad=True
) # [1,1,3,3]
optimizer_smpl = torch.optim.SGD(
[optimed_pose, optimed_trans, optimed_betas, optimed_orient],
lr=1e-3,
momentum=0.9,
)
scheduler_smpl = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer_smpl, mode="min", factor=0.5, verbose=0, min_lr=1e-5, patience=5
)
loop_smpl = range(50)
for i in loop_smpl:
optimizer_smpl.zero_grad()
# prior_loss, optimed_pose = dataset.vposer_prior(optimed_pose)
smpl_out = smpl_model(
betas=optimed_betas,
body_pose=optimed_pose,
global_orient=optimed_orient,
transl=optimed_trans,
return_verts=True,
)
smpl_verts = smpl_out.vertices[0] * 100.0
smpl_verts = projection(
smpl_verts, batch["calib"][0], format="tensor")
smpl_verts[:, 1] *= -1
# render optimized mesh (normal, T_normal, image [-1,1])
self.render.load_meshes(
smpl_verts, in_tensor_dict["smpl_faces"])
(
in_tensor_dict["T_normal_F"],
in_tensor_dict["T_normal_B"],
) = self.render.get_rgb_image()
T_mask_F, T_mask_B = self.render.get_silhouette_image()
with torch.no_grad():
(
in_tensor_dict["normal_F"],
in_tensor_dict["normal_B"],
) = self.netG.normal_filter(in_tensor_dict)
# mask = torch.abs(in_tensor['T_normal_F']).sum(dim=0, keepdims=True) > 0.0
diff_F_smpl = torch.abs(
in_tensor_dict["T_normal_F"] - in_tensor_dict["normal_F"]
)
diff_B_smpl = torch.abs(
in_tensor_dict["T_normal_B"] - in_tensor_dict["normal_B"]
)
loss = (diff_F_smpl + diff_B_smpl).mean()
# silhouette loss
smpl_arr = torch.cat([T_mask_F, T_mask_B], dim=-1)[0]
gt_arr = torch.cat(
[in_tensor_dict["normal_F"][0], in_tensor_dict["normal_B"][0]], dim=2
).permute(1, 2, 0)
gt_arr = ((gt_arr + 1.0) * 0.5).to(self.device)
bg_color = (
torch.Tensor([0.5, 0.5, 0.5]).unsqueeze(
0).unsqueeze(0).to(self.device)
)
gt_arr = ((gt_arr - bg_color).sum(dim=-1) != 0.0).float()
loss += torch.abs(smpl_arr - gt_arr).mean()
# Image.fromarray(((in_tensor_dict['T_normal_F'][0].permute(1,2,0)+1.0)*0.5*255.0).detach().cpu().numpy().astype(np.uint8)).show()
# loop_smpl.set_description(f"smpl = {loss:.3f}")
loss.backward(retain_graph=True)
optimizer_smpl.step()
scheduler_smpl.step(loss)
in_tensor_dict["smpl_verts"] = smpl_verts.unsqueeze(0)
in_tensor_dict.update(
self.compute_vis_cmap(
batch["type"][0],
in_tensor_dict["smpl_verts"][0],
in_tensor_dict["smpl_faces"][0],
)
)
features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
return features, inter, in_tensor_dict
@torch.enable_grad()
def optim_cloth(self, verts_pr, faces_pr, inter):
# convert from GT to SDF
verts_pr -= (self.resolutions[-1] - 1) / 2.0
verts_pr /= (self.resolutions[-1] - 1) / 2.0
losses = {
"cloth": {"weight": 5.0, "value": 0.0},
"edge": {"weight": 100.0, "value": 0.0},
"normal": {"weight": 0.2, "value": 0.0},
"laplacian": {"weight": 100.0, "value": 0.0},
"smpl": {"weight": 1.0, "value": 0.0},
"deform": {"weight": 20.0, "value": 0.0},
}
deform_verts = torch.full(
verts_pr.shape, 0.0, device=self.device, requires_grad=True
)
optimizer_cloth = torch.optim.SGD(
[deform_verts], lr=1e-1, momentum=0.9)
scheduler_cloth = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer_cloth, mode="min", factor=0.1, verbose=0, min_lr=1e-3, patience=5
)
# cloth optimization
loop_cloth = range(100)
for i in loop_cloth:
optimizer_cloth.zero_grad()
self.render.load_meshes(
verts_pr.unsqueeze(0).to(self.device),
faces_pr.unsqueeze(0).to(self.device).long(),
deform_verts,
)
P_normal_F, P_normal_B = self.render.get_rgb_image()
update_mesh_shape_prior_losses(self.render.mesh, losses)
diff_F_cloth = torch.abs(P_normal_F[0] - inter[:3])
diff_B_cloth = torch.abs(P_normal_B[0] - inter[3:])
losses["cloth"]["value"] = (diff_F_cloth + diff_B_cloth).mean()
losses["deform"]["value"] = torch.topk(
torch.abs(deform_verts.flatten()), 30
)[0].mean()
# Weighted sum of the losses
cloth_loss = torch.tensor(0.0, device=self.device)
pbar_desc = ""
for k in losses.keys():
if k != "smpl":
cloth_loss_per_cls = losses[k]["value"] * \
losses[k]["weight"]
pbar_desc += f"{k}: {cloth_loss_per_cls:.3f} | "
cloth_loss += cloth_loss_per_cls
# loop_cloth.set_description(pbar_desc)
cloth_loss.backward(retain_graph=True)
optimizer_cloth.step()
scheduler_cloth.step(cloth_loss)
# convert from GT to SDF
deform_verts = deform_verts.flatten().detach()
deform_verts[torch.topk(torch.abs(deform_verts), 30)[
1]] = deform_verts.mean()
deform_verts = deform_verts.view(-1, 3).cpu()
verts_pr += deform_verts
verts_pr *= (self.resolutions[-1] - 1) / 2.0
verts_pr += (self.resolutions[-1] - 1) / 2.0
return verts_pr
def test_step(self, batch, batch_idx):
# dict_keys(['dataset', 'subject', 'rotation', 'scale', 'calib',
# 'normal_F', 'normal_B', 'image', 'T_normal_F', 'T_normal_B',
# 'z-trans', 'verts', 'faces', 'samples_geo', 'labels_geo',
# 'smpl_verts', 'smpl_faces', 'smpl_vis', 'smpl_cmap', 'pts_signs',
# 'type', 'gender', 'age', 'body_pose', 'global_orient', 'betas', 'transl'])
if self.evaluator._normal_render is None:
self.evaluator.init_gl()
self.netG.eval()
self.netG.training = False
in_tensor_dict = {}
# export paths
mesh_name = batch["subject"][0]
mesh_rot = batch["rotation"][0].item()
ckpt_dir = self.cfg.name
for kid, key in enumerate(self.cfg.dataset.noise_type):
ckpt_dir += f"_{key}_{self.cfg.dataset.noise_scale[kid]}"
if self.cfg.optim_cloth:
ckpt_dir += "_optim_cloth"
if self.cfg.optim_body:
ckpt_dir += "_optim_body"
self.export_dir = osp.join(self.cfg.results_path, ckpt_dir, mesh_name)
os.makedirs(self.export_dir, exist_ok=True)
for name in self.in_total:
if name in batch.keys():
in_tensor_dict.update({name: batch[name]})
# update the new T_normal_F/B
in_tensor_dict.update(
self.evaluator.render_normal(
batch["smpl_verts"], batch["smpl_faces"])
)
# update the new smpl_vis
(xy, z) = batch["smpl_verts"][0].split([2, 1], dim=1)
smpl_vis = get_visibility(
xy,
z,
torch.as_tensor(self.smpl_data.faces).type_as(
batch["smpl_verts"]).long(),
)
in_tensor_dict.update({"smpl_vis": smpl_vis.unsqueeze(0)})
if self.prior_type == "icon":
for key in self.icon_keys:
in_tensor_dict.update({key: batch[key]})
elif self.prior_type == "pamir":
for key in self.pamir_keys:
in_tensor_dict.update({key: batch[key]})
else:
pass
with torch.no_grad():
if self.cfg.optim_body:
features, inter, in_tensor_dict = self.optim_body(
in_tensor_dict, batch)
else:
features, inter = self.netG.filter(
in_tensor_dict, return_inter=True)
sdf = self.reconEngine(
opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
)
# save inter results
image = (
in_tensor_dict["image"][0].permute(
1, 2, 0).detach().cpu().numpy() + 1.0
) * 0.5
smpl_F = (
in_tensor_dict["T_normal_F"][0].permute(
1, 2, 0).detach().cpu().numpy()
+ 1.0
) * 0.5
smpl_B = (
in_tensor_dict["T_normal_B"][0].permute(
1, 2, 0).detach().cpu().numpy()
+ 1.0
) * 0.5
image_inter = np.concatenate(
self.tensor2image(512, inter[0]) + [smpl_F, smpl_B, image], axis=1
)
Image.fromarray((image_inter * 255.0).astype(np.uint8)).save(
osp.join(self.export_dir, f"{mesh_rot}_inter.png")
)
verts_pr, faces_pr = self.reconEngine.export_mesh(sdf)
if self.clean_mesh_flag:
verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr)
if self.cfg.optim_cloth:
verts_pr = self.optim_cloth(verts_pr, faces_pr, inter[0].detach())
verts_gt = batch["verts"][0]
faces_gt = batch["faces"][0]
self.result_eval.update(
{
"verts_gt": verts_gt,
"faces_gt": faces_gt,
"verts_pr": verts_pr,
"faces_pr": faces_pr,
"recon_size": (self.resolutions[-1] - 1.0),
"calib": batch["calib"][0],
}
)
self.evaluator.set_mesh(self.result_eval, scale_factor=1.0)
self.evaluator.space_transfer()
chamfer, p2s = self.evaluator.calculate_chamfer_p2s(
sampled_points=1000)
normal_consist = self.evaluator.calculate_normal_consist(
save_demo_img=osp.join(self.export_dir, f"{mesh_rot}_nc.png")
)
test_log = {"chamfer": chamfer, "p2s": p2s, "NC": normal_consist}
return test_log
def test_epoch_end(self, outputs):
# make_test_gif("/".join(self.export_dir.split("/")[:-2]))
accu_outputs = accumulate(
outputs,
rot_num=3,
split={
"thuman2": (0, 5),
},
)
print(colored(self.cfg.name, "green"))
print(colored(self.cfg.dataset.noise_scale, "green"))
self.logger.experiment.add_hparams(
hparam_dict={"lr_G": self.lr_G, "bsize": self.batch_size},
metric_dict=accu_outputs,
)
np.save(
osp.join(self.export_dir, "../test_results.npy"),
accu_outputs,
allow_pickle=True,
)
return accu_outputs
def tensor2image(self, height, inter):
all = []
for dim in self.in_geo_dim:
img = resize(
np.tile(
((inter[:dim].cpu().numpy() + 1.0) /
2.0).transpose(1, 2, 0),
(1, 1, int(3 / dim)),
),
(height, height),
anti_aliasing=True,
)
all.append(img)
inter = inter[dim:]
return all
def render_func(self, in_tensor_dict, dataset="title", idx=0):
for name in in_tensor_dict.keys():
in_tensor_dict[name] = in_tensor_dict[name][0:1]
self.netG.eval()
features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
sdf = self.reconEngine(
opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
)
if sdf is not None:
render = self.reconEngine.display(sdf)
image_pred = np.flip(render[:, :, ::-1], axis=0)
height = image_pred.shape[0]
image_gt = resize(
((in_tensor_dict["image"].cpu().numpy()[0] + 1.0) / 2.0).transpose(
1, 2, 0
),
(height, height),
anti_aliasing=True,
)
image_inter = self.tensor2image(height, inter[0])
image = np.concatenate(
[image_pred, image_gt] + image_inter, axis=1)
step_id = self.global_step if dataset == "train" else self.global_step + idx
self.logger.experiment.add_image(
tag=f"Occupancy-{dataset}/{step_id}",
img_tensor=image.transpose(2, 0, 1),
global_step=step_id,
)
def test_single(self, batch):
self.netG.eval()
self.netG.training = False
in_tensor_dict = {}
for name in self.in_total:
if name in batch.keys():
in_tensor_dict.update({name: batch[name]})
if self.prior_type == "icon":
for key in self.icon_keys:
in_tensor_dict.update({key: batch[key]})
elif self.prior_type == "pamir":
for key in self.pamir_keys:
in_tensor_dict.update({key: batch[key]})
else:
pass
features, inter = self.netG.filter(in_tensor_dict, return_inter=True)
sdf = self.reconEngine(
opt=self.cfg, netG=self.netG, features=features, proj_matrix=None
)
verts_pr, faces_pr = self.reconEngine.export_mesh(sdf)
if self.clean_mesh_flag:
verts_pr, faces_pr = clean_mesh(verts_pr, faces_pr)
verts_pr -= (self.resolutions[-1] - 1) / 2.0
verts_pr /= (self.resolutions[-1] - 1) / 2.0
return verts_pr, faces_pr, inter