lib/dataset/EvalDataset.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: ps-license@tuebingen.mpg.de

import os
import os.path as osp

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import trimesh
from PIL import Image

from lib.common.render import Render
from lib.dataset.mesh_util import SMPLX, HoppeMesh, projection, rescale_smpl

cape_gender = {
    "male":
    ['00032', '00096', '00122', '00127', '00145', '00215', '02474', '03284', '03375',
     '03394'], "female": ['00134', '00159', '03223', '03331', '03383']
}


class EvalDataset:
    def __init__(self, cfg, device):

        self.root = cfg.root
        self.bsize = cfg.batch_size

        self.opt = cfg.dataset
        self.datasets = self.opt.types
        self.input_size = self.opt.input_size
        self.scales = self.opt.scales
        self.vol_res = cfg.vol_res

        # [(feat_name, channel_num),...]
        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_nml_dim = [item[1] for item in cfg.net.in_nml]

        self.in_total = self.in_geo + self.in_nml
        self.in_total_dim = self.in_geo_dim + self.in_nml_dim

        self.rotations = range(0, 360, 120)

        self.datasets_dict = {}

        for dataset_id, dataset in enumerate(self.datasets):

            dataset_dir = osp.join(self.root, dataset)

            mesh_dir = osp.join(dataset_dir, "scans")
            smplx_dir = osp.join(dataset_dir, "smplx")
            smpl_dir = osp.join(dataset_dir, "smpl")

            self.datasets_dict[dataset] = {
                "smplx_dir": smplx_dir,
                "smpl_dir": smpl_dir,
                "mesh_dir": mesh_dir,
                "scale": self.scales[dataset_id],
            }

            self.datasets_dict[dataset].update({
                "subjects":
                np.loadtxt(osp.join(dataset_dir, "all.txt"), dtype=str)
            })

        self.subject_list = self.get_subject_list()
        self.smplx = SMPLX()

        # PIL to tensor
        self.image_to_tensor = transforms.Compose([
            transforms.Resize(self.input_size),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])

        # PIL to tensor
        self.mask_to_tensor = transforms.Compose([
            transforms.Resize(self.input_size),
            transforms.ToTensor(),
            transforms.Normalize((0.0, ), (1.0, )),
        ])

        self.device = device
        self.render = Render(size=512, device=self.device)

    def render_normal(self, verts, faces):

        # render optimized mesh (normal, T_normal, image [-1,1])
        self.render.load_meshes(verts, faces)
        return self.render.get_image()

    def get_subject_list(self):

        subject_list = []

        for dataset in self.datasets:

            split_txt = ""

            if dataset == 'renderpeople':
                split_txt = osp.join(self.root, dataset, "loose.txt")
            elif dataset == 'cape':
                split_txt = osp.join(self.root, dataset, "pose.txt")

            if osp.exists(split_txt) and osp.getsize(split_txt) > 0:
                print(f"load from {split_txt}")
                subject_list += np.loadtxt(split_txt, dtype=str).tolist()

        return subject_list

    def __len__(self):
        return len(self.subject_list) * len(self.rotations)

    def __getitem__(self, index):

        rid = index % len(self.rotations)
        mid = index // len(self.rotations)

        rotation = self.rotations[rid]
        subject = self.subject_list[mid].split("/")[1]
        dataset = self.subject_list[mid].split("/")[0]
        render_folder = "/".join([dataset + f"_{self.opt.rotation_num}views", subject])

        if not osp.exists(osp.join(self.root, render_folder)):
            render_folder = "/".join([dataset + "_36views", subject])

        # setup paths
        data_dict = {
            "dataset": dataset,
            "subject": subject,
            "rotation": rotation,
            "scale": self.datasets_dict[dataset]["scale"],
            "calib_path": osp.join(self.root, render_folder, "calib", f"{rotation:03d}.txt"),
            "image_path": osp.join(self.root, render_folder, "render", f"{rotation:03d}.png"),
        }

        if dataset == "cape":
            data_dict.update({
                "mesh_path":
                osp.join(self.datasets_dict[dataset]["mesh_dir"], f"{subject}.obj"),
                "smpl_path":
                osp.join(self.datasets_dict[dataset]["smpl_dir"], f"{subject}.obj"),
            })
        else:

            data_dict.update({
                "mesh_path":
                osp.join(
                    self.datasets_dict[dataset]["mesh_dir"],
                    f"{subject}.obj",
                ),
                "smplx_path":
                osp.join(self.datasets_dict[dataset]["smplx_dir"], f"{subject}.obj"),
            })

        # load training data
        data_dict.update(self.load_calib(data_dict))

        # image/normal/depth loader
        for name, channel in zip(self.in_total, self.in_total_dim):

            if f"{name}_path" not in data_dict.keys():
                data_dict.update({
                    f"{name}_path":
                    osp.join(self.root, render_folder, name, f"{rotation:03d}.png")
                })

            # tensor update
            if os.path.exists(data_dict[f"{name}_path"]):
                data_dict.update({
                    name:
                    self.imagepath2tensor(data_dict[f"{name}_path"], channel, inv=False)
                })

        data_dict.update(self.load_mesh(data_dict))
        data_dict.update(self.load_smpl(data_dict))

        del data_dict["mesh"]

        return data_dict

    def imagepath2tensor(self, path, channel=3, inv=False):

        rgba = Image.open(path).convert("RGBA")

        # remove CAPE's noisy outliers using OpenCV's inpainting
        if "cape" in path and "T_" not in path:
            mask = cv2.imread(path.replace(path.split("/")[-2], "mask"), 0) > 1
            img = np.asarray(rgba)[:, :, :3]
            fill_mask = ((mask & (img.sum(axis=2) == 0))).astype(np.uint8)
            image = Image.fromarray(
                cv2.inpaint(img * mask[..., None], fill_mask, 3, cv2.INPAINT_TELEA)
            )
            mask = Image.fromarray(mask)
        else:
            mask = rgba.split()[-1]
            image = rgba.convert("RGB")

        image = self.image_to_tensor(image)
        mask = self.mask_to_tensor(mask)
        image = (image * mask)[:channel]

        return (image * (0.5 - inv) * 2.0).float()

    def load_calib(self, data_dict):
        calib_data = np.loadtxt(data_dict["calib_path"], dtype=float)
        extrinsic = calib_data[:4, :4]
        intrinsic = calib_data[4:8, :4]
        calib_mat = np.matmul(intrinsic, extrinsic)
        calib_mat = torch.from_numpy(calib_mat).float()
        return {"calib": calib_mat}

    def load_mesh(self, data_dict):

        mesh_path = data_dict["mesh_path"]
        scale = data_dict["scale"]

        scan_mesh = trimesh.load(mesh_path)
        verts = scan_mesh.vertices
        faces = scan_mesh.faces

        mesh = HoppeMesh(verts * scale, faces)

        return {
            "mesh": mesh,
            "verts": torch.as_tensor(verts * scale).float(),
            "faces": torch.as_tensor(faces).long(),
        }

    def load_smpl(self, data_dict):

        smpl_type = ("smplx" if ("smplx_path" in data_dict.keys()) else "smpl")

        smplx_verts = rescale_smpl(data_dict[f"{smpl_type}_path"], scale=100.0)
        smplx_faces = torch.as_tensor(getattr(self.smplx, f"{smpl_type}_faces")).long()
        smplx_verts = projection(smplx_verts, data_dict["calib"]).float()

        return_dict = {
            "smpl_verts": smplx_verts,
            "smpl_faces": smplx_faces,
        }

        return return_dict

    def depth_to_voxel(self, data_dict):

        data_dict["depth_F"] = transforms.Resize(self.vol_res)(data_dict["depth_F"])
        data_dict["depth_B"] = transforms.Resize(self.vol_res)(data_dict["depth_B"])

        depth_mask = (~torch.isnan(data_dict['depth_F']))
        depth_FB = torch.cat([data_dict['depth_F'], data_dict['depth_B']], dim=0)
        depth_FB[:, ~depth_mask[0]] = 0.

        # Important: index_long = depth_value - 1
        index_z = (((depth_FB + 1.) * 0.5 * self.vol_res) - 1).clip(0, self.vol_res -
                                                                    1).permute(1, 2, 0)
        index_z_ceil = torch.ceil(index_z).long()
        index_z_floor = torch.floor(index_z).long()
        index_z_frac = torch.frac(index_z)

        index_mask = index_z[..., 0] == torch.tensor(self.vol_res * 0.5 - 1).long()
        voxels = F.one_hot(index_z_ceil[..., 0], self.vol_res) * index_z_frac[..., 0] + \
            F.one_hot(index_z_floor[..., 0], self.vol_res) * (1.0-index_z_frac[..., 0]) + \
            F.one_hot(index_z_ceil[..., 1], self.vol_res) * index_z_frac[..., 1]+ \
            F.one_hot(index_z_floor[..., 1], self.vol_res) * (1.0 - index_z_frac[..., 1])

        voxels[index_mask] *= 0
        voxels = torch.flip(voxels, [2]).permute(2, 0, 1).float()    #[x-2, y-0, z-1]

        return {
            "depth_voxels": voxels.flip([
                0,
            ]).unsqueeze(0).to(self.device),
        }

    def render_depth(self, verts, faces):

        # render optimized mesh (normal, T_normal, image [-1,1])
        self.render.load_meshes(verts, faces)
        return self.render.get_image(type="depth")