apps/ICON.py

# -*- 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: [email protected]

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 torch
import lib.smplx as smplx
import numpy as np
from torch import nn
from skimage.transform import resize
import pytorch_lightning as pl

torch.backends.cudnn.benchmark = True


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(),
        )

        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.base_keys = ["smpl_verts", "smpl_faces"]
        self.feat_names = self.cfg.net.smpl_feats

        self.icon_keys = self.base_keys + [
            f"smpl_{feat_name}" for feat_name in self.feat_names
        ]
        self.keypoint_keys = self.base_keys + [
            f"smpl_{feat_name}" for feat_name in self.feat_names
        ]
        self.pamir_keys = [
            "voxel_verts", "voxel_faces", "pad_v_num", "pad_f_num"
        ]
        self.pifu_keys = []

        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=osp.join(
                osp.realpath(self.smpl_data.model_dir),
                f"{smpl_type}/{smpl_type}_kid_template.npy",
            ),
            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]})

        in_tensor_dict.update({
            k: batch[k] if k in batch.keys() else None
            for k in getattr(self, f"{self.prior_type}_keys")
        })

        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]})

        in_tensor_dict.update({
            k: batch[k] if k in batch.keys() else None
            for k in getattr(self, f"{self.prior_type}_keys")
        })
        
        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())
        smpl_cmap = self.smpl_data.cmap_smpl_vids(smpl_type)

        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):

        self.netG.eval()
        self.netG.training = False
        in_tensor_dict = {}

        # export paths
        mesh_name = batch["subject"][0]
        mesh_rot = batch["rotation"][0].item()

        self.export_dir = osp.join(self.cfg.results_path, self.cfg.name,
                                   "-".join(self.cfg.dataset.types), 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]})

        in_tensor_dict.update({
            k: batch[k] if k in batch.keys() else None
            for k in getattr(self, f"{self.prior_type}_keys")
        })

        if "T_normal_F" not in in_tensor_dict.keys(
        ) or "T_normal_B" not in in_tensor_dict.keys():

            # update the new T_normal_F/B
            self.render.load_meshes(
                batch["smpl_verts"] *
                torch.tensor([1.0, -1.0, 1.0]).to(self.device),
                batch["smpl_faces"])
            T_normal_F, T_noraml_B = self.render.get_rgb_image()
            in_tensor_dict.update({
                'T_normal_F': T_normal_F,
                'T_normal_B': T_noraml_B
            })

        with torch.no_grad():
            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)

        def tensor2arr(x):
            return (x[0].permute(1, 2, 0).detach().cpu().numpy() +
                    1.0) * 0.5 * 255.0

        # save inter results
        image = tensor2arr(in_tensor_dict["image"])
        smpl_F = tensor2arr(in_tensor_dict["T_normal_F"])
        smpl_B = tensor2arr(in_tensor_dict["T_normal_B"])
        image_inter = np.concatenate(self.tensor2image(512, inter[0]) +
                                     [smpl_F, smpl_B, image],
                                     axis=1)
        Image.fromarray((image_inter).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)

        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)
        chamfer, p2s = self.evaluator.calculate_chamfer_p2s(num_samples=1000)
        normal_consist = self.evaluator.calculate_normal_consist(
            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={
                "cape-easy": (0, 50),
                "cape-hard": (50, 100)
            },
        )

        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 *
                     255.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():
            if in_tensor_dict[name] is not None:
                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 *
                 255.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]})

        in_tensor_dict.update({
            k: batch[k] if k in batch.keys() else None
            for k in getattr(self, f"{self.prior_type}_keys")
        })

        with torch.no_grad():
            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