Spaces:
Paused
Paused
| from torch.utils.data import Dataset | |
| import os | |
| import json | |
| import numpy as np | |
| import torch | |
| import imageio | |
| import math | |
| import cv2 | |
| from torchvision import transforms | |
| def cartesian_to_spherical(xyz): | |
| ptsnew = np.hstack((xyz, np.zeros(xyz.shape))) | |
| xy = xyz[:,0]**2 + xyz[:,1]**2 | |
| z = np.sqrt(xy + xyz[:,2]**2) | |
| theta = np.arctan2(np.sqrt(xy), xyz[:,2]) # for elevation angle defined from Z-axis down | |
| #ptsnew[:,4] = np.arctan2(xyz[:,2], np.sqrt(xy)) # for elevation angle defined from XY-plane up | |
| azimuth = np.arctan2(xyz[:,1], xyz[:,0]) | |
| return np.array([theta, azimuth, z]) | |
| def get_T(T_target, T_cond): | |
| theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :]) | |
| theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :]) | |
| d_theta = theta_target - theta_cond | |
| d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi) | |
| d_z = z_target - z_cond | |
| d_T = torch.tensor([d_theta.item(), math.sin(d_azimuth.item()), math.cos(d_azimuth.item()), d_z.item()]) | |
| return d_T | |
| def get_spherical(T_target, T_cond): | |
| theta_cond, azimuth_cond, z_cond = cartesian_to_spherical(T_cond[None, :]) | |
| theta_target, azimuth_target, z_target = cartesian_to_spherical(T_target[None, :]) | |
| d_theta = theta_target - theta_cond | |
| d_azimuth = (azimuth_target - azimuth_cond) % (2 * math.pi) | |
| d_z = z_target - z_cond | |
| d_T = torch.tensor([math.degrees(d_theta.item()), math.degrees(d_azimuth.item()), d_z.item()]) | |
| return d_T | |
| class RTMV(Dataset): | |
| def __init__(self, root_dir='datasets/RTMV/google_scanned',\ | |
| first_K=64, resolution=256, load_target=False): | |
| self.root_dir = root_dir | |
| self.scene_list = sorted(next(os.walk(root_dir))[1]) | |
| self.resolution = resolution | |
| self.first_K = first_K | |
| self.load_target = load_target | |
| def __len__(self): | |
| return len(self.scene_list) | |
| def __getitem__(self, idx): | |
| scene_dir = os.path.join(self.root_dir, self.scene_list[idx]) | |
| with open(os.path.join(scene_dir, 'transforms.json'), "r") as f: | |
| meta = json.load(f) | |
| imgs = [] | |
| poses = [] | |
| for i_img in range(self.first_K): | |
| meta_img = meta['frames'][i_img] | |
| if i_img == 0 or self.load_target: | |
| img_path = os.path.join(scene_dir, meta_img['file_path']) | |
| img = imageio.imread(img_path) | |
| img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR) | |
| imgs.append(img) | |
| c2w = meta_img['transform_matrix'] | |
| poses.append(c2w) | |
| imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs | |
| imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2) | |
| imgs = imgs * 2 - 1. # convert to stable diffusion range | |
| poses = torch.tensor(np.array(poses).astype(np.float32)) | |
| return imgs, poses | |
| def blend_rgba(self, img): | |
| img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB | |
| return img | |
| class GSO(Dataset): | |
| def __init__(self, root_dir='datasets/GoogleScannedObjects',\ | |
| split='val', first_K=5, resolution=256, load_target=False, name='render_mvs'): | |
| self.root_dir = root_dir | |
| with open(os.path.join(root_dir, '%s.json' % split), "r") as f: | |
| self.scene_list = json.load(f) | |
| self.resolution = resolution | |
| self.first_K = first_K | |
| self.load_target = load_target | |
| self.name = name | |
| def __len__(self): | |
| return len(self.scene_list) | |
| def __getitem__(self, idx): | |
| scene_dir = os.path.join(self.root_dir, self.scene_list[idx]) | |
| with open(os.path.join(scene_dir, 'transforms_%s.json' % self.name), "r") as f: | |
| meta = json.load(f) | |
| imgs = [] | |
| poses = [] | |
| for i_img in range(self.first_K): | |
| meta_img = meta['frames'][i_img] | |
| if i_img == 0 or self.load_target: | |
| img_path = os.path.join(scene_dir, meta_img['file_path']) | |
| img = imageio.imread(img_path) | |
| img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR) | |
| imgs.append(img) | |
| c2w = meta_img['transform_matrix'] | |
| poses.append(c2w) | |
| imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs | |
| mask = imgs[:, :, :, -1] | |
| imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2) | |
| imgs = imgs * 2 - 1. # convert to stable diffusion range | |
| poses = torch.tensor(np.array(poses).astype(np.float32)) | |
| return imgs, poses | |
| def blend_rgba(self, img): | |
| img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB | |
| return img | |
| class WILD(Dataset): | |
| def __init__(self, root_dir='data/nerf_wild',\ | |
| first_K=33, resolution=256, load_target=False): | |
| self.root_dir = root_dir | |
| self.scene_list = sorted(next(os.walk(root_dir))[1]) | |
| self.resolution = resolution | |
| self.first_K = first_K | |
| self.load_target = load_target | |
| def __len__(self): | |
| return len(self.scene_list) | |
| def __getitem__(self, idx): | |
| scene_dir = os.path.join(self.root_dir, self.scene_list[idx]) | |
| with open(os.path.join(scene_dir, 'transforms_train.json'), "r") as f: | |
| meta = json.load(f) | |
| imgs = [] | |
| poses = [] | |
| for i_img in range(self.first_K): | |
| meta_img = meta['frames'][i_img] | |
| if i_img == 0 or self.load_target: | |
| img_path = os.path.join(scene_dir, meta_img['file_path']) | |
| img = imageio.imread(img_path + '.png') | |
| img = cv2.resize(img, (self.resolution, self.resolution), interpolation = cv2.INTER_LINEAR) | |
| imgs.append(img) | |
| c2w = meta_img['transform_matrix'] | |
| poses.append(c2w) | |
| imgs = (np.array(imgs) / 255.).astype(np.float32) # (RGBA) imgs | |
| imgs = torch.tensor(self.blend_rgba(imgs)).permute(0, 3, 1, 2) | |
| imgs = imgs * 2 - 1. # convert to stable diffusion range | |
| poses = torch.tensor(np.array(poses).astype(np.float32)) | |
| return imgs, poses | |
| def blend_rgba(self, img): | |
| img = img[..., :3] * img[..., -1:] + (1. - img[..., -1:]) # blend A to RGB | |
| return img |