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