File size: 5,401 Bytes
c29de8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | '''
-----------------------------------------------------------------------------
Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
'''
import json
import numpy as np
import torch
import torchvision.transforms.functional as torchvision_F
from PIL import Image, ImageFile
from projects.nerf.datasets import base
from projects.nerf.utils import camera
from projects.neuralangelo.utils.misc import gl_to_cv
ImageFile.LOAD_TRUNCATED_IMAGES = True
class Dataset(base.Dataset):
def __init__(self, cfg, is_inference=False):
super().__init__(cfg, is_inference=is_inference, is_test=False)
cfg_data = cfg.data
self.root = cfg_data.root
self.preload = cfg_data.preload
self.H, self.W = cfg_data.val.image_size if is_inference else cfg_data.train.image_size
meta_fname = f"{cfg_data.root}/transforms.json"
with open(meta_fname) as file:
self.meta = json.load(file)
self.list = self.meta["frames"]
if cfg_data[self.split].subset:
subset = cfg_data[self.split].subset
subset_idx = np.linspace(0, len(self.list), subset+1)[:-1].astype(int)
self.list = [self.list[i] for i in subset_idx]
self.num_rays = cfg.model.render.rand_rays
self.readjust = getattr(cfg_data, "readjust", None)
# Preload dataset if possible.
if cfg_data.preload:
self.images = self.preload_threading(self.get_image, cfg_data.num_workers)
self.cameras = self.preload_threading(self.get_camera, cfg_data.num_workers, data_str="cameras")
def __getitem__(self, idx):
"""Process raw data and return processed data in a dictionary.
Args:
idx: The index of the sample of the dataset.
Returns: A dictionary containing the data.
idx (scalar): The index of the sample of the dataset.
image (R tensor): Image idx for per-image embedding.
image (Rx3 tensor): Image with pixel values in [0,1] for supervision.
intr (3x3 tensor): The camera intrinsics of `image`.
pose (3x4 tensor): The camera extrinsics [R,t] of `image`.
"""
# Keep track of sample index for convenience.
sample = dict(idx=idx)
# Get the images.
image, image_size_raw = self.images[idx] if self.preload else self.get_image(idx)
image = self.preprocess_image(image)
# Get the cameras (intrinsics and pose).
intr, pose = self.cameras[idx] if self.preload else self.get_camera(idx)
intr, pose = self.preprocess_camera(intr, pose, image_size_raw)
# Pre-sample ray indices.
if self.split == "train":
ray_idx = torch.randperm(self.H * self.W)[:self.num_rays] # [R]
image_sampled = image.flatten(1, 2)[:, ray_idx].t() # [R,3]
sample.update(
ray_idx=ray_idx,
image_sampled=image_sampled,
intr=intr,
pose=pose,
)
else: # keep image during inference
sample.update(
image=image,
intr=intr,
pose=pose,
)
return sample
def get_image(self, idx):
fpath = self.list[idx]["file_path"]
image_fname = f"{self.root}/{fpath}"
image = Image.open(image_fname)
image.load()
image_size_raw = image.size
return image, image_size_raw
def preprocess_image(self, image):
# Resize the image.
image = image.resize((self.W, self.H))
image = torchvision_F.to_tensor(image)
rgb = image[:3]
return rgb
def get_camera(self, idx):
# Camera intrinsics.
intr = torch.tensor([[self.meta["fl_x"], self.meta["sk_x"], self.meta["cx"]],
[self.meta["sk_y"], self.meta["fl_y"], self.meta["cy"]],
[0, 0, 1]]).float()
# Camera pose.
c2w_gl = torch.tensor(self.list[idx]["transform_matrix"], dtype=torch.float32)
c2w = gl_to_cv(c2w_gl)
if not self.meta['centered']:
# center scene
center = np.array(self.meta["sphere_center"])
if self.readjust:
center += np.array(getattr(self.readjust, "center", [0]))
c2w[:3, -1] -= center
if not self.meta['scaled']:
# scale scene
scale = np.array(self.meta["sphere_radius"])
if self.readjust:
scale *= getattr(self.readjust, "scale", 1.)
c2w[:3, -1] /= scale
w2c = camera.Pose().invert(c2w[:3])
return intr, w2c
def preprocess_camera(self, intr, pose, image_size_raw):
# Adjust the intrinsics according to the resized image.
intr = intr.clone()
raw_W, raw_H = image_size_raw
intr[0] *= self.W / raw_W
intr[1] *= self.H / raw_H
return intr, pose
|