<reponame>leehsiu/neuralbody
import os
import glob

import cv2
import imageio
import scipy.io
import neuralbody.utils.geometry as geometry
import numpy as np
import torch

# AMASS dataset.
# https://amass.is.tue.mpg.de/


class AMASSData:
    def __init__(self, file_path):
        data = np.load(file_path)
        poses = data['poses']
        trans = data['trans']
        beta = data['betas']

        self.R = poses[..., :3].astype(np.float32)
        self.T = trans.astype(np.float32)

        self.theta = poses.astype(np.float32)
        self.theta[..., :3] = 0
        self.beta = beta.astype(np.float32)

    def __len__(self):
        return self.theta.shape[0]

    def __getitem__(self, idx):
        return self.smpl(idx)

    def smpl(self, i_frame):
        h2w_R = torch.tensor(self.R[i_frame])
        h2w_T = torch.tensor(self.T[i_frame])
        h2w_R = geometry.quat2mat(geometry.rodrigues(h2w_R))
        h2w = torch.cat([h2w_R, h2w_T[:, None]], dim=-1)
        h2w = torch.cat([h2w,torch.tensor([[0,0,0,1]])],dim=0).unsqueeze(0)

        pose = torch.tensor(self.theta[i_frame])
        pose = pose.reshape(-1,3).unsqueeze(0)
        shape = torch.tensor(self.beta).unsqueeze(0)

        return h2w, pose, shape

# ZJU mocap data
# https://github.com/zju3dv/neuralbody/blob/master/INSTALL.md
class MoCapData:
    def __init__(self, data_root):
        self.data_root = data_root
        annots = np.load(os.path.join(data_root, 'annots.npy'),
                         allow_pickle=True).item()

        self.camK = np.array(annots['cams']['K']).astype(np.float32)
        self.camR = np.array(annots['cams']['R']).astype(np.float32)
        self.camT = np.array(annots['cams']['T'])/1000.0
        self.camT = self.camT.astype(np.float32)
        self.n_cam = len(self.camK)
        self.n_frame = len(annots['ims'])

        self.imgs = [ims['ims'] for ims in annots['ims']]
        # load params
        self.R = []
        self.T = []
        self.beta = []
        self.theta = []

        for i in range(self.n_frame):
            dd = np.load(os.path.join(
                self.data_root, 'params/{}.npy'.format(i+1)), allow_pickle=True).item()
            self.R.append(np.array(dd['Rh'][0]))
            self.T.append(np.array(dd['Th'][0]))
            self.beta.append(np.array(dd['shapes'])[0])
            self.theta.append(np.reshape(dd['poses'], (24, 3)))
        self.R = np.array(self.R).astype(np.float32)
        self.T = np.array(self.T).astype(np.float32)
        self.beta = np.array(self.beta).astype(np.float32)
        self.theta = np.array(self.theta).astype(np.float32)

        print('{} frames with {} views loaded'.format(self.n_frame, self.n_cam))

    def camera(self, i_view, scale=1.0):
        invR = self.camR[i_view].T
        invT = np.matmul(invR, -self.camT[i_view])
        c2w = np.hstack([invR, invT])
        c2w = torch.tensor(c2w)

        camK = self.camK[i_view]*scale
        camK[2, 2] = 1.0
        camK = torch.tensor(camK)

        return camK, c2w

    def smpl(self, i_frame):
        h2w_R = torch.tensor(self.R[i_frame])
        h2w_T = torch.tensor(self.T[i_frame])
        h2w_R = geometry.quat2mat(geometry.rodrigues(h2w_R))
        h2w = torch.cat([h2w_R, h2w_T[:, None]], dim=-1)
        h2w = torch.cat([h2w,torch.tensor([[0,0,0,1]])],dim=0)[None,...]

        pose = torch.tensor(self.theta[i_frame])[None,...]
        shape = torch.tensor(self.beta[i_frame])[None,...]

        return h2w, pose, shape

    def image(self, i_frame, i_view, scale=-1):
        file_path = self.imgs[i_frame][i_view]
        img = imageio.imread(os.path.join(
            self.data_root, file_path)).astype(np.float32)/255.0
        if scale > 0:
            img = cv2.resize(img, dsize=None, fx=scale, fy=scale)
        return torch.tensor(img)

    def mask(self, i_frame, i_view, scale=-1):
        file_name = self.imgs[i_frame][i_view][:-4]
        mask = imageio.imread(os.path.join(
            self.data_root, 'mask', file_name+'.png'))
        mask = (mask != 0).astype(np.uint8)

        mask_cihp = imageio.imread(os.path.join(
            self.data_root, 'mask_cihp', file_name+'.png'))
        mask_cihp = (mask_cihp != 0).astype(np.uint8)
        mask = (mask | mask_cihp)
        if scale > 0:
            mask = cv2.resize(mask, dsize=None, fx=scale,
                              fy=scale, interpolation=cv2.INTER_NEAREST)
        return torch.tensor(mask)

    def __len__(self):
        return self.n_frame



# DoubleFusion
# https://arxiv.org/pdf/1804.06023.pdf
class DoubleFusion:
    def __init__(self, data_root):
        self.data_root = data_root
        data = scipy.io.loadmat(os.path.join(data_root, 'camera.mat'))
        _camK = np.eye(3)
        _camK[:2, :2] = np.diag(data['color_focal_length'][0])
        _camK[:2, 2] = data['color_center'][0]
        self.camK = _camK.astype(np.float32)
        self.c2w = data['c2w'].astype(np.float32)

        rgb_frames = glob.glob(os.path.join(data_root, 'rgb/*.png'))
        self.rgb_format = 'rgb_frame_{}.png'
        self.obj_format = 'smpl_frame_{}.obj'
        indices = []
        for fname in rgb_frames:
            fname = os.path.basename(fname)
            fname = os.path.splitext(fname)[0]
            idx = int(fname[10:])
            if os.path.isfile(
                    os.path.join(data_root, 'obj',
                                 self.obj_format.format(idx))):
                indices.append(idx)

        self.indices = sorted(indices)

    @staticmethod
    def load_obj(file_path):
        v = []
        with open(file_path, 'r') as f:
            lines = f.readlines()
        for line in lines:
            elem = line.split()
            if elem[0] == 'v':
                v.append(elem[1:])
        v = np.array(v).astype(np.float32)
        return v

    def camera(self, idx, scale=1.0):
        camK = self.camK*scale
        camK[2, 2] = 1.0
        return torch.tensor(camK), torch.tensor(self.c2w)

    def image(self, idx, scale=1.0):
        img_path = os.path.join(self.data_root, 'rgb',
                                self.rgb_format.format(self.indices[idx]))
        img = imageio.imread(img_path)
        img = img.astype(np.float32)/255.0

        if scale < 1.0:
            img = cv2.resize(img, dsize=None,
                             fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
        # the last channel is useless
        return torch.tensor(img[..., :3])

    def smpl_v(self, idx):

        obj_path = os.path.join(self.data_root, 'obj',
                                self.obj_format.format(self.indices[idx]))
        v = self.load_obj(obj_path)
        v = torch.tensor(v)

        return v

    def __len__(self):
        return len(self.indices)

import os
import glob

import cv2
import imageio
import scipy.io
import neuralbody.utils.geometry as geometry
import numpy as np
import torch

# AMASS dataset.
# https://amass.is.tue.mpg.de/


class AMASSData:
    def __init__(self, file_path):
        data = np.load(file_path)
        poses = data['poses']
        trans = data['trans']
        beta = data['betas']

        self.R = poses[..., :3].astype(np.float32)
        self.T = trans.astype(np.float32)

        self.theta = poses.astype(np.float32)
        self.theta[..., :3] = 0
        self.beta = beta.astype(np.float32)

    def __len__(self):
        return self.theta.shape[0]

    def __getitem__(self, idx):
        return self.smpl(idx)

    def smpl(self, i_frame):
        h2w_R = torch.tensor(self.R[i_frame])
        h2w_T = torch.tensor(self.T[i_frame])
        h2w_R = geometry.quat2mat(geometry.rodrigues(h2w_R))
        h2w = torch.cat([h2w_R, h2w_T[:, None]], dim=-1)
        h2w = torch.cat([h2w,torch.tensor([[0,0,0,1]])],dim=0).unsqueeze(0)

        pose = torch.tensor(self.theta[i_frame])
        pose = pose.reshape(-1,3).unsqueeze(0)
        shape = torch.tensor(self.beta).unsqueeze(0)

        return h2w, pose, shape

# ZJU mocap data
# https://github.com/zju3dv/neuralbody/blob/master/INSTALL.md
class MoCapData:
    def __init__(self, data_root):
        self.data_root = data_root
        annots = np.load(os.path.join(data_root, 'annots.npy'),
                         allow_pickle=True).item()

        self.camK = np.array(annots['cams']['K']).astype(np.float32)
        self.camR = np.array(annots['cams']['R']).astype(np.float32)
        self.camT = np.array(annots['cams']['T'])/1000.0
        self.camT = self.camT.astype(np.float32)
        self.n_cam = len(self.camK)
        self.n_frame = len(annots['ims'])

        self.imgs = [ims['ims'] for ims in annots['ims']]
        # load params
        self.R = []
        self.T = []
        self.beta = []
        self.theta = []

        for i in range(self.n_frame):
            dd = np.load(os.path.join(
                self.data_root, 'params/{}.npy'.format(i+1)), allow_pickle=True).item()
            self.R.append(np.array(dd['Rh'][0]))
            self.T.append(np.array(dd['Th'][0]))
            self.beta.append(np.array(dd['shapes'])[0])
            self.theta.append(np.reshape(dd['poses'], (24, 3)))
        self.R = np.array(self.R).astype(np.float32)
        self.T = np.array(self.T).astype(np.float32)
        self.beta = np.array(self.beta).astype(np.float32)
        self.theta = np.array(self.theta).astype(np.float32)

        print('{} frames with {} views loaded'.format(self.n_frame, self.n_cam))

    def camera(self, i_view, scale=1.0):
        invR = self.camR[i_view].T
        invT = np.matmul(invR, -self.camT[i_view])
        c2w = np.hstack([invR, invT])
        c2w = torch.tensor(c2w)

        camK = self.camK[i_view]*scale
        camK[2, 2] = 1.0
        camK = torch.tensor(camK)

        return camK, c2w

    def smpl(self, i_frame):
        h2w_R = torch.tensor(self.R[i_frame])
        h2w_T = torch.tensor(self.T[i_frame])
        h2w_R = geometry.quat2mat(geometry.rodrigues(h2w_R))
        h2w = torch.cat([h2w_R, h2w_T[:, None]], dim=-1)
        h2w = torch.cat([h2w,torch.tensor([[0,0,0,1]])],dim=0)[None,...]

        pose = torch.tensor(self.theta[i_frame])[None,...]
        shape = torch.tensor(self.beta[i_frame])[None,...]

        return h2w, pose, shape

    def image(self, i_frame, i_view, scale=-1):
        file_path = self.imgs[i_frame][i_view]
        img = imageio.imread(os.path.join(
            self.data_root, file_path)).astype(np.float32)/255.0
        if scale > 0:
            img = cv2.resize(img, dsize=None, fx=scale, fy=scale)
        return torch.tensor(img)

    def mask(self, i_frame, i_view, scale=-1):
        file_name = self.imgs[i_frame][i_view][:-4]
        mask = imageio.imread(os.path.join(
            self.data_root, 'mask', file_name+'.png'))
        mask = (mask != 0).astype(np.uint8)

        mask_cihp = imageio.imread(os.path.join(
            self.data_root, 'mask_cihp', file_name+'.png'))
        mask_cihp = (mask_cihp != 0).astype(np.uint8)
        mask = (mask | mask_cihp)
        if scale > 0:
            mask = cv2.resize(mask, dsize=None, fx=scale,
                              fy=scale, interpolation=cv2.INTER_NEAREST)
        return torch.tensor(mask)

    def __len__(self):
        return self.n_frame



# DoubleFusion
# https://arxiv.org/pdf/1804.06023.pdf
class DoubleFusion:
    def __init__(self, data_root):
        self.data_root = data_root
        data = scipy.io.loadmat(os.path.join(data_root, 'camera.mat'))
        _camK = np.eye(3)
        _camK[:2, :2] = np.diag(data['color_focal_length'][0])
        _camK[:2, 2] = data['color_center'][0]
        self.camK = _camK.astype(np.float32)
        self.c2w = data['c2w'].astype(np.float32)

        rgb_frames = glob.glob(os.path.join(data_root, 'rgb/*.png'))
        self.rgb_format = 'rgb_frame_{}.png'
        self.obj_format = 'smpl_frame_{}.obj'
        indices = []
        for fname in rgb_frames:
            fname = os.path.basename(fname)
            fname = os.path.splitext(fname)[0]
            idx = int(fname[10:])
            if os.path.isfile(
                    os.path.join(data_root, 'obj',
                                 self.obj_format.format(idx))):
                indices.append(idx)

        self.indices = sorted(indices)

    @staticmethod
    def load_obj(file_path):
        v = []
        with open(file_path, 'r') as f:
            lines = f.readlines()
        for line in lines:
            elem = line.split()
            if elem[0] == 'v':
                v.append(elem[1:])
        v = np.array(v).astype(np.float32)
        return v

    def camera(self, idx, scale=1.0):
        camK = self.camK*scale
        camK[2, 2] = 1.0
        return torch.tensor(camK), torch.tensor(self.c2w)

    def image(self, idx, scale=1.0):
        img_path = os.path.join(self.data_root, 'rgb',
                                self.rgb_format.format(self.indices[idx]))
        img = imageio.imread(img_path)
        img = img.astype(np.float32)/255.0

        if scale < 1.0:
            img = cv2.resize(img, dsize=None,
                             fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
        # the last channel is useless
        return torch.tensor(img[..., :3])

    def smpl_v(self, idx):

        obj_path = os.path.join(self.data_root, 'obj',
                                self.obj_format.format(self.indices[idx]))
        v = self.load_obj(obj_path)
        v = torch.tensor(v)

        return v

    def __len__(self):
        return len(self.indices)
# AMASS dataset. # https://amass.is.tue.mpg.de/ # ZJU mocap data # https://github.com/zju3dv/neuralbody/blob/master/INSTALL.md # load params # DoubleFusion # https://arxiv.org/pdf/1804.06023.pdf # the last channel is useless