bluestyle97
6 months ago
3 changed files with 326 additions and 1 deletions
@ -0,0 +1,326 @@ |
|||||
|
import os |
||||
|
import math |
||||
|
import json |
||||
|
from pathlib import Path |
||||
|
|
||||
|
import cv2 |
||||
|
import numpy as np |
||||
|
from PIL import Image |
||||
|
import webdataset as wds |
||||
|
import pytorch_lightning as pl |
||||
|
|
||||
|
import torch |
||||
|
import torch.nn.functional as F |
||||
|
from torch.utils.data import Dataset |
||||
|
from torch.utils.data import DataLoader |
||||
|
from torch.utils.data.distributed import DistributedSampler |
||||
|
from torchvision import transforms |
||||
|
|
||||
|
from src.utils.train_util import instantiate_from_config |
||||
|
from src.utils.camera_util import ( |
||||
|
FOV_to_intrinsics, |
||||
|
center_looking_at_camera_pose, |
||||
|
get_circular_camera_poses, |
||||
|
) |
||||
|
|
||||
|
|
||||
|
class DataModuleFromConfig(pl.LightningDataModule): |
||||
|
def __init__( |
||||
|
self, |
||||
|
batch_size=8, |
||||
|
num_workers=4, |
||||
|
train=None, |
||||
|
validation=None, |
||||
|
test=None, |
||||
|
**kwargs, |
||||
|
): |
||||
|
super().__init__() |
||||
|
|
||||
|
self.batch_size = batch_size |
||||
|
self.num_workers = num_workers |
||||
|
|
||||
|
self.dataset_configs = dict() |
||||
|
if train is not None: |
||||
|
self.dataset_configs['train'] = train |
||||
|
if validation is not None: |
||||
|
self.dataset_configs['validation'] = validation |
||||
|
if test is not None: |
||||
|
self.dataset_configs['test'] = test |
||||
|
|
||||
|
def setup(self, stage): |
||||
|
|
||||
|
if stage in ['fit']: |
||||
|
self.datasets = dict((k, instantiate_from_config(self.dataset_configs[k])) for k in self.dataset_configs) |
||||
|
else: |
||||
|
raise NotImplementedError |
||||
|
|
||||
|
def train_dataloader(self): |
||||
|
|
||||
|
sampler = DistributedSampler(self.datasets['train']) |
||||
|
return wds.WebLoader(self.datasets['train'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler) |
||||
|
|
||||
|
def val_dataloader(self): |
||||
|
|
||||
|
sampler = DistributedSampler(self.datasets['validation']) |
||||
|
return wds.WebLoader(self.datasets['validation'], batch_size=1, num_workers=self.num_workers, shuffle=False, sampler=sampler) |
||||
|
|
||||
|
def test_dataloader(self): |
||||
|
|
||||
|
return wds.WebLoader(self.datasets['test'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False) |
||||
|
|
||||
|
|
||||
|
class ObjaverseData(Dataset): |
||||
|
def __init__(self, |
||||
|
root_dir='objaverse/', |
||||
|
meta_fname='valid_paths.json', |
||||
|
input_image_dir='rendering_random_32views', |
||||
|
target_image_dir='rendering_random_32views', |
||||
|
input_view_num=6, |
||||
|
target_view_num=2, |
||||
|
total_view_n=32, |
||||
|
fov=50, |
||||
|
camera_rotation=True, |
||||
|
validation=False, |
||||
|
): |
||||
|
self.root_dir = Path(root_dir) |
||||
|
self.input_image_dir = input_image_dir |
||||
|
self.target_image_dir = target_image_dir |
||||
|
|
||||
|
self.input_view_num = input_view_num |
||||
|
self.target_view_num = target_view_num |
||||
|
self.total_view_n = total_view_n |
||||
|
self.fov = fov |
||||
|
self.camera_rotation = camera_rotation |
||||
|
|
||||
|
with open(os.path.join(root_dir, meta_fname)) as f: |
||||
|
filtered_dict = json.load(f) |
||||
|
paths = filtered_dict['good_objs'] |
||||
|
self.paths = paths |
||||
|
|
||||
|
self.depth_scale = 5.0 |
||||
|
|
||||
|
total_objects = len(self.paths) |
||||
|
print('============= length of dataset %d =============' % len(self.paths)) |
||||
|
|
||||
|
def __len__(self): |
||||
|
return len(self.paths) |
||||
|
|
||||
|
def load_im(self, path, color): |
||||
|
''' |
||||
|
replace background pixel with random color in rendering |
||||
|
''' |
||||
|
pil_img = Image.open(path) |
||||
|
|
||||
|
image = np.asarray(pil_img, dtype=np.float32) / 255. |
||||
|
alpha = image[:, :, 3:] |
||||
|
image = image[:, :, :3] * alpha + color * (1 - alpha) |
||||
|
|
||||
|
image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float() |
||||
|
alpha = torch.from_numpy(alpha).permute(2, 0, 1).contiguous().float() |
||||
|
return image, alpha |
||||
|
|
||||
|
def __getitem__(self, index): |
||||
|
# load data |
||||
|
while True: |
||||
|
input_image_path = os.path.join(self.root_dir, self.input_image_dir, self.paths[index]) |
||||
|
target_image_path = os.path.join(self.root_dir, self.target_image_dir, self.paths[index]) |
||||
|
|
||||
|
indices = np.random.choice(range(self.total_view_n), self.input_view_num + self.target_view_num, replace=False) |
||||
|
input_indices = indices[:self.input_view_num] |
||||
|
target_indices = indices[self.input_view_num:] |
||||
|
|
||||
|
'''background color, default: white''' |
||||
|
bg_white = [1., 1., 1.] |
||||
|
bg_black = [0., 0., 0.] |
||||
|
|
||||
|
image_list = [] |
||||
|
alpha_list = [] |
||||
|
depth_list = [] |
||||
|
normal_list = [] |
||||
|
pose_list = [] |
||||
|
|
||||
|
try: |
||||
|
input_cameras = np.load(os.path.join(input_image_path, 'cameras.npz'))['cam_poses'] |
||||
|
for idx in input_indices: |
||||
|
image, alpha = self.load_im(os.path.join(input_image_path, '%03d.png' % idx), bg_white) |
||||
|
normal, _ = self.load_im(os.path.join(input_image_path, '%03d_normal.png' % idx), bg_black) |
||||
|
depth = cv2.imread(os.path.join(input_image_path, '%03d_depth.png' % idx), cv2.IMREAD_UNCHANGED) / 255.0 * self.depth_scale |
||||
|
depth = torch.from_numpy(depth).unsqueeze(0) |
||||
|
pose = input_cameras[idx] |
||||
|
pose = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0) |
||||
|
|
||||
|
image_list.append(image) |
||||
|
alpha_list.append(alpha) |
||||
|
depth_list.append(depth) |
||||
|
normal_list.append(normal) |
||||
|
pose_list.append(pose) |
||||
|
|
||||
|
target_cameras = np.load(os.path.join(target_image_path, 'cameras.npz'))['cam_poses'] |
||||
|
for idx in target_indices: |
||||
|
image, alpha = self.load_im(os.path.join(target_image_path, '%03d.png' % idx), bg_white) |
||||
|
normal, _ = self.load_im(os.path.join(target_image_path, '%03d_normal.png' % idx), bg_black) |
||||
|
depth = cv2.imread(os.path.join(target_image_path, '%03d_depth.png' % idx), cv2.IMREAD_UNCHANGED) / 255.0 * self.depth_scale |
||||
|
depth = torch.from_numpy(depth).unsqueeze(0) |
||||
|
pose = target_cameras[idx] |
||||
|
pose = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0) |
||||
|
|
||||
|
image_list.append(image) |
||||
|
alpha_list.append(alpha) |
||||
|
depth_list.append(depth) |
||||
|
normal_list.append(normal) |
||||
|
pose_list.append(pose) |
||||
|
|
||||
|
except Exception as e: |
||||
|
print(e) |
||||
|
index = np.random.randint(0, len(self.paths)) |
||||
|
continue |
||||
|
|
||||
|
break |
||||
|
|
||||
|
images = torch.stack(image_list, dim=0).float() # (6+V, 3, H, W) |
||||
|
alphas = torch.stack(alpha_list, dim=0).float() # (6+V, 1, H, W) |
||||
|
depths = torch.stack(depth_list, dim=0).float() # (6+V, 1, H, W) |
||||
|
normals = torch.stack(normal_list, dim=0).float() # (6+V, 3, H, W) |
||||
|
w2cs = torch.from_numpy(np.stack(pose_list, axis=0)).float() # (6+V, 4, 4) |
||||
|
c2ws = torch.linalg.inv(w2cs).float() |
||||
|
|
||||
|
normals = normals * 2.0 - 1.0 |
||||
|
normals = F.normalize(normals, dim=1) |
||||
|
normals = (normals + 1.0) / 2.0 |
||||
|
normals = torch.lerp(torch.zeros_like(normals), normals, alphas) |
||||
|
|
||||
|
# random rotation along z axis |
||||
|
if self.camera_rotation: |
||||
|
degree = np.random.uniform(0, math.pi * 2) |
||||
|
rot = torch.tensor([ |
||||
|
[np.cos(degree), -np.sin(degree), 0, 0], |
||||
|
[np.sin(degree), np.cos(degree), 0, 0], |
||||
|
[0, 0, 1, 0], |
||||
|
[0, 0, 0, 1], |
||||
|
]).unsqueeze(0).float() |
||||
|
c2ws = torch.matmul(rot, c2ws) |
||||
|
|
||||
|
# rotate normals |
||||
|
N, _, H, W = normals.shape |
||||
|
normals = normals * 2.0 - 1.0 |
||||
|
normals = torch.matmul(rot[:, :3, :3], normals.view(N, 3, -1)).view(N, 3, H, W) |
||||
|
normals = F.normalize(normals, dim=1) |
||||
|
normals = (normals + 1.0) / 2.0 |
||||
|
normals = torch.lerp(torch.zeros_like(normals), normals, alphas) |
||||
|
|
||||
|
# random scaling |
||||
|
if np.random.rand() < 0.5: |
||||
|
scale = np.random.uniform(0.8, 1.0) |
||||
|
c2ws[:, :3, 3] *= scale |
||||
|
depths *= scale |
||||
|
|
||||
|
# instrinsics of perspective cameras |
||||
|
K = FOV_to_intrinsics(self.fov) |
||||
|
Ks = K.unsqueeze(0).repeat(self.input_view_num + self.target_view_num, 1, 1).float() |
||||
|
|
||||
|
data = { |
||||
|
'input_images': images[:self.input_view_num], # (6, 3, H, W) |
||||
|
'input_alphas': alphas[:self.input_view_num], # (6, 1, H, W) |
||||
|
'input_depths': depths[:self.input_view_num], # (6, 1, H, W) |
||||
|
'input_normals': normals[:self.input_view_num], # (6, 3, H, W) |
||||
|
'input_c2ws': c2ws[:self.input_view_num], # (6, 4, 4) |
||||
|
'input_Ks': Ks[:self.input_view_num], # (6, 3, 3) |
||||
|
|
||||
|
# lrm generator input and supervision |
||||
|
'target_images': images[self.input_view_num:], # (V, 3, H, W) |
||||
|
'target_alphas': alphas[self.input_view_num:], # (V, 1, H, W) |
||||
|
'target_depths': depths[self.input_view_num:], # (V, 1, H, W) |
||||
|
'target_normals': normals[self.input_view_num:], # (V, 3, H, W) |
||||
|
'target_c2ws': c2ws[self.input_view_num:], # (V, 4, 4) |
||||
|
'target_Ks': Ks[self.input_view_num:], # (V, 3, 3) |
||||
|
|
||||
|
'depth_available': 1, |
||||
|
} |
||||
|
return data |
||||
|
|
||||
|
|
||||
|
class ValidationData(Dataset): |
||||
|
def __init__(self, |
||||
|
root_dir='objaverse/', |
||||
|
input_view_num=6, |
||||
|
input_image_size=256, |
||||
|
fov=30, |
||||
|
): |
||||
|
self.root_dir = Path(root_dir) |
||||
|
self.input_view_num = input_view_num |
||||
|
self.input_image_size = input_image_size |
||||
|
self.fov = fov |
||||
|
|
||||
|
self.paths = sorted(os.listdir(self.root_dir)) |
||||
|
print('============= length of dataset %d =============' % len(self.paths)) |
||||
|
|
||||
|
cam_distance = 4.0 |
||||
|
azimuths = np.array([30, 90, 150, 210, 270, 330]) |
||||
|
elevations = np.array([20, -10, 20, -10, 20, -10]) |
||||
|
azimuths = np.deg2rad(azimuths) |
||||
|
elevations = np.deg2rad(elevations) |
||||
|
|
||||
|
x = cam_distance * np.cos(elevations) * np.cos(azimuths) |
||||
|
y = cam_distance * np.cos(elevations) * np.sin(azimuths) |
||||
|
z = cam_distance * np.sin(elevations) |
||||
|
|
||||
|
cam_locations = np.stack([x, y, z], axis=-1) |
||||
|
cam_locations = torch.from_numpy(cam_locations).float() |
||||
|
c2ws = center_looking_at_camera_pose(cam_locations) |
||||
|
self.c2ws = c2ws.float() |
||||
|
self.Ks = FOV_to_intrinsics(self.fov).unsqueeze(0).repeat(6, 1, 1).float() |
||||
|
|
||||
|
render_c2ws = get_circular_camera_poses(M=8, radius=cam_distance, elevation=20.0) |
||||
|
render_Ks = FOV_to_intrinsics(self.fov).unsqueeze(0).repeat(render_c2ws.shape[0], 1, 1) |
||||
|
self.render_c2ws = render_c2ws.float() |
||||
|
self.render_Ks = render_Ks.float() |
||||
|
|
||||
|
def __len__(self): |
||||
|
return len(self.paths) |
||||
|
|
||||
|
def load_im(self, path, color): |
||||
|
''' |
||||
|
replace background pixel with random color in rendering |
||||
|
''' |
||||
|
pil_img = Image.open(path) |
||||
|
pil_img = pil_img.resize((self.input_image_size, self.input_image_size), resample=Image.BICUBIC) |
||||
|
|
||||
|
image = np.asarray(pil_img, dtype=np.float32) / 255. |
||||
|
if image.shape[-1] == 4: |
||||
|
alpha = image[:, :, 3:] |
||||
|
image = image[:, :, :3] * alpha + color * (1 - alpha) |
||||
|
else: |
||||
|
alpha = np.ones_like(image[:, :, :1]) |
||||
|
|
||||
|
image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float() |
||||
|
alpha = torch.from_numpy(alpha).permute(2, 0, 1).contiguous().float() |
||||
|
return image, alpha |
||||
|
|
||||
|
def __getitem__(self, index): |
||||
|
# load data |
||||
|
input_image_path = os.path.join(self.root_dir, self.paths[index]) |
||||
|
|
||||
|
'''background color, default: white''' |
||||
|
bkg_color = [1.0, 1.0, 1.0] |
||||
|
|
||||
|
image_list = [] |
||||
|
alpha_list = [] |
||||
|
|
||||
|
for idx in range(self.input_view_num): |
||||
|
image, alpha = self.load_im(os.path.join(input_image_path, f'{idx:03d}.png'), bkg_color) |
||||
|
image_list.append(image) |
||||
|
alpha_list.append(alpha) |
||||
|
|
||||
|
images = torch.stack(image_list, dim=0).float() |
||||
|
alphas = torch.stack(alpha_list, dim=0).float() |
||||
|
|
||||
|
data = { |
||||
|
'input_images': images, |
||||
|
'input_alphas': alphas, |
||||
|
'input_c2ws': self.c2ws, |
||||
|
'input_Ks': self.Ks, |
||||
|
|
||||
|
'render_c2ws': self.render_c2ws, |
||||
|
'render_Ks': self.render_Ks, |
||||
|
} |
||||
|
return data |
||||
Loading…
Reference in new issue