recon/dataloader/dataset_recon.py

# Dataloader for training Portrait4D, modified from EG3D: https://github.com/NVlabs/eg3d

# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

"""Streaming images and labels from datasets created with dataset_tool.py."""

import os
import sys
import numpy as np
import zipfile
import PIL.Image
import json
import torch
import dnnlib
from training.dataloader.protocols import datum_portrait_genhead_pb2 as datum_pb2_genhead
from training.dataloader.protocols import datum_portrait_genhead_static_pb2 as datum_pb2_genhead_static
from training.dataloader.protocols import datum_portrait_genhead_new_pb2 as datum_pb2_genhead_new
from training.dataloader.protocols import datum_portrait_genhead_static_new_pb2 as datum_pb2_genhead_static_new
from training.dataloader.protocols import datum_portrait_vfhq_pb2 as datum_pb2_vfhq
from training.dataloader.protocols import datum_portrait_ffhq_pb2 as datum_pb2_ffhq
import lmdb
import cv2
from scipy.spatial.transform import Rotation as R
# try:
#     import pyspng
# except ImportError:
pyspng = None

#----------------------------------------------------------------------------

class Dataset(torch.utils.data.Dataset):
    def __init__(self,
        name,                   # Name of the dataset.
        raw_shape,              # Shape of the raw image data (NCHW).
        max_size    = None,     # Artificially limit the size of the dataset. None = no limit. Applied before xflip.
        use_labels  = False,    # Enable conditioning labels? False = label dimension is zero.
        xflip       = False,    # Artificially double the size of the dataset via x-flips. Applied after max_size.
        random_seed = 0,        # Random seed to use when applying max_size.
    ):
        self._name = name
        self._raw_shape = list(raw_shape)
        self._use_labels = use_labels
        self._raw_labels = None
        self._label_shape = None

        # Apply max_size.
        self._raw_idx = np.arange(self._raw_shape[0], dtype=np.int64)
        if (max_size is not None) and (self._raw_idx.size > max_size):
            np.random.RandomState(random_seed).shuffle(self._raw_idx)
            self._raw_idx = np.sort(self._raw_idx[:max_size])

        # Apply xflip.
        self._xflip = np.zeros(self._raw_idx.size, dtype=np.uint8)
        if xflip:
            self._raw_idx = np.tile(self._raw_idx, 2)
            self._xflip = np.concatenate([self._xflip, np.ones_like(self._xflip)])

    def _get_raw_labels(self):
        if self._raw_labels is None:
            self._raw_labels = self._load_raw_labels() if self._use_labels else None
            if self._raw_labels is None:
                self._raw_labels = np.zeros([self._raw_shape[0], 0], dtype=np.float32)
            assert isinstance(self._raw_labels, np.ndarray)
            assert self._raw_labels.shape[0] == self._raw_shape[0]
            assert self._raw_labels.dtype in [np.float32, np.int64]
            if self._raw_labels.dtype == np.int64:
                assert self._raw_labels.ndim == 1
                assert np.all(self._raw_labels >= 0)
            self._raw_labels_std = self._raw_labels.std(0)
        return self._raw_labels

    def close(self): # to be overridden by subclass
        pass

    def _load_raw_image(self, raw_idx): # to be overridden by subclass
        raise NotImplementedError

    def _load_raw_labels(self): # to be overridden by subclass
        raise NotImplementedError

    def __getstate__(self):
        return dict(self.__dict__, _raw_labels=None)

    def __del__(self):
        try:
            self.close()
        except:
            pass

    def __len__(self):
        return self._raw_idx.size

    def __getitem__(self, idx):
        image = self._load_raw_image(self._raw_idx[idx])
        assert isinstance(image, np.ndarray)
        assert list(image.shape) == self.image_shape
        assert image.dtype == np.uint8
        if self._xflip[idx]:
            assert image.ndim == 3 # CHW
            image = image[:, :, ::-1]
        return image.copy(), self.get_label(idx)

    def get_label(self, idx):
        label = self._get_raw_labels()[self._raw_idx[idx]]
        if label.dtype == np.int64:
            onehot = np.zeros(self.label_shape, dtype=np.float32)
            onehot[label] = 1
            label = onehot
        return label.copy()

    def get_details(self, idx):
        d = dnnlib.EasyDict()
        d.raw_idx = int(self._raw_idx[idx])
        d.xflip = (int(self._xflip[idx]) != 0)
        d.raw_label = self._get_raw_labels()[d.raw_idx].copy()
        return d

    def get_label_std(self):
        return self._raw_labels_std

    @property
    def name(self):
        return self._name

    @property
    def image_shape(self):
        return list(self._raw_shape[1:])

    @property
    def num_channels(self):
        assert len(self.image_shape) == 3 # CHW
        return self.image_shape[0]

    @property
    def resolution(self):
        assert len(self.image_shape) == 3 # CHW
        assert self.image_shape[1] == self.image_shape[2]
        return self.image_shape[1]

    @property
    def label_shape(self):
        if self._label_shape is None:
            raw_labels = self._get_raw_labels()
            if raw_labels.dtype == np.int64:
                self._label_shape = [int(np.max(raw_labels)) + 1]
            else:
                self._label_shape = raw_labels.shape[1:]
        return list(self._label_shape)

    @property
    def label_dim(self):
        assert len(self.label_shape) == 1
        return self.label_shape[0]

    @property
    def has_labels(self):
        return any(x != 0 for x in self.label_shape)

    @property
    def has_onehot_labels(self):
        return self._get_raw_labels().dtype == np.int64

#----------------------------------------------------------------------------
# Dataloader for real video set
class OneShotReconSegLmdbFolderDataset(Dataset):
    def __init__(self,
        path,                   # Path to datalist.
        resolution      = None, # Ensure specific resolution, None = highest available.
        data_type      = "vfhq",# Deprecated
        static         = False, # if True, only select multiview static frames instead of frames with different motions
        use_flame_mot  = False, # whether or not to use FLAME parameters as motion embedding
        max_num = None,
        rescale_camera = False, # Rescale camera extrinsics and intrinscs to align with an older version of camera labels
        **super_kwargs,         # Additional arguments for the Dataset base class.
    ):
        self._path = path
        print(self._path)
        self._resolution = resolution
        self._zipfile = None
        self._data_type = data_type
        self.static = static
        self.use_flame_mot = use_flame_mot
        self.rescale_camera = rescale_camera
        
        # initialize lmdb
        if os.path.isdir(self._path):
            self.db = None
            self.txn = None
            self.num = None
            self.datum = None
        else:
            raise IOError('Path must point to a directory or zip')

        img_size = int(self._path.split("/")[-2].split("_")[-2])
        num = int(self._path.split("/")[-2].split("_")[-1])
        img_shape = [3, img_size, img_size]
        raw_shape = [num] + img_shape
        if max_num is not None:
            raw_shape = [max_num] + img_shape
        if resolution is None:
            self._resolution = raw_shape[2]
        name = os.path.splitext(os.path.basename(self._path))[0]
        super().__init__(name=name, raw_shape=raw_shape, **super_kwargs)

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_vfhq.Datum_vfhq()

    def get_details(self, idx):
        return None

    def get_label_std(self):
        return 0

    @property
    def resolution(self):
        return self._resolution

    @property
    def label_shape(self):
        if self._label_shape is None:
            raw_labels, _ = self._load_raw_labels(0,[0,0])
            self._label_shape = raw_labels.shape
        return list(self._label_shape)

    @property
    def label_dim(self):
        assert len(self.label_shape) == 1
        return self.label_shape[0]

    @property
    def has_labels(self):
        return any(x != 0 for x in self.label_shape)

    @property
    def has_onehot_labels(self):
        return None

    def __getstate__(self):
        return dict(super().__getstate__(), _zipfile=None)

    def _load_raw_image(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        images = np.fromstring(self.datum.images, dtype=np.uint8).reshape(self.datum.num,-1)
        frame_idx = np.random.randint(self.datum.num,size=2)

        source_image = images[frame_idx[0]]
        driving_image = images[frame_idx[1]]
        source_image = source_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        driving_image = driving_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        
        if source_image.ndim == 2:
            source_image = source_image[:, :, np.newaxis] # HW => HWC
        if driving_image.ndim == 2:
            driving_image = driving_image[:, :, np.newaxis] # HW => HWC

        source_image = source_image.transpose(2, 0, 1) # HWC => CHW
        driving_image = driving_image.transpose(2, 0, 1) # HWC => CHW
        
        return source_image, driving_image, frame_idx

    def _load_raw_seg(self, raw_idx, frame_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs, dtype=np.uint8).reshape(self.datum.num,-1)
        
        source_seg = seg[frame_idx[0]]
        driving_seg = seg[frame_idx[1]]
        source_seg = source_seg.reshape(self.datum.height,self.datum.width,-1)
        driving_seg = driving_seg.reshape(self.datum.height,self.datum.width,-1)
        
        if source_seg.ndim == 2:
            source_seg = source_seg[:, :, np.newaxis] # HW => HWC
        if driving_seg.ndim == 2:
            driving_seg = driving_seg[:, :, np.newaxis] # HW => HWC
        source_seg = source_seg.transpose(2, 0, 1) # HWC => CHW
        driving_seg = driving_seg.transpose(2, 0, 1) # HWC => CHW

        if source_seg.shape[0] == 1:
            source_seg = np.tile(source_seg, (3, 1, 1))

        if driving_seg.shape[0] == 1:
            driving_seg = np.tile(driving_seg, (3, 1, 1))
        return source_seg, driving_seg

    def _load_raw_labels(self, raw_idx, frame_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        labels = np.fromstring(self.datum.labels, dtype=np.float32).reshape(self.datum.num,-1)

        # Use average yaw and roll angles across all frames in a video clip as the camera pose to rectify extracted head pose
        head_poses = labels[:,425:428]
        compose_rot = R.from_rotvec(head_poses)
        compose_euler = compose_rot.as_euler('xyz')
        compose_euler_mean = np.mean(compose_euler,axis=0,keepdims=True)
        cam_euler = np.stack([np.zeros_like(compose_euler_mean[...,0]), compose_euler_mean[...,1], compose_euler_mean[...,2]], axis=-1)
        cam_rot = R.from_euler('xyz', cam_euler).as_matrix()
        cam_rot_inverse = np.transpose(cam_rot,[0,2,1])
        head_rot_rectified = np.matmul(cam_rot_inverse,compose_rot.as_matrix())
        head_poses = R.from_matrix(head_rot_rectified).as_rotvec()
        labels[:,425:428] = head_poses

        source_labels = labels[frame_idx[0]]
        driving_labels = labels[frame_idx[1]]
        
        if self.rescale_camera:
            if frame_idx[0] == frame_idx[1]:
                intrinsics = source_labels[16:25].reshape(3,3)
                
                # normalize intrinsics
                if self._resolution != intrinsics[0,2]*2:
                    intrinsics[:2,:] *= (0.5*self._resolution/intrinsics[0,2])

                intrinsics[0, 0] /= self._resolution
                intrinsics[1, 1] /= self._resolution
                intrinsics[0, 2] /= self._resolution
                intrinsics[1, 2] /= self._resolution       
                
                # rescale extrinsics
                extrinsics = source_labels[:16].reshape(4,4) # Our face scale is around 0.1~0.2. Multiply by 3 to match the scale of EG3D
                extrinsics[:3,3] *= 3
            else:    
                for labels in [source_labels, driving_labels]:
                    intrinsics = labels[16:25].reshape(3,3)
                    if self._resolution != intrinsics[0,2]*2:
                        intrinsics[:2,:] *= (0.5*self._resolution/intrinsics[0,2])

                    intrinsics[0, 0] /= self._resolution
                    intrinsics[1, 1] /= self._resolution
                    intrinsics[0, 2] /= self._resolution
                    intrinsics[1, 2] /= self._resolution       

                    extrinsics = labels[:16].reshape(4,4)  
                    extrinsics[:3,3] *= 3

        return source_labels, driving_labels

    def _load_raw_motions(self, raw_idx, frame_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        motions = np.fromstring(self.datum.mots, dtype=np.float32).reshape(self.datum.num,-1)
        motions = motions[:,:548]
        source_motions = motions[frame_idx[0]]
        driving_motions = motions[frame_idx[1]]

        return source_motions, driving_motions

    def get_label(self, raw_labels):
        label = raw_labels[:25]
        return label.copy()

    def get_shape_param(self, raw_labels):
        label = raw_labels[25:325]
        return label.copy()

    def get_exp_param(self, raw_labels):
        label = raw_labels[325:425]
        return label.copy()

    def get_exp_param_w_jaw_pose(self, raw_labels):
        label = np.concatenate([raw_labels[325:425],raw_labels[428:431]],axis=0)
        return label.copy()

    def get_pose_param(self, raw_labels):
        label = raw_labels[425:431]
        return label.copy()

    def get_eye_pose_param(self, raw_labels):
        label = raw_labels[431:437]
        return label.copy()

    def get_label_all(self,raw_labels):
        c = self.get_label(raw_labels)
        shape_param = self.get_shape_param(raw_labels)
        exp_param = self.get_exp_param(raw_labels)
        pose_param = self.get_pose_param(raw_labels)
        eye_pose_param = self.get_eye_pose_param(raw_labels)
        return c, shape_param, exp_param, pose_param, eye_pose_param

    def __getitem__(self, idx):
        source_image, driving_image, frame_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg = self._load_raw_seg(self._raw_idx[idx], frame_idx)
        source_labels, driving_labels = self._load_raw_labels(self._raw_idx[idx], frame_idx)
        source_motions, driving_motions = self._load_raw_motions(self._raw_idx[idx], frame_idx)
        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8
        
        # use flame parameters as motion embedding if use_flame_mot is true
        if self.use_flame_mot:
            source_motions = np.concatenate([self.get_exp_param_w_jaw_pose(source_labels),self.get_eye_pose_param(source_labels)],axis=0)
            driving_motions = np.concatenate([self.get_exp_param_w_jaw_pose(driving_labels),self.get_eye_pose_param(driving_labels)],axis=0)

        return source_image.copy(), driving_image.copy(), driving_image.copy(), driving_seg.copy(), np.zeros([]), np.zeros([]), np.zeros([]), np.zeros([]), np.zeros([]), self.get_label(driving_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), source_motions.copy(), driving_motions.copy()

#----------------------------------------------------------------------------
# Dataloader for real image set
class OneShotReconSegLmdbFolderDataset_Static(OneShotReconSegLmdbFolderDataset):

    @property
    def label_shape(self):
        if self._label_shape is None:
            raw_labels = self._load_raw_labels(0)
            self._label_shape = raw_labels.shape
        return list(self._label_shape)

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_ffhq.Datum_ffhq()

    def _load_raw_image(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        image = np.fromstring(self.datum.images, dtype=np.uint8)
        image = cv2.imdecode(image, cv2.IMREAD_COLOR)
        image = image[:, :, [2, 1, 0]] # bgr -> rgb
        
        image = image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        if image.ndim == 2:
            image = image[:, :, np.newaxis] # HW => HWC
        image = image.transpose(2, 0, 1) # HWC => CHW
        return image

    def _load_raw_seg(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs, dtype=np.uint8)
        seg = cv2.imdecode(seg, cv2.IMREAD_COLOR)        

        if seg.ndim == 2:
            seg = seg[:, :, np.newaxis] # HW => HWC
        seg = seg.transpose(2, 0, 1) # HWC => CHW

        if seg.shape[0] == 1:
            seg = np.tile(seg, (3, 1, 1))

        return seg

    def _load_raw_labels(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        labels = np.fromstring(self.datum.labels, dtype=np.float32).reshape(-1)
        labels[425:428] = 0 # set head pose to zero and use camera pose instead
        
        if self.rescale_camera:
            intrinsics = labels[16:25].reshape(3,3)
            
            # normalize intrinsics
            if self._resolution != intrinsics[0,2]*2:
                intrinsics[:2,:] *= (0.5*self._resolution/intrinsics[0,2])

            intrinsics[0, 0] /= self._resolution
            intrinsics[1, 1] /= self._resolution
            intrinsics[0, 2] /= self._resolution
            intrinsics[1, 2] /= self._resolution       
            
            # rescale extrinsics
            extrinsics = labels[:16].reshape(4,4) # Our face scale is around 0.1~0.2. Multiply by 3 to match the scale of EG3D
            extrinsics[:3,3] *= 3

        return labels

    def _load_raw_motions(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        motions = np.fromstring(self.datum.mots, dtype=np.float32).reshape(-1)
        motions = motions[:548]

        return motions

    def __getitem__(self, idx):
        image = self._load_raw_image(self._raw_idx[idx])
        seg = self._load_raw_seg(self._raw_idx[idx])
        labels = self._load_raw_labels(self._raw_idx[idx])
        motions = self._load_raw_motions(self._raw_idx[idx])
        assert isinstance(image, np.ndarray)
        assert isinstance(seg, np.ndarray)
        assert list(image.shape) == self.image_shape
        assert seg.shape[1] == self.image_shape[1] and seg.shape[2] == self.image_shape[2]
        assert image.dtype == np.uint8

        return image.copy(), image.copy(), image.copy(), seg.copy(), np.zeros([]), np.zeros([]), np.zeros([]), np.zeros([]), np.zeros([]), self.get_label(labels), self.get_shape_param(labels), self.get_exp_param(labels), self.get_pose_param(labels), self.get_eye_pose_param(labels), motions.copy(), motions.copy()

#-----------------------------------------------------------------------------
# Dataloader for GenHead generated 4D data (base version)
class GenHeadReconSegLmdbFolderDataset(OneShotReconSegLmdbFolderDataset):

    @property
    def label_shape(self):
        if self._label_shape is None:
            raw_labels, _, _ = self._load_raw_labels(0,[0,0],[0,0,0])
            self._label_shape = raw_labels.shape
        return list(self._label_shape)

    def _load_raw_image(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        images = np.fromstring(self.datum.images, dtype=np.uint8).reshape(self.datum.num//5,5,-1)
        frame_idx = np.random.randint(self.datum.num//5,size=2)
        
        # use same motion for source and driving if static is true
        if self.static:
            frame_idx[1] = frame_idx[0]

        view_idx = np.random.randint(5,size=3)
        source_image = images[frame_idx[0],view_idx[0]]
        driving_image = images[frame_idx[1],view_idx[1]]
        recon_image = images[frame_idx[1],view_idx[2]]

        source_image = source_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        driving_image = driving_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        recon_image = recon_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        if source_image.ndim == 2:
            source_image = source_image[:, :, np.newaxis] # HW => HWC
        if driving_image.ndim == 2:
            driving_image = driving_image[:, :, np.newaxis] # HW => HWC
        if recon_image.ndim == 2:
            recon_image = recon_image[:, :, np.newaxis] # HW => HWC

        source_image = source_image.transpose(2, 0, 1) # HWC => CHW
        driving_image = driving_image.transpose(2, 0, 1) # HWC => CHW
        recon_image = recon_image.transpose(2, 0, 1) # HWC => CHW
        
        return source_image, driving_image, recon_image, frame_idx, view_idx

    def _load_raw_seg(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs, dtype=np.uint8).reshape(self.datum.num//5,5,-1)
    
        source_seg = seg[frame_idx[0],view_idx[0]]
        driving_seg = seg[frame_idx[1],view_idx[1]]
        recon_seg = seg[frame_idx[1],view_idx[2]]
        source_seg = source_seg.reshape(self.datum.height,self.datum.width,-1)
        driving_seg = driving_seg.reshape(self.datum.height,self.datum.width,-1)
        recon_seg = recon_seg.reshape(self.datum.height,self.datum.width,-1)
        if source_seg.ndim == 2:
            source_seg = source_seg[:, :, np.newaxis] # HW => HWC
        if driving_seg.ndim == 2:
            driving_seg = driving_seg[:, :, np.newaxis] # HW => HWC
        if recon_seg.ndim == 2:
            recon_seg = recon_seg[:, :, np.newaxis] # HW => HWC
        source_seg = source_seg.transpose(2, 0, 1) # HWC => CHW
        driving_seg = driving_seg.transpose(2, 0, 1) # HWC => CHW
        recon_seg = recon_seg.transpose(2, 0, 1) # HWC => CHW

        if source_seg.shape[0] == 1:
            source_seg = np.tile(source_seg, (3, 1, 1))

        if driving_seg.shape[0] == 1:
            driving_seg = np.tile(driving_seg, (3, 1, 1))

        if recon_seg.shape[0] == 1:
            recon_seg = np.tile(recon_seg, (3, 1, 1))

        return source_seg, driving_seg, recon_seg

    def _load_raw_labels(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        labels = np.fromstring(self.datum.labels, dtype=np.float32).reshape(self.datum.num//5,5,-1)

        source_labels = labels[frame_idx[0], view_idx[0]]
        driving_labels = labels[frame_idx[1], view_idx[1]]
        recon_labels = labels[frame_idx[1], view_idx[2]]

        return source_labels, driving_labels, recon_labels

    def _load_raw_motions(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        motions = np.fromstring(self.datum.mots, dtype=np.float32).reshape(self.datum.num//5,5,-1)
        motions = motions[...,:548]
        source_motions = motions[frame_idx[0], view_idx[0]]
        driving_motions = motions[frame_idx[1], view_idx[1]]
        recon_motios = motions[frame_idx[1], view_idx[2]]

        return source_motions, driving_motions, recon_motios

    def __getitem__(self, idx):
        source_image, driving_image, recon_image, frame_idx, view_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg, recon_seg = self._load_raw_seg(self._raw_idx[idx], frame_idx, view_idx)
        source_labels, driving_labels, recon_labels = self._load_raw_labels(self._raw_idx[idx], frame_idx, view_idx)
        source_motions, driving_motions, recon_motions = self._load_raw_motions(self._raw_idx[idx], frame_idx, view_idx)

        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(recon_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert isinstance(recon_seg, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8

        return source_image.copy(), driving_image.copy(), recon_image.copy(), recon_seg.copy(), self.get_label(recon_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), driving_motions.copy()

#-----------------------------------------------------------------------------
# Dataloader for GenHead generated 4D data (include depth, feature maps, and triplane features)
class GenHeadReconSegLmdbFolderDatasetV2(GenHeadReconSegLmdbFolderDataset):

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_genhead.Datum()

    def _load_raw_depth(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        depth = np.fromstring(self.datum.depths, dtype=np.float16).reshape(self.datum.num//5,5,-1)
    
        source_depth = depth[frame_idx[0],view_idx[0]]
        driving_depth = depth[frame_idx[1],view_idx[1]]
        recon_depth = depth[frame_idx[1],view_idx[2]]
        source_depth = source_depth.reshape(64,64,-1)
        driving_depth = driving_depth.reshape(64,64,-1)
        recon_depth = recon_depth.reshape(64,64,-1)

        source_depth = source_depth.transpose(2, 0, 1) # HWC => CHW
        driving_depth = driving_depth.transpose(2, 0, 1) # HWC => CHW
        recon_depth = recon_depth.transpose(2, 0, 1) # HWC => CHW


        return source_depth, driving_depth, recon_depth

    def _load_raw_feature(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        feature = np.fromstring(self.datum.features, dtype=np.float16).reshape(self.datum.num//5,5,-1)
    
        source_feature = feature[frame_idx[0],view_idx[0]]
        driving_feature = feature[frame_idx[1],view_idx[1]]
        recon_feature = feature[frame_idx[1],view_idx[2]]
        
        source_feature = source_feature.reshape(64,64,-1)
        driving_feature = driving_feature.reshape(64,64,-1)
        recon_feature = recon_feature.reshape(64,64,-1)

        source_feature = source_feature.transpose(2, 0, 1) # HWC => CHW
        driving_feature = driving_feature.transpose(2, 0, 1) # HWC => CHW
        recon_feature = recon_feature.transpose(2, 0, 1) # HWC => CHW


        return source_feature, driving_feature, recon_feature

    def _load_raw_triplane(self, raw_idx, frame_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        triplane = np.fromstring(self.datum.triplanes, dtype=np.float16).reshape(self.datum.num//5,-1)
    
        source_triplane = triplane[frame_idx[0]]
        driving_triplane = triplane[frame_idx[1]]
        recon_triplane = triplane[frame_idx[1]]

        source_triplane = source_triplane.reshape(-1,35) # [N,3+32]
        driving_triplane = driving_triplane.reshape(-1,35)
        recon_triplane = recon_triplane.reshape(-1,35)

        return source_triplane, driving_triplane, recon_triplane

    def __getitem__(self, idx):
        source_image, driving_image, recon_image, frame_idx, view_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg, recon_seg = self._load_raw_seg(self._raw_idx[idx], frame_idx, view_idx)
        source_depth, driving_depth, recon_depth = self._load_raw_depth(self._raw_idx[idx], frame_idx, view_idx)
        source_feature, driving_feature, recon_feature = self._load_raw_feature(self._raw_idx[idx], frame_idx, view_idx)
        source_labels, driving_labels, recon_labels = self._load_raw_labels(self._raw_idx[idx], frame_idx, view_idx)
        source_motions, driving_motions, recon_motions = self._load_raw_motions(self._raw_idx[idx], frame_idx, view_idx)
        source_triplane, driving_triplane, recon_triplane = self._load_raw_triplane(self._raw_idx[idx], frame_idx)

        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(recon_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert isinstance(recon_seg, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8

        return source_image.copy(), driving_image.copy(), recon_image.copy(), recon_seg.copy(), recon_depth.copy(), recon_feature.copy(), recon_triplane.copy(), self.get_label(recon_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), source_motions.copy(), driving_motions.copy()

#----------------------------------------------------------------------------
# Dataloader for GenHead generated static multiview data (include depth, feature maps, and triplane features)
class GenHeadReconSegLmdbFolderDatasetV2_Static(OneShotReconSegLmdbFolderDataset):

    @property
    def label_shape(self):
        if self._label_shape is None:
            raw_labels, _, _ = self._load_raw_labels(0,[0,0,0])
            self._label_shape = raw_labels.shape
        return list(self._label_shape)

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_genhead_static.Datum_static()

    def _load_raw_image(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        images = np.fromstring(self.datum.images, dtype=np.uint8).reshape(self.datum.num,-1)

        view_idx = np.random.randint(self.datum.num,size=3)

        source_image = images[view_idx[0]]
        driving_image = images[view_idx[1]]
        recon_image = images[view_idx[2]]

        source_image = source_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        driving_image = driving_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        recon_image = recon_image.reshape(self.datum.height,self.datum.width,self.datum.channels)
        if source_image.ndim == 2:
            source_image = source_image[:, :, np.newaxis] # HW => HWC
        if driving_image.ndim == 2:
            driving_image = driving_image[:, :, np.newaxis] # HW => HWC
        if recon_image.ndim == 2:
            recon_image = recon_image[:, :, np.newaxis] # HW => HWC

        source_image = source_image.transpose(2, 0, 1) # HWC => CHW
        driving_image = driving_image.transpose(2, 0, 1) # HWC => CHW
        recon_image = recon_image.transpose(2, 0, 1) # HWC => CHW
        return source_image, driving_image, recon_image, view_idx

    def _load_raw_seg(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs, dtype=np.uint8).reshape(self.datum.num,-1)
    
        source_seg = seg[view_idx[0]]
        driving_seg = seg[view_idx[1]]
        recon_seg = seg[view_idx[2]]
        
        source_seg = source_seg.reshape(self.datum.height,self.datum.width,-1)
        driving_seg = driving_seg.reshape(self.datum.height,self.datum.width,-1)
        recon_seg = recon_seg.reshape(self.datum.height,self.datum.width,-1)
        if source_seg.ndim == 2:
            source_seg = source_seg[:, :, np.newaxis] # HW => HWC
        if driving_seg.ndim == 2:
            driving_seg = driving_seg[:, :, np.newaxis] # HW => HWC
        if recon_seg.ndim == 2:
            recon_seg = recon_seg[:, :, np.newaxis] # HW => HWC
        source_seg = source_seg.transpose(2, 0, 1) # HWC => CHW
        driving_seg = driving_seg.transpose(2, 0, 1) # HWC => CHW
        recon_seg = recon_seg.transpose(2, 0, 1) # HWC => CHW

        if source_seg.shape[0] == 1:
            source_seg = np.tile(source_seg, (3, 1, 1))

        if driving_seg.shape[0] == 1:
            driving_seg = np.tile(driving_seg, (3, 1, 1))

        if recon_seg.shape[0] == 1:
            recon_seg = np.tile(recon_seg, (3, 1, 1))

        return source_seg, driving_seg, recon_seg

    def _load_raw_labels(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        labels = np.fromstring(self.datum.labels, dtype=np.float32).reshape(self.datum.num,-1)

        source_labels = labels[view_idx[0]]
        driving_labels = labels[view_idx[1]]
        recon_labels = labels[view_idx[2]]

        return source_labels, driving_labels, recon_labels

    def _load_raw_motions(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        motions = np.fromstring(self.datum.mots, dtype=np.float32).reshape(self.datum.num,-1)
        motions = motions[...,:548]
        source_motions = motions[view_idx[0]]
        driving_motions = motions[view_idx[1]]
        recon_motios = motions[view_idx[2]]

        return source_motions, driving_motions, recon_motios

    def _load_raw_depth(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        depth = np.fromstring(self.datum.depths, dtype=np.float16).reshape(self.datum.num,-1)
    
        source_depth = depth[view_idx[0]]
        driving_depth = depth[view_idx[1]]
        recon_depth = depth[view_idx[2]]

        source_depth = source_depth.reshape(64,64,-1)
        driving_depth = driving_depth.reshape(64,64,-1)
        recon_depth = recon_depth.reshape(64,64,-1)

        source_depth = source_depth.transpose(2, 0, 1) # HWC => CHW
        driving_depth = driving_depth.transpose(2, 0, 1) # HWC => CHW
        recon_depth = recon_depth.transpose(2, 0, 1) # HWC => CHW

        return source_depth, driving_depth, recon_depth

    def _load_raw_feature(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        feature = np.fromstring(self.datum.features, dtype=np.float16).reshape(self.datum.num,-1)
    
        source_feature = feature[view_idx[0]]
        driving_feature = feature[view_idx[1]]
        recon_feature = feature[view_idx[2]]

        source_feature = source_feature.reshape(64,64,-1)
        driving_feature = driving_feature.reshape(64,64,-1)
        recon_feature = recon_feature.reshape(64,64,-1)

        source_feature = source_feature.transpose(2, 0, 1) # HWC => CHW
        driving_feature = driving_feature.transpose(2, 0, 1) # HWC => CHW
        recon_feature = recon_feature.transpose(2, 0, 1) # HWC => CHW

        return source_feature, driving_feature, recon_feature

    def _load_raw_triplane(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        triplane = np.fromstring(self.datum.triplanes, dtype=np.float16).reshape(-1)
    
        source_triplane = triplane
        driving_triplane = triplane
        recon_triplane = triplane

        source_triplane = source_triplane.reshape(-1,35) # [N,3+32]
        driving_triplane = driving_triplane.reshape(-1,35)
        recon_triplane = recon_triplane.reshape(-1,35)

        return source_triplane, driving_triplane, recon_triplane

    def __getitem__(self, idx):
        source_image, driving_image, recon_image, view_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg, recon_seg = self._load_raw_seg(self._raw_idx[idx], view_idx)
        source_depth, driving_depth, recon_depth = self._load_raw_depth(self._raw_idx[idx], view_idx)
        source_feature, driving_feature, recon_feature = self._load_raw_feature(self._raw_idx[idx], view_idx)
        source_labels, driving_labels, recon_labels = self._load_raw_labels(self._raw_idx[idx], view_idx)
        source_motions, driving_motions, recon_motions = self._load_raw_motions(self._raw_idx[idx], view_idx)
        source_triplane, driving_triplane, recon_triplane = self._load_raw_triplane(self._raw_idx[idx])

        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(recon_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert isinstance(recon_seg, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8

        return source_image.copy(), driving_image.copy(), recon_image.copy(), recon_seg.copy(), recon_depth.copy(), recon_feature.copy(), recon_triplane.copy(), self.get_label(recon_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), source_motions.copy(), driving_motions.copy()

#--------------------------------------------------------------------------------------------------------
# Dataloader for GenHead generated 4D data (include depth, feature maps, triplane features, background, and rendered segmentation)
class GenHeadReconSegLmdbFolderDatasetV2New(GenHeadReconSegLmdbFolderDatasetV2):

    def __init__(self,
        path,                   # Path to datalist.
        resolution      = None, # Ensure specific resolution, None = highest available.
        data_type      = "vfhq",# Deprecated
        static         = False, # if True, only select multiview static frames instead of frames with different motions
        use_flame_mot  = False, # whether or not to use FLAME parameters as motion embedding
        **super_kwargs,         # Additional arguments for the Dataset base class.
        ):
        super().__init__(path=path, resolution=resolution, data_type=data_type, static=static, **super_kwargs)
        self.use_flame_mot = use_flame_mot

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_genhead_new.Datum_new()

    def _load_raw_bg_feature(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        feature = np.fromstring(self.datum.bgs, dtype=np.float16)
    
        recon_feature = feature
        recon_feature = recon_feature.reshape(64,64,-1)
        recon_feature = recon_feature.transpose(2, 0, 1) # HWC => CHW

        return recon_feature

    def _load_raw_seg_render(self, raw_idx, frame_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs_render, dtype=np.uint8).reshape(self.datum.num//5,5,-1)
    
        source_seg = seg[frame_idx[0],view_idx[0]]
        driving_seg = seg[frame_idx[1],view_idx[1]]
        recon_seg = seg[frame_idx[1],view_idx[2]]
        source_seg = source_seg.reshape(64,64,-1)
        driving_seg = driving_seg.reshape(64,64,-1)
        recon_seg = recon_seg.reshape(64,64,-1)
        if source_seg.ndim == 2:
            source_seg = source_seg[:, :, np.newaxis] # HW => HWC
        if driving_seg.ndim == 2:
            driving_seg = driving_seg[:, :, np.newaxis] # HW => HWC
        if recon_seg.ndim == 2:
            recon_seg = recon_seg[:, :, np.newaxis] # HW => HWC
        source_seg = source_seg.transpose(2, 0, 1) # HWC => CHW
        driving_seg = driving_seg.transpose(2, 0, 1) # HWC => CHW
        recon_seg = recon_seg.transpose(2, 0, 1) # HWC => CHW

        if source_seg.shape[0] == 1:
            source_seg = np.tile(source_seg, (3, 1, 1))

        if driving_seg.shape[0] == 1:
            driving_seg = np.tile(driving_seg, (3, 1, 1))

        if recon_seg.shape[0] == 1:
            recon_seg = np.tile(recon_seg, (3, 1, 1))

        return source_seg, driving_seg, recon_seg

    def __getitem__(self, idx):
        source_image, driving_image, recon_image, frame_idx, view_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg, recon_seg = self._load_raw_seg(self._raw_idx[idx], frame_idx, view_idx)
        source_seg_render, driving_seg_render, recon_seg_render = self._load_raw_seg_render(self._raw_idx[idx], frame_idx, view_idx)
        source_depth, driving_depth, recon_depth = self._load_raw_depth(self._raw_idx[idx], frame_idx, view_idx)
        source_feature, driving_feature, recon_feature = self._load_raw_feature(self._raw_idx[idx], frame_idx, view_idx)
        recon_feature_bg = self._load_raw_bg_feature(self._raw_idx[idx])
        source_labels, driving_labels, recon_labels = self._load_raw_labels(self._raw_idx[idx], frame_idx, view_idx)
        source_motions, driving_motions, recon_motions = self._load_raw_motions(self._raw_idx[idx], frame_idx, view_idx)
        source_triplane, driving_triplane, recon_triplane = self._load_raw_triplane(self._raw_idx[idx], frame_idx)

        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(recon_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert isinstance(recon_seg, np.ndarray)
        assert isinstance(source_seg_render, np.ndarray)
        assert isinstance(driving_seg_render, np.ndarray)
        assert isinstance(source_seg_render, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8
        
        # use flame parameters as motion embedding if use_flame_mot is true
        if self.use_flame_mot:
            source_motions = np.concatenate([self.get_exp_param_w_jaw_pose(source_labels),self.get_eye_pose_param(source_labels)],axis=0)
            driving_motions = np.concatenate([self.get_exp_param_w_jaw_pose(driving_labels),self.get_eye_pose_param(driving_labels)],axis=0)

        return source_image.copy(), driving_image.copy(), recon_image.copy(), recon_seg.copy(), recon_seg_render.copy(), recon_depth.copy(), recon_feature.copy(), recon_feature_bg.copy(), recon_triplane.copy(), self.get_label(recon_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), source_motions.copy(), driving_motions.copy()

#--------------------------------------------------------------------------------------------------------
# Dataloader for GenHead generated static multiview data (include depth, feature maps, triplane features, background, and rendered segmentation)
class GenHeadReconSegLmdbFolderDatasetV2_StaticNew(GenHeadReconSegLmdbFolderDatasetV2_Static):

    def __init__(self,
        path,                   # Path to datalist.
        resolution      = None, # Ensure specific resolution, None = highest available.
        data_type      = "vfhq",# Deprecated 
        static         = False, # if True, only select multiview static frames instead of frames with different motions
        use_flame_mot  = False, # whether or not to use FLAME parameters as motion embedding
        **super_kwargs,         # Additional arguments for the Dataset base class.
        ):
        super().__init__(path=path, resolution=resolution, data_type=data_type, static=static, **super_kwargs)
        self.use_flame_mot = use_flame_mot

    def open_lmdb(self):
        self.db = lmdb.open(self._path, map_size=1024 ** 4, readonly=True, lock=False)
        self.txn = self.db.begin()
        self.num = int(self.txn.get('num_samples'.encode()))
        self.datum = datum_pb2_genhead_static_new.Datum_static_new()

    def _load_raw_bg_feature(self, raw_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        feature = np.fromstring(self.datum.bgs, dtype=np.float16)
    
        recon_feature = feature
        recon_feature = recon_feature.reshape(64,64,-1)
        recon_feature = recon_feature.transpose(2, 0, 1) # HWC => CHW

        return recon_feature

    def _load_raw_seg_render(self, raw_idx, view_idx):
        if self.txn is None:
            self.open_lmdb()
        value = self.txn.get('{:0>8d}'.format(raw_idx).encode())
        self.datum.ParseFromString(value)
        seg = np.fromstring(self.datum.segs_render, dtype=np.uint8).reshape(self.datum.num,-1)
    
        source_seg = seg[view_idx[0]]
        driving_seg = seg[view_idx[1]]
        recon_seg = seg[view_idx[2]]
        source_seg = source_seg.reshape(64,64,-1)
        driving_seg = driving_seg.reshape(64,64,-1)
        recon_seg = recon_seg.reshape(64,64,-1)
        if source_seg.ndim == 2:
            source_seg = source_seg[:, :, np.newaxis] # HW => HWC
        if driving_seg.ndim == 2:
            driving_seg = driving_seg[:, :, np.newaxis] # HW => HWC
        if recon_seg.ndim == 2:
            recon_seg = recon_seg[:, :, np.newaxis] # HW => HWC
        source_seg = source_seg.transpose(2, 0, 1) # HWC => CHW
        driving_seg = driving_seg.transpose(2, 0, 1) # HWC => CHW
        recon_seg = recon_seg.transpose(2, 0, 1) # HWC => CHW

        if source_seg.shape[0] == 1:
            source_seg = np.tile(source_seg, (3, 1, 1))

        if driving_seg.shape[0] == 1:
            driving_seg = np.tile(driving_seg, (3, 1, 1))

        if recon_seg.shape[0] == 1:
            recon_seg = np.tile(recon_seg, (3, 1, 1))

        return source_seg, driving_seg, recon_seg

    def __getitem__(self, idx):
        source_image, driving_image, recon_image, view_idx = self._load_raw_image(self._raw_idx[idx])
        source_seg, driving_seg, recon_seg = self._load_raw_seg(self._raw_idx[idx], view_idx)
        source_seg_render, driving_seg_render, recon_seg_render = self._load_raw_seg_render(self._raw_idx[idx], view_idx)
        source_depth, driving_depth, recon_depth = self._load_raw_depth(self._raw_idx[idx], view_idx)
        source_feature, driving_feature, recon_feature = self._load_raw_feature(self._raw_idx[idx], view_idx)
        recon_feature_bg = self._load_raw_bg_feature(self._raw_idx[idx])
        source_labels, driving_labels, recon_labels = self._load_raw_labels(self._raw_idx[idx], view_idx)
        source_motions, driving_motions, recon_motions = self._load_raw_motions(self._raw_idx[idx], view_idx)
        source_triplane, driving_triplane, recon_triplane = self._load_raw_triplane(self._raw_idx[idx])

        assert isinstance(source_image, np.ndarray)
        assert isinstance(driving_image, np.ndarray)
        assert isinstance(recon_image, np.ndarray)
        assert isinstance(source_seg, np.ndarray)
        assert isinstance(driving_seg, np.ndarray)
        assert isinstance(recon_seg, np.ndarray)
        assert isinstance(source_seg_render, np.ndarray)
        assert isinstance(driving_seg_render, np.ndarray)
        assert isinstance(source_seg_render, np.ndarray)
        assert list(source_image.shape) == self.image_shape
        assert source_seg.shape[1] == self.image_shape[1] and source_seg.shape[2] == self.image_shape[2]
        assert source_image.dtype == np.uint8
        
        # use flame parameters as motion embedding if use_flame_mot is true
        if self.use_flame_mot:
            source_motions = np.concatenate([self.get_exp_param_w_jaw_pose(source_labels),self.get_eye_pose_param(source_labels)],axis=0)
            driving_motions = np.concatenate([self.get_exp_param_w_jaw_pose(driving_labels),self.get_eye_pose_param(driving_labels)],axis=0)

        return source_image.copy(), driving_image.copy(), recon_image.copy(), recon_seg.copy(), recon_seg_render.copy(), recon_depth.copy(), recon_feature.copy(), recon_feature_bg.copy(), recon_triplane.copy(), self.get_label(recon_labels), self.get_shape_param(source_labels), self.get_exp_param(driving_labels), self.get_pose_param(driving_labels), self.get_eye_pose_param(driving_labels), source_motions.copy(), driving_motions.copy()