datasets/shapenet.py

import os
import torchvision
import pickle
from typing import Any
import lmdb
import cv2
import imageio
import numpy as np
from PIL import Image
import Imath
import OpenEXR
from pdb import set_trace as st
from pathlib import Path

from functools import partial
import io
import gzip
import random
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torch.utils.data.distributed import DistributedSampler
from pathlib import Path

from guided_diffusion import logger

def load_dataset(
    file_path="",
    reso=64,
    reso_encoder=224,
    batch_size=1,
    #   shuffle=True,
    num_workers=6,
    load_depth=False,
    preprocess=None,
    imgnet_normalize=True,
    dataset_size=-1,
    trainer_name='input_rec',
    use_lmdb=False,
    infi_sampler=True
):
    # st()
    # dataset_cls = {
    #     'input_rec': MultiViewDataset,
    #     'nv': NovelViewDataset,
    # }[trainer_name]
    # st()
    if use_lmdb:
        logger.log('using LMDB dataset')
        # dataset_cls = LMDBDataset_MV #  2.5-3iter/s, but unstable, drops to 1 later.
        if 'nv' in trainer_name:
            dataset_cls = LMDBDataset_NV_Compressed  #  2.5-3iter/s, but unstable, drops to 1 later.
        else:
            dataset_cls = LMDBDataset_MV_Compressed  #  2.5-3iter/s, but unstable, drops to 1 later.
        # dataset = dataset_cls(file_path)
    else:
        if 'nv' in trainer_name:
            dataset_cls = NovelViewDataset  # 1.5-2iter/s
        else:
            dataset_cls = MultiViewDataset

    dataset = dataset_cls(file_path,
                          reso,
                          reso_encoder,
                          test=False,
                          preprocess=preprocess,
                          load_depth=load_depth,
                          imgnet_normalize=imgnet_normalize,
                          dataset_size=dataset_size)

    logger.log('dataset_cls: {}, dataset size: {}'.format(
        trainer_name, len(dataset)))
    
    loader = DataLoader(dataset,
                        batch_size=batch_size,
                        num_workers=num_workers,
                        drop_last=False,
                        pin_memory=True,
                        persistent_workers=num_workers > 0, 
                        shuffle=False)
    return loader


def load_data(
    file_path="",
    reso=64,
    reso_encoder=224,
    batch_size=1,
    #   shuffle=True,
    num_workers=6,
    load_depth=False,
    preprocess=None,
    imgnet_normalize=True,
    dataset_size=-1,
    trainer_name='input_rec',
    use_lmdb=False,
    infi_sampler=True
):
    # st()
    # dataset_cls = {
    #     'input_rec': MultiViewDataset,
    #     'nv': NovelViewDataset,
    # }[trainer_name]
    # st()
    if use_lmdb:
        logger.log('using LMDB dataset')
        # dataset_cls = LMDBDataset_MV #  2.5-3iter/s, but unstable, drops to 1 later.
        if 'nv' in trainer_name:
            dataset_cls = LMDBDataset_NV_Compressed  #  2.5-3iter/s, but unstable, drops to 1 later.
        else:
            dataset_cls = LMDBDataset_MV_Compressed  #  2.5-3iter/s, but unstable, drops to 1 later.
        # dataset = dataset_cls(file_path)
    else:
        if 'nv' in trainer_name:
            dataset_cls = NovelViewDataset  # 1.5-2iter/s
        else:
            dataset_cls = MultiViewDataset

    dataset = dataset_cls(file_path,
                          reso,
                          reso_encoder,
                          test=False,
                          preprocess=preprocess,
                          load_depth=load_depth,
                          imgnet_normalize=imgnet_normalize,
                          dataset_size=dataset_size)

    logger.log('dataset_cls: {}, dataset size: {}'.format(
        trainer_name, len(dataset)))
    
    # st()

    if infi_sampler:
        train_sampler = DistributedSampler(dataset=dataset,
                                        shuffle=True,
                                        drop_last=True)

        loader = DataLoader(dataset,
                            batch_size=batch_size,
                            num_workers=num_workers,
                            drop_last=True,
                            pin_memory=True,
                            persistent_workers=num_workers > 0,
                            sampler=train_sampler)

        while True:
            yield from loader

    else:
        # loader = DataLoader(dataset,
        #                     batch_size=batch_size,
        #                     num_workers=num_workers,
        #                     drop_last=False,
        #                     pin_memory=True,
        #                     persistent_workers=num_workers > 0, 
        #                     shuffle=False)
        st()
        return dataset


def load_eval_rays(file_path="",
                   reso=64,
                   reso_encoder=224,
                   imgnet_normalize=True):
    dataset = MultiViewDataset(file_path,
                               reso,
                               reso_encoder,
                               imgnet_normalize=imgnet_normalize)
    pose_list = dataset.single_pose_list
    ray_list = []
    for pose_fname in pose_list:
        # c2w = dataset.get_c2w(pose_fname).reshape(1,4,4)  #[1, 4, 4]
        # rays_o, rays_d = dataset.gen_rays(c2w)
        # ray_list.append(
        #     [rays_o.unsqueeze(0),
        #      rays_d.unsqueeze(0),
        #      c2w.reshape(-1, 16)])

        c2w = dataset.get_c2w(pose_fname).reshape(16)  #[1, 4, 4]

        c = torch.cat([c2w, dataset.intrinsics],
                      dim=0).reshape(25)  # 25, no '1' dim needed.
        ray_list.append(c)

    return ray_list


def load_eval_data(file_path="",
                   reso=64,
                   reso_encoder=224,
                   batch_size=1,
                   num_workers=1,
                   load_depth=False,
                   preprocess=None,
                   imgnet_normalize=True, 
                   interval=1, **kwargs):

    dataset = MultiViewDataset(file_path,
                               reso,
                               reso_encoder,
                               preprocess=preprocess,
                               load_depth=load_depth,
                               test=True,
                               imgnet_normalize=imgnet_normalize,
                               interval=interval)
    print('eval dataset size: {}'.format(len(dataset)))
    # train_sampler = DistributedSampler(dataset=dataset)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        drop_last=False,
        shuffle=False,
    )
    # sampler=train_sampler)
    return loader


def load_memory_data(file_path="",
                     reso=64,
                     reso_encoder=224,
                     batch_size=1,
                     num_workers=1,
                     load_depth=True,
                     preprocess=None,
                     imgnet_normalize=True):
    # load a single-instance into the memory to speed up training IO
    dataset = MultiViewDataset(file_path,
                               reso,
                               reso_encoder,
                               preprocess=preprocess,
                               load_depth=True,
                               test=False,
                               overfitting=True,
                               imgnet_normalize=imgnet_normalize,
                               overfitting_bs=batch_size)
    logger.log('!!!!!!! memory dataset size: {} !!!!!!'.format(len(dataset)))
    # train_sampler = DistributedSampler(dataset=dataset)
    loader = DataLoader(
        dataset,
        batch_size=len(dataset),
        num_workers=num_workers,
        drop_last=False,
        shuffle=False,
    )

    all_data: dict = next(iter(loader))
    while True:
        start_idx = np.random.randint(0, len(dataset) - batch_size + 1)
        yield {
            k: v[start_idx:start_idx + batch_size]
            for k, v in all_data.items()
        }


class MultiViewDataset(Dataset):

    def __init__(self,
                 file_path,
                 reso,
                 reso_encoder,
                 preprocess=None,
                 classes=False,
                 load_depth=False,
                 test=False,
                 scene_scale=1,
                 overfitting=False,
                 imgnet_normalize=True,
                 dataset_size=-1,
                 overfitting_bs=-1, 
                 interval=1):
        self.file_path = file_path
        self.overfitting = overfitting
        self.scene_scale = scene_scale
        self.reso = reso
        self.reso_encoder = reso_encoder
        self.classes = False
        self.load_depth = load_depth
        self.preprocess = preprocess
        assert not self.classes, "Not support class condition now."

        # self.ins_list = os.listdir(self.file_path)
        # if test: # TODO

        dataset_name = Path(self.file_path).stem.split('_')[0]
        self.dataset_name = dataset_name

        if test:
            # ins_list_file = Path(self.file_path).parent / f'{dataset_name}_test_list.txt' # ? in domain
            if dataset_name == 'chair':
                self.ins_list = sorted(os.listdir(
                    self.file_path))[1:2]  # more diversity
            else:
                self.ins_list = sorted(os.listdir(self.file_path))[
                    0:1]  # the first 1 instance for evaluation reference.
        else:
            # self.ins_list = sorted(Path(self.file_path).glob('[0-8]*'))
            # self.ins_list = Path(self.file_path).glob('*')
            # self.ins_list = list(Path(self.file_path).glob('*'))[:dataset_size]

            # ins_list_file = Path(
            #     self.file_path).parent / f'{dataset_name}s_train_list.txt'
            # assert ins_list_file.exists(), 'add training list for ShapeNet'
            # with open(ins_list_file, 'r') as f:
            #     self.ins_list = [name.strip() for name in f.readlines()]

            # if dataset_name == 'chair':
            ins_list_file = Path(
                self.file_path).parent / f'{dataset_name}_train_list.txt'
            # st()
            assert ins_list_file.exists(), 'add training list for ShapeNet'
            with open(ins_list_file, 'r') as f:
                self.ins_list = [name.strip()
                                 for name in f.readlines()][:dataset_size]
            # else:
            #     self.ins_list = Path(self.file_path).glob('*')

        if overfitting:
            self.ins_list = self.ins_list[:1]

        self.rgb_list = []
        self.pose_list = []
        self.depth_list = []
        self.data_ins_list = []
        self.instance_data_length = -1
        for ins in self.ins_list:
            cur_rgb_path = os.path.join(self.file_path, ins, 'rgb')
            cur_pose_path = os.path.join(self.file_path, ins, 'pose')

            cur_all_fname = sorted([
                t.split('.')[0] for t in os.listdir(cur_rgb_path)
                if 'depth' not in t
            ][::interval])
            if self.instance_data_length == -1:
                self.instance_data_length = len(cur_all_fname)
            else:
                assert len(cur_all_fname) == self.instance_data_length

            # ! check filtered data
            # for idx in range(len(cur_all_fname)):
            #     fname = cur_all_fname[idx]
            #     if not Path(os.path.join(cur_rgb_path, fname + '.png') ).exists():
            #         cur_all_fname.remove(fname)

            # del cur_all_fname[idx]

            if test:
                mid_index = len(cur_all_fname) // 3 * 2
                cur_all_fname.insert(0, cur_all_fname[mid_index])

            self.pose_list += ([
                os.path.join(cur_pose_path, fname + '.txt')
                for fname in cur_all_fname
            ])
            self.rgb_list += ([
                os.path.join(cur_rgb_path, fname + '.png')
                for fname in cur_all_fname
            ])

            self.depth_list += ([
                os.path.join(cur_rgb_path, fname + '_depth0001.exr')
                for fname in cur_all_fname
            ])
            self.data_ins_list += ([ins] * len(cur_all_fname))

        # validate overfitting on images
        if overfitting:
            # bs=9
            # self.pose_list = self.pose_list[::50//9+1]
            # self.rgb_list = self.rgb_list[::50//9+1]
            # self.depth_list = self.depth_list[::50//9+1]
            # bs=6
            # self.pose_list = self.pose_list[::50//6+1]
            # self.rgb_list = self.rgb_list[::50//6+1]
            # self.depth_list = self.depth_list[::50//6+1]
            # bs=3
            assert overfitting_bs != -1
            # bs=1
            # self.pose_list = self.pose_list[25:26]
            # self.rgb_list = self.rgb_list[25:26]
            # self.depth_list = self.depth_list[25:26]

            # uniform pose sampling
            self.pose_list = self.pose_list[::50//overfitting_bs+1]
            self.rgb_list = self.rgb_list[::50//overfitting_bs+1]
            self.depth_list = self.depth_list[::50//overfitting_bs+1]

            # sequentially sampling pose
            # self.pose_list = self.pose_list[25:25+overfitting_bs]
            # self.rgb_list = self.rgb_list[25:25+overfitting_bs]
            # self.depth_list = self.depth_list[25:25+overfitting_bs]

            # duplicate the same pose
            # self.pose_list = [self.pose_list[25]] * overfitting_bs
            # self.rgb_list = [self.rgb_list[25]] * overfitting_bs
            # self.depth_list = [self.depth_list[25]] * overfitting_bs
            # self.pose_list = [self.pose_list[28]] * overfitting_bs
            # self.rgb_list = [self.rgb_list[28]] * overfitting_bs
            # self.depth_list = [self.depth_list[28]] * overfitting_bs

        self.single_pose_list = [
            os.path.join(cur_pose_path, fname + '.txt')
            for fname in cur_all_fname
        ]

        # st()

        # if imgnet_normalize:
        transformations = [
            transforms.ToTensor(),  # [0,1] range
        ]
        if imgnet_normalize:
            transformations.append(
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225))  # type: ignore
            )
        else:
            transformations.append(
                transforms.Normalize((0.5, 0.5, 0.5),
                                     (0.5, 0.5, 0.5)))  # type: ignore

        self.normalize = transforms.Compose(transformations)

        # self.normalize_normalrange = transforms.Compose([
        #     transforms.ToTensor(),# [0,1] range
        #     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        # ])

        fx = fy = 525
        cx = cy = 256  # rendering default K
        factor = self.reso / (cx * 2)  # 128 / 512
        self.fx = fx * factor
        self.fy = fy * factor
        self.cx = cx * factor
        self.cy = cy * factor

        # ! fix scale for triplane ray_sampler(), here we adopt [0,1] uv range, not [0, w] img space range.
        self.cx /= self.reso  # 0.5
        self.cy /= self.reso  # 0.5
        self.fx /= self.reso
        self.fy /= self.reso

        intrinsics = np.array([[self.fx, 0, self.cx], [0, self.fy, self.cy],
                               [0, 0, 1]]).reshape(9)
        # self.intrinsics = torch.from_numpy(intrinsics).float()
        self.intrinsics = intrinsics

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

    def get_c2w(self, pose_fname):
        with open(pose_fname, 'r') as f:
            cam2world = f.readline().strip()
            cam2world = [float(t) for t in cam2world.split(' ')]
        c2w = torch.tensor(cam2world, dtype=torch.float32).reshape(4, 4)
        return c2w

    def gen_rays(self, c2w):
        # Generate rays
        self.h = self.reso
        self.w = self.reso
        yy, xx = torch.meshgrid(
            torch.arange(self.h, dtype=torch.float32) + 0.5,
            torch.arange(self.w, dtype=torch.float32) + 0.5,
            indexing='ij')
        xx = (xx - self.cx) / self.fx
        yy = (yy - self.cy) / self.fy
        zz = torch.ones_like(xx)
        dirs = torch.stack((xx, yy, zz), dim=-1)  # OpenCV convention
        dirs /= torch.norm(dirs, dim=-1, keepdim=True)
        dirs = dirs.reshape(1, -1, 3, 1)
        del xx, yy, zz
        dirs = (c2w[:, None, :3, :3] @ dirs)[..., 0]

        origins = c2w[:, None, :3, 3].expand(-1, self.h * self.w,
                                             -1).contiguous()
        origins = origins.view(-1, 3)
        dirs = dirs.view(-1, 3)

        return origins, dirs

    def read_depth(self, idx):
        depth_path = self.depth_list[idx]
        # image_path = os.path.join(depth_fname, self.image_names[index])
        exr = OpenEXR.InputFile(depth_path)
        header = exr.header()
        size = (header['dataWindow'].max.x - header['dataWindow'].min.x + 1,
                header['dataWindow'].max.y - header['dataWindow'].min.y + 1)
        FLOAT = Imath.PixelType(Imath.PixelType.FLOAT)
        depth_str = exr.channel('B', FLOAT)
        depth = np.frombuffer(depth_str,
                              dtype=np.float32).reshape(size[1],
                                                        size[0])  # H W
        depth = np.nan_to_num(depth, posinf=0, neginf=0)
        depth = depth.reshape(size)

        def resize_depth_mask(depth_to_resize, resolution):
            depth_resized = cv2.resize(depth_to_resize,
                                       (resolution, resolution),
                                       interpolation=cv2.INTER_LANCZOS4)
            #    interpolation=cv2.INTER_AREA)
            return depth_resized > 0  # type: ignore

        fg_mask_reso = resize_depth_mask(depth, self.reso)
        fg_mask_sr = resize_depth_mask(depth, 128)

        # depth = cv2.resize(depth, (self.reso, self.reso),
        #                    interpolation=cv2.INTER_LANCZOS4)
        #    interpolation=cv2.INTER_AREA)
        # depth_mask = depth > 0
        # depth = np.expand_dims(depth, axis=0).reshape(size)
        # return torch.from_numpy(depth)
        return torch.from_numpy(depth), torch.from_numpy(
            fg_mask_reso), torch.from_numpy(fg_mask_sr)

    def load_bbox(self, mask):
        nonzero_value = torch.nonzero(mask)
        height, width = nonzero_value.max(dim=0)[0]
        top, left = nonzero_value.min(dim=0)[0]
        bbox = torch.tensor([top, left, height, width], dtype=torch.float32)
        return bbox

    def __getitem__(self, idx):
        rgb_fname = self.rgb_list[idx]
        pose_fname = self.pose_list[idx]

        raw_img = imageio.imread(rgb_fname)

        if self.preprocess is None:
            img_to_encoder = cv2.resize(raw_img,
                                        (self.reso_encoder, self.reso_encoder),
                                        interpolation=cv2.INTER_LANCZOS4)
            # interpolation=cv2.INTER_AREA)
            img_to_encoder = img_to_encoder[
                ..., :3]  #[3, reso_encoder, reso_encoder]
            img_to_encoder = self.normalize(img_to_encoder)
        else:
            img_to_encoder = self.preprocess(Image.open(rgb_fname))  # clip

        img = cv2.resize(raw_img, (self.reso, self.reso),
                         interpolation=cv2.INTER_LANCZOS4)
        #  interpolation=cv2.INTER_AREA)

        # img_sr = cv2.resize(raw_img, (512, 512), interpolation=cv2.INTER_AREA)
        # img_sr = cv2.resize(raw_img, (256, 256), interpolation=cv2.INTER_AREA) # just as refinement, since eg3d uses 64->128 final resolution
        # img_sr = cv2.resize(raw_img, (128, 128), interpolation=cv2.INTER_AREA) # just as refinement, since eg3d uses 64->128 final resolution
        img_sr = cv2.resize(
            raw_img, (128, 128), interpolation=cv2.INTER_LANCZOS4
        )  # just as refinement, since eg3d uses 64->128 final resolution

        # img = torch.from_numpy(img)[..., :3].permute(
        #     2, 0, 1) / 255.0  #[3, reso, reso]

        img = torch.from_numpy(img)[..., :3].permute(
            2, 0, 1
        ) / 127.5 - 1  #[3, reso, reso], normalize to [-1,1], follow triplane range

        img_sr = torch.from_numpy(img_sr)[..., :3].permute(
            2, 0, 1
        ) / 127.5 - 1  #[3, reso, reso], normalize to [-1,1], follow triplane range

        # c2w = self.get_c2w(pose_fname).reshape(1, 4, 4) #[1, 4, 4]
        # rays_o, rays_d = self.gen_rays(c2w)
        # return img_to_encoder, img, rays_o, rays_d, c2w.reshape(-1)

        c2w = self.get_c2w(pose_fname).reshape(16)  #[1, 4, 4] -> [1, 16]
        # c = np.concatenate([c2w, self.intrinsics], axis=0).reshape(25)  # 25, no '1' dim needed.
        c = torch.cat([c2w, torch.from_numpy(self.intrinsics)],
                      dim=0).reshape(25)  # 25, no '1' dim needed.
        ret_dict = {
            # 'rgb_fname': rgb_fname,
            'img_to_encoder': img_to_encoder,
            'img': img,
            'c': c,
            'img_sr': img_sr,
            # 'ins_name': self.data_ins_list[idx]
        }
        if self.load_depth:
            depth, depth_mask, depth_mask_sr = self.read_depth(idx)
            bbox = self.load_bbox(depth_mask)
            ret_dict.update({
                'depth': depth,
                'depth_mask': depth_mask,
                'depth_mask_sr': depth_mask_sr,
                'bbox': bbox
            })
        # rays_o, rays_d = self.gen_rays(c2w)
        # return img_to_encoder, img, c
        return ret_dict


class MultiViewDatasetforLMDB(MultiViewDataset):

    def __init__(self,
                 file_path,
                 reso,
                 reso_encoder,
                 preprocess=None,
                 classes=False,
                 load_depth=False,
                 test=False,
                 scene_scale=1,
                 overfitting=False,
                 imgnet_normalize=True,
                 dataset_size=-1,
                 overfitting_bs=-1):
        super().__init__(file_path, reso, reso_encoder, preprocess, classes,
                         load_depth, test, scene_scale, overfitting,
                         imgnet_normalize, dataset_size, overfitting_bs)

    def __len__(self):
        return super().__len__()
        # return 100 # for speed debug

    def __getitem__(self, idx):
        # ret_dict = super().__getitem__(idx)
        rgb_fname = self.rgb_list[idx]
        pose_fname = self.pose_list[idx]
        raw_img = imageio.imread(rgb_fname)[..., :3]

        c2w = self.get_c2w(pose_fname).reshape(16)  #[1, 4, 4] -> [1, 16]
        # c = np.concatenate([c2w, self.intrinsics], axis=0).reshape(25)  # 25, no '1' dim needed.
        c = torch.cat([c2w, torch.from_numpy(self.intrinsics)],
                      dim=0).reshape(25)  # 25, no '1' dim needed.

        depth, depth_mask, depth_mask_sr = self.read_depth(idx)
        bbox = self.load_bbox(depth_mask)
        ret_dict = {
            'raw_img': raw_img,
            'c': c,
            'depth': depth,
            # 'depth_mask': depth_mask, # 64x64 here?
            'bbox': bbox
        }
        return ret_dict


def load_data_dryrun(
        file_path="",
        reso=64,
        reso_encoder=224,
        batch_size=1,
        #   shuffle=True,
        num_workers=6,
        load_depth=False,
        preprocess=None,
        imgnet_normalize=True):
    # st()
    dataset = MultiViewDataset(file_path,
                               reso,
                               reso_encoder,
                               test=False,
                               preprocess=preprocess,
                               load_depth=load_depth,
                               imgnet_normalize=imgnet_normalize)
    print('dataset size: {}'.format(len(dataset)))
    # st()
    # train_sampler = DistributedSampler(dataset=dataset)
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        # shuffle=shuffle,
        drop_last=False,
    )
    # sampler=train_sampler)

    return loader


class NovelViewDataset(MultiViewDataset):
    """novel view prediction version.
    """

    def __init__(self,
                 file_path,
                 reso,
                 reso_encoder,
                 preprocess=None,
                 classes=False,
                 load_depth=False,
                 test=False,
                 scene_scale=1,
                 overfitting=False,
                 imgnet_normalize=True,
                 dataset_size=-1,
                 overfitting_bs=-1):
        super().__init__(file_path, reso, reso_encoder, preprocess, classes,
                         load_depth, test, scene_scale, overfitting,
                         imgnet_normalize, dataset_size, overfitting_bs)

    def __getitem__(self, idx):
        input_view = super().__getitem__(
            idx)  # get previous input view results

        # get novel view of the same instance
        novel_view = super().__getitem__(
            (idx // self.instance_data_length) * self.instance_data_length +
            random.randint(0, self.instance_data_length - 1))

        # assert input_view['ins_name'] == novel_view['ins_name'], 'should sample novel view from the same instance'

        input_view.update({f'nv_{k}': v for k, v in novel_view.items()})
        return input_view


def load_data_for_lmdb(
        file_path="",
        reso=64,
        reso_encoder=224,
        batch_size=1,
        #   shuffle=True,
        num_workers=6,
        load_depth=False,
        preprocess=None,
        imgnet_normalize=True,
        dataset_size=-1,
        trainer_name='input_rec'):
    # st()
    # dataset_cls = {
    #     'input_rec': MultiViewDataset,
    #     'nv': NovelViewDataset,
    # }[trainer_name]
    # if 'nv' in trainer_name:
    #     dataset_cls = NovelViewDataset
    # else:
    # dataset_cls = MultiViewDataset
    dataset_cls = MultiViewDatasetforLMDB

    dataset = dataset_cls(file_path,
                          reso,
                          reso_encoder,
                          test=False,
                          preprocess=preprocess,
                          load_depth=load_depth,
                          imgnet_normalize=imgnet_normalize,
                          dataset_size=dataset_size)

    logger.log('dataset_cls: {}, dataset size: {}'.format(
        trainer_name, len(dataset)))
    # train_sampler = DistributedSampler(dataset=dataset, shuffle=True, drop_last=True)
    loader = DataLoader(
        dataset,
        shuffle=False,
        batch_size=batch_size,
        num_workers=num_workers,
        drop_last=False,
        prefetch_factor=2,
        # prefetch_factor=3,
        pin_memory=True,
        persistent_workers=True,
    )
    # sampler=train_sampler)

    # while True:
    #     yield from loader
    return loader, dataset.dataset_name, len(dataset)


class LMDBDataset(Dataset):

    def __init__(self, lmdb_path):
        self.env = lmdb.open(
            lmdb_path,
            readonly=True,
            max_readers=32,
            lock=False,
            readahead=False,
            meminit=False,
        )
        self.num_samples = self.env.stat()['entries']
        # self.start_idx = self.env.stat()['start_idx']
        # self.end_idx = self.env.stat()['end_idx']

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        with self.env.begin(write=False) as txn:
            key = str(idx).encode('utf-8')
            value = txn.get(key)

        sample = pickle.loads(value)
        return sample


def resize_depth_mask(depth_to_resize, resolution):
    depth_resized = cv2.resize(depth_to_resize, (resolution, resolution),
                               interpolation=cv2.INTER_LANCZOS4)
    #    interpolation=cv2.INTER_AREA)
    return depth_resized, depth_resized > 0  # type: ignore


class LMDBDataset_MV(LMDBDataset):

    def __init__(self,
                 lmdb_path,
                 reso,
                 reso_encoder,
                 imgnet_normalize=True,
                 **kwargs):
        super().__init__(lmdb_path)

        self.reso_encoder = reso_encoder
        self.reso = reso

        transformations = [
            transforms.ToTensor(),  # [0,1] range
        ]
        if imgnet_normalize:
            transformations.append(
                transforms.Normalize((0.485, 0.456, 0.406),
                                     (0.229, 0.224, 0.225))  # type: ignore
            )
        else:
            transformations.append(
                transforms.Normalize((0.5, 0.5, 0.5),
                                     (0.5, 0.5, 0.5)))  # type: ignore

        self.normalize = transforms.Compose(transformations)

    def _post_process_sample(self, raw_img, depth):
        
        # if raw_img.shape[-1] == 4: # ! set bg to white
        #     alpha_mask = raw_img[..., -1:] > 0
        #     raw_img = alpha_mask * raw_img[..., :3] + (1-alpha_mask) * np.ones_like(raw_img[..., :3]) * 255
        #     raw_img = raw_img.astype(np.uint8)

        # img_to_encoder = cv2.resize(sample.pop('raw_img'),
        img_to_encoder = cv2.resize(raw_img,
                                    (self.reso_encoder, self.reso_encoder),
                                    interpolation=cv2.INTER_LANCZOS4)
        # interpolation=cv2.INTER_AREA)
        img_to_encoder = img_to_encoder[..., :
                                        3]  #[3, reso_encoder, reso_encoder]
        img_to_encoder = self.normalize(img_to_encoder)

        img = cv2.resize(raw_img, (self.reso, self.reso),
                         interpolation=cv2.INTER_LANCZOS4)

        if img.shape[-1] == 4:
            alpha_mask = img[..., -1:] > 0
            img = alpha_mask * img[..., :3] + (1-alpha_mask) * np.ones_like(img[..., :3]) * 255

        img = torch.from_numpy(img)[..., :3].permute(
            2, 0, 1
        ) / 127.5 - 1  #[3, reso, reso], normalize to [-1,1], follow triplane range

        img_sr = torch.from_numpy(raw_img)[..., :3].permute(
            2, 0, 1
        ) / 127.5 - 1  #[3, reso, reso], normalize to [-1,1], follow triplane range

        # depth
        # fg_mask_reso = resize_depth_mask(sample['depth'], self.reso)
        depth_reso, fg_mask_reso = resize_depth_mask(depth, self.reso)

        return {
            # **sample,
            'img_to_encoder': img_to_encoder,
            'img': img,
            'depth_mask': fg_mask_reso,
            'img_sr': img_sr, 
            'depth': depth_reso,
            # ! no need to load img_sr for now
        }

    def __getitem__(self, idx):
        sample = super().__getitem__(idx)
        # do transformations online

        return self._post_process_sample(sample['raw_img'], sample['depth'])
        # return sample

def load_bytes(inp_bytes, dtype, shape):
    return np.frombuffer(inp_bytes, dtype=dtype).reshape(shape).copy()

# Function to decompress an image using gzip and open with imageio
def decompress_and_open_image_gzip(compressed_data, is_img=False):
    # Decompress the image data using gzip
    decompressed_data = gzip.decompress(compressed_data)

    # Read the decompressed image using imageio
    if is_img:
        image = imageio.v3.imread(io.BytesIO(decompressed_data)).copy()
        return image
    return decompressed_data


# Function to decompress an array using gzip
def decompress_array(compressed_data, shape, dtype):
    # Decompress the array data using gzip
    decompressed_data = gzip.decompress(compressed_data)

    # Convert the decompressed data to a NumPy array
    # arr = np.frombuffer(decompressed_data, dtype=dtype).reshape(shape)

    return load_bytes(decompressed_data, dtype, shape)


class LMDBDataset_MV_Compressed(LMDBDataset_MV):

    def __init__(self,
                 lmdb_path,
                 reso,
                 reso_encoder,
                 imgnet_normalize=True,
                 **kwargs):
        super().__init__(lmdb_path, reso, reso_encoder, imgnet_normalize,
                         **kwargs)
        with self.env.begin(write=False) as txn:
            self.length = int(
                txn.get('length'.encode('utf-8')).decode('utf-8')) - 40

        self.load_image_fn = partial(decompress_and_open_image_gzip,
                                     is_img=True)

    def __len__(self):
        return self.length
    
    def _load_lmdb_data(self, idx):

        with self.env.begin(write=False) as txn:
            raw_img_key = f'{idx}-raw_img'.encode('utf-8')
            raw_img = self.load_image_fn(txn.get(raw_img_key))

            depth_key = f'{idx}-depth'.encode('utf-8')
            depth = decompress_array(txn.get(depth_key), (512,512), np.float32)

            c_key = f'{idx}-c'.encode('utf-8')
            c = decompress_array(txn.get(c_key), (25, ), np.float32)

            bbox_key = f'{idx}-bbox'.encode('utf-8')
            bbox = decompress_array(txn.get(bbox_key), (4, ), np.float32)

        return raw_img, depth, c, bbox

    def __getitem__(self, idx):
        # sample = super(LMDBDataset).__getitem__(idx)

        # do gzip uncompress online
        raw_img, depth, c, bbox  = self._load_lmdb_data(idx)

        return {
            **self._post_process_sample(raw_img, depth), 'c': c,
            'bbox': bbox*(self.reso/64.0),
            # 'depth': depth,
        }


class LMDBDataset_NV_Compressed(LMDBDataset_MV_Compressed):
    def __init__(self, lmdb_path, reso, reso_encoder, imgnet_normalize=True, **kwargs):
        super().__init__(lmdb_path, reso, reso_encoder, imgnet_normalize, **kwargs)
        self.instance_data_length = 50 # 

    def __getitem__(self, idx):
        input_view = super().__getitem__(
            idx)  # get previous input view results

        # get novel view of the same instance
        try:
            novel_view = super().__getitem__(
                (idx // self.instance_data_length) * self.instance_data_length +
                random.randint(0, self.instance_data_length - 1))
        except Exception as e:
            raise NotImplementedError(idx)

        assert input_view['ins_name'] == novel_view['ins_name'], 'should sample novel view from the same instance'

        input_view.update({f'nv_{k}': v for k, v in novel_view.items()})
        return input_view