from typing import Dict
import numpy as np
from omegaconf import DictConfig, ListConfig
import torch
from torch.utils.data import Dataset
from pathlib import Path
import json
from PIL import Image
from torchvision import transforms
from einops import rearrange
from typing import Literal, Tuple, Optional, Any
import cv2
import random
import json
import os, sys
import math
import PIL.Image
from .normal_utils import trans_normal, normal2img, img2normal
import pdb
from .depth_utils import scale_depth_to_model
import traceback
class ObjaverseDataset(Dataset):
def __init__(self,
root_dir_ortho: str,
root_dir_persp: str,
pred_ortho: bool,
pred_persp: bool,
num_views: int,
bg_color: Any,
img_wh: Tuple[int, int],
object_list: str,
groups_num: int=1,
validation: bool = False,
data_view_num: int = 6,
num_validation_samples: int = 64,
num_samples: Optional[int] = None,
invalid_list: Optional[str] = None,
trans_norm_system: bool = True, # if True, transform all normals map into the cam system of front view
augment_data: bool = False,
read_normal: bool = True,
read_color: bool = False,
read_depth: bool = False,
read_mask: bool = False,
pred_type: str = 'color',
suffix: str = 'png',
subscene_tag: int = 2,
load_cam_type: bool = False,
backup_scene: str = "0306b42594fb447ca574f597352d4b56",
ortho_crop_size: int = 360,
persp_crop_size: int = 440,
load_switcher: bool = True
) -> None:
"""Create a dataset from a folder of images.
If you pass in a root directory it will be searched for images
ending in ext (ext can be a list)
"""
self.load_cam_type = load_cam_type
self.root_dir_ortho = Path(root_dir_ortho)
self.root_dir_persp = Path(root_dir_persp)
self.pred_ortho = pred_ortho
self.pred_persp = pred_persp
self.num_views = num_views
self.bg_color = bg_color
self.validation = validation
self.num_samples = num_samples
self.trans_norm_system = trans_norm_system
self.augment_data = augment_data
self.invalid_list = invalid_list
self.groups_num = groups_num
print("augment data: ", self.augment_data)
self.img_wh = img_wh
self.read_normal = read_normal
self.read_color = read_color
self.read_depth = read_depth
self.read_mask = read_mask
self.pred_type = pred_type # load type
self.suffix = suffix
self.subscene_tag = subscene_tag
self.view_types = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
self.fix_cam_pose_dir = "./mvdiffusion/data/fixed_poses/nine_views"
self.fix_cam_poses = self.load_fixed_poses() # world2cam matrix
self.ortho_crop_size = ortho_crop_size
self.persp_crop_size = persp_crop_size
self.load_switcher = load_switcher
if object_list is not None:
with open(object_list) as f:
self.objects = json.load(f)
self.objects = [os.path.basename(o).replace(".glb", "") for o in self.objects]
else:
self.objects = os.listdir(self.root_dir)
self.objects = sorted(self.objects)
if self.invalid_list is not None:
with open(self.invalid_list) as f:
self.invalid_objects = json.load(f)
self.invalid_objects = [os.path.basename(o).replace(".glb", "") for o in self.invalid_objects]
else:
self.invalid_objects = []
self.all_objects = set(self.objects) - (set(self.invalid_objects) & set(self.objects))
self.all_objects = list(self.all_objects)
if not validation:
self.all_objects = self.all_objects[:-num_validation_samples]
else:
self.all_objects = self.all_objects[-num_validation_samples:]
if num_samples is not None:
self.all_objects = self.all_objects[:num_samples]
print("loading ", len(self.all_objects), " objects in the dataset")
if self.pred_type == 'color':
self.backup_data = self.__getitem_color__(0, backup_scene)
elif self.pred_type == 'normal_depth':
self.backup_data = self.__getitem_normal_depth__(0, backup_scene)
elif self.pred_type == 'mixed_rgb_normal_depth':
self.backup_data = self.__getitem_mixed__(0, backup_scene)
elif self.pred_type == 'mixed_color_normal':
self.backup_data = self.__getitem_image_normal_mixed__(0, backup_scene)
elif self.pred_type == 'mixed_rgb_noraml_mask':
self.backup_data = self.__getitem_mixed_rgb_noraml_mask__(0, backup_scene)
elif self.pred_type == 'joint_color_normal':
self.backup_data = self.__getitem_joint_rgb_noraml__(0, backup_scene)
def __len__(self):
return len(self.objects)*self.total_view
def load_fixed_poses(self):
poses = {}
for face in self.view_types:
RT = np.loadtxt(os.path.join(self.fix_cam_pose_dir,'%03d_%s_RT.txt'%(0, face)))
poses[face] = RT
return poses
def cartesian_to_spherical(self, 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(self, target_RT, cond_RT):
R, T = target_RT[:3, :3], target_RT[:, -1]
T_target = -R.T @ T # change to cam2world
R, T = cond_RT[:3, :3], cond_RT[:, -1]
T_cond = -R.T @ T
theta_cond, azimuth_cond, z_cond = self.cartesian_to_spherical(T_cond[None, :])
theta_target, azimuth_target, z_target = self.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_theta, d_azimuth
def get_bg_color(self):
if self.bg_color == 'white':
bg_color = np.array([1., 1., 1.], dtype=np.float32)
elif self.bg_color == 'black':
bg_color = np.array([0., 0., 0.], dtype=np.float32)
elif self.bg_color == 'gray':
bg_color = np.array([0.5, 0.5, 0.5], dtype=np.float32)
elif self.bg_color == 'random':
bg_color = np.random.rand(3)
elif self.bg_color == 'three_choices':
white = np.array([1., 1., 1.], dtype=np.float32)
black = np.array([0., 0., 0.], dtype=np.float32)
gray = np.array([0.5, 0.5, 0.5], dtype=np.float32)
bg_color = random.choice([white, black, gray])
elif isinstance(self.bg_color, float):
bg_color = np.array([self.bg_color] * 3, dtype=np.float32)
else:
raise NotImplementedError
return bg_color
def load_mask(self, img_path, return_type='np'):
# not using cv2 as may load in uint16 format
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
# pil always returns uint8
img = np.array(Image.open(img_path).resize(self.img_wh))
img = np.float32(img > 0)
assert len(np.shape(img)) == 2
if return_type == "np":
pass
elif return_type == "pt":
img = torch.from_numpy(img)
else:
raise NotImplementedError
return img
def load_mask_from_rgba(self, img_path, camera_type):
img = Image.open(img_path)
if camera_type == 'ortho':
left = (img.width - self.ortho_crop_size) // 2
right = (img.width + self.ortho_crop_size) // 2
top = (img.height - self.ortho_crop_size) // 2
bottom = (img.height + self.ortho_crop_size) // 2
img = img.crop((left, top, right, bottom))
if camera_type == 'persp':
left = (img.width - self.persp_crop_size) // 2
right = (img.width + self.persp_crop_size) // 2
top = (img.height - self.persp_crop_size) // 2
bottom = (img.height + self.persp_crop_size) // 2
img = img.crop((left, top, right, bottom))
img = img.resize(self.img_wh)
img = np.array(img).astype(np.float32) / 255. # [0, 1]
assert img.shape[-1] == 4 # must RGBA
alpha = img[:, :, 3:]
if alpha.shape[-1] != 1:
alpha = alpha[:, :, None]
return alpha
def load_image(self, img_path, bg_color, alpha, return_type='np', camera_type=None, read_depth=False, center_crop_size=None):
# not using cv2 as may load in uint16 format
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
# pil always returns uint8
img = Image.open(img_path)
if center_crop_size == None:
if camera_type == 'ortho':
left = (img.width - self.ortho_crop_size) // 2
right = (img.width + self.ortho_crop_size) // 2
top = (img.height - self.ortho_crop_size) // 2
bottom = (img.height + self.ortho_crop_size) // 2
img = img.crop((left, top, right, bottom))
if camera_type == 'persp':
left = (img.width - self.persp_crop_size) // 2
right = (img.width + self.persp_crop_size) // 2
top = (img.height - self.persp_crop_size) // 2
bottom = (img.height + self.persp_crop_size) // 2
img = img.crop((left, top, right, bottom))
else:
center_crop_size = min(center_crop_size, 512)
left = (img.width - center_crop_size) // 2
right = (img.width + center_crop_size) // 2
top = (img.height - center_crop_size) // 2
bottom = (img.height + center_crop_size) // 2
img = img.crop((left, top, right, bottom))
img = img.resize(self.img_wh)
img = np.array(img).astype(np.float32) / 255. # [0, 1]
assert img.shape[-1] == 3 or img.shape[-1] == 4 # RGB or RGBA
if alpha is None and img.shape[-1] == 4:
alpha = img[:, :, 3:]
img = img[:, :, :3]
if alpha.shape[-1] != 1:
alpha = alpha[:, :, None]
if read_depth:
bg_color = np.array([1., 1., 1.], dtype=np.float32)
img = img[...,:3] * alpha + bg_color * (1 - alpha)
if return_type == "np":
pass
elif return_type == "pt":
img = torch.from_numpy(img)
else:
raise NotImplementedError
return img
def load_depth(self, img_path, bg_color, alpha, return_type='np', camera_type=None):
# not using cv2 as may load in uint16 format
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
# pil always returns uint8
depth_bg_color = np.array([1., 1., 1.], dtype=np.float32) # white for depth
depth_map = Image.open(img_path)
if camera_type == 'ortho':
left = (depth_map.width - self.ortho_crop_size) // 2
right = (depth_map.width + self.ortho_crop_size) // 2
top = (depth_map.height - self.ortho_crop_size) // 2
bottom = (depth_map.height + self.ortho_crop_size) // 2
depth_map = depth_map.crop((left, top, right, bottom))
if camera_type == 'persp':
left = (depth_map.width - self.persp_crop_size) // 2
right = (depth_map.width + self.persp_crop_size) // 2
top = (depth_map.height - self.persp_crop_size) // 2
bottom = (depth_map.height + self.persp_crop_size) // 2
depth_map = depth_map.crop((left, top, right, bottom))
depth_map = depth_map.resize(self.img_wh)
depth_map = np.array(depth_map)
# scale the depth map:
depth_map = scale_depth_to_model(depth_map.astype(np.float32))
# depth_map = depth_map / 65535. # [0, 1]
# depth_map[depth_map > 0.4] = 0
# depth_map = depth_map / 0.4
assert depth_map.ndim == 2 # depth
img = np.stack([depth_map]*3, axis=-1)
if alpha.shape[-1] != 1:
alpha = alpha[:, :, None]
# print(np.max(img[:, :, 0]))
# print(np.min(img[...,:3]), np.max(img[...,:3]))
img = img[...,:3] * alpha + depth_bg_color * (1 - alpha)
if return_type == "np":
pass
elif return_type == "pt":
img = torch.from_numpy(img)
else:
raise NotImplementedError
return img
def transform_mask_as_input(self, mask, return_type='np'):
# mask = mask * 255
# print(np.max(mask))
# mask = mask.resize(self.img_wh)
mask = np.squeeze(mask, axis=-1)
assert mask.ndim == 2 #
mask = np.stack([mask]*3, axis=-1)
if return_type == "np":
pass
elif return_type == "pt":
mask = torch.from_numpy(mask)
else:
raise NotImplementedError
return mask
def load_normal(self, img_path, bg_color, alpha, RT_w2c=None, RT_w2c_cond=None, return_type='np', camera_type=None, center_crop_size=None):
# not using cv2 as may load in uint16 format
# img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) # [0, 255]
# img = cv2.resize(img, self.img_wh, interpolation=cv2.INTER_CUBIC)
# pil always returns uint8
# normal = Image.open(img_path)
img = Image.open(img_path)
if center_crop_size == None:
if camera_type == 'ortho':
left = (img.width - self.ortho_crop_size) // 2
right = (img.width + self.ortho_crop_size) // 2
top = (img.height - self.ortho_crop_size) // 2
bottom = (img.height + self.ortho_crop_size) // 2
img = img.crop((left, top, right, bottom))
if camera_type == 'persp':
left = (img.width - self.persp_crop_size) // 2
right = (img.width + self.persp_crop_size) // 2
top = (img.height - self.persp_crop_size) // 2
bottom = (img.height + self.persp_crop_size) // 2
img = img.crop((left, top, right, bottom))
else:
center_crop_size = min(center_crop_size, 512)
left = (img.width - center_crop_size) // 2
right = (img.width + center_crop_size) // 2
top = (img.height - center_crop_size) // 2
bottom = (img.height + center_crop_size) // 2
img = img.crop((left, top, right, bottom))
normal = np.array(img.resize(self.img_wh))
assert normal.shape[-1] == 3 or normal.shape[-1] == 4 # RGB or RGBA
if alpha is None and normal.shape[-1] == 4:
alpha = normal[:, :, 3:] / 255.
normal = normal[:, :, :3]
normal = trans_normal(img2normal(normal), RT_w2c, RT_w2c_cond)
img = (normal*0.5 + 0.5).astype(np.float32) # [0, 1]
if alpha.shape[-1] != 1:
alpha = alpha[:, :, None]
img = img[...,:3] * alpha + bg_color * (1 - alpha)
if return_type == "np":
pass
elif return_type == "pt":
img = torch.from_numpy(img)
else:
raise NotImplementedError
return img
def __len__(self):
return len(self.all_objects)
def __getitem_color__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index % len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
# ! if you would like predict depth; modify here
read_color, read_normal, read_depth = True, False, False
assert (read_color and (read_normal or read_depth)) is False
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
return_type='np')
else:
cond_alpha = None
img_tensors_in = [
self.load_image(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
bg_color, cond_alpha, return_type='pt', camera_type=load_cam_type).permute(2, 0, 1)
] * self.num_views
img_tensors_out = []
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"mask_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"normals_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if read_normal:
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c,
return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
img_tensors_out.append(normal_tensor)
if read_depth:
depth_tensor = self.load_depth(depth_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
img_tensors_out.append(depth_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
# if self.pred_ortho and self.pred_persp:
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class] * self.num_views, dim=0) # (Nv, 2)
color_class = torch.tensor([0, 1]).float()
color_task_embeddings = torch.stack([color_class] * self.num_views, dim=0) # (Nv, 2)
if read_normal or read_depth:
task_embeddings = normal_task_embeddings
if read_color:
task_embeddings = color_task_embeddings
# print(elevations)
# print(azimuths)
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'camera_embeddings': camera_embeddings,
'task_embeddings': task_embeddings
}
def __getitem_normal_depth__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index%len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
else:
cond_alpha = None
# img_tensors_in = [
# self.load_image(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
# ] * self.num_views
img_tensors_out = []
normal_tensors_out = []
depth_tensors_out = []
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if self.read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if self.read_normal:
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
normal_tensors_out.append(normal_tensor)
if self.read_depth:
if alpha is None:
alpha = self.load_mask_from_rgba(img_path, camera_type=load_cam_type)
depth_tensor = self.load_depth(depth_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
depth_tensors_out.append(depth_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
img_tensors_in = img_tensors_out
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
if self.read_color:
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
if self.read_normal:
normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
if self.read_depth:
depth_tensors_out = torch.stack(depth_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
# if self.pred_ortho and self.pred_persp:
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
color_class = torch.tensor([0, 1]).float()
depth_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'normals_out': normal_tensors_out,
'depth_out': depth_tensors_out,
'camera_embeddings': camera_embeddings,
'normal_task_embeddings': normal_task_embeddings,
'depth_task_embeddings': depth_task_embeddings
}
def __getitem_mixed_rgb_noraml_mask__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index%len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
else:
cond_alpha = None
img_tensors_out = []
normal_tensors_out = []
depth_tensors_out = []
random_select = random.random()
read_color, read_normal, read_mask = [random_select < 1 / 3, 1 / 3 <= random_select <= 2 / 3,
random_select > 2 / 3]
# print(read_color, read_normal, read_depth)
assert sum([read_color, read_normal, read_mask]) == 1, "Only one variable should be True"
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type, read_depth=False)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if read_normal:
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
img_tensors_out.append(normal_tensor)
if read_mask:
if alpha is None:
alpha = self.load_mask_from_rgba(img_path, camera_type=load_cam_type)
mask_tensor = self.transform_mask_as_input(alpha, return_type='pt').permute(2, 0, 1)
img_tensors_out.append(mask_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
if self.load_switcher: # rgb input, use domain switcher to control the output type
img_tensors_in = [
self.load_image(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"normals_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
bg_color, cond_alpha, RT_w2c=cond_w2c, RT_w2c_cond=cond_w2c, return_type='pt', camera_type=load_cam_type).permute(
2, 0, 1)
] * self.num_views
color_class = torch.tensor([0, 1]).float()
color_task_embeddings = torch.stack([color_class] * self.num_views, dim=0) # (Nv, 2)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class] * self.num_views, dim=0) # (Nv, 2)
mask_class = torch.tensor([1, 1]).float()
mask_task_embeddings = torch.stack([mask_class] * self.num_views, dim=0)
if read_color:
task_embeddings = color_task_embeddings
# img_tensors_out = depth_tensors_out
elif read_normal:
task_embeddings = normal_task_embeddings
# img_tensors_out = normal_tensors_out
elif read_mask:
task_embeddings = mask_task_embeddings
# img_tensors_out = depth_tensors_out
else: # for stage 1 training, the input and the output are in the same domain
img_tensors_in = [img_tensors_out[0]] * self.num_views
empty_class = torch.tensor([0, 0]).float() # empty task
empty_task_embeddings = torch.stack([empty_class] * self.num_views, dim=0)
task_embeddings = empty_task_embeddings
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'normals_out': normal_tensors_out,
'depth_out': depth_tensors_out,
'camera_embeddings': camera_embeddings,
'task_embeddings': task_embeddings,
}
def __getitem_mixed__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index%len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
else:
cond_alpha = None
# img_tensors_in = [
# self.load_image(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
# ] * self.num_views
img_tensors_out = []
normal_tensors_out = []
depth_tensors_out = []
random_select = random.random()
read_color, read_normal, read_depth = [random_select < 1 / 3, 1 / 3 <= random_select <= 2 / 3,
random_select > 2 / 3]
# print(read_color, read_normal, read_depth)
assert sum([read_color, read_normal, read_depth]) == 1, "Only one variable should be True"
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type, read_depth=read_depth)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if read_normal:
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
img_tensors_out.append(normal_tensor)
if read_depth:
if alpha is None:
alpha = self.load_mask_from_rgba(img_path, camera_type=load_cam_type)
depth_tensor = self.load_depth(depth_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
img_tensors_out.append(depth_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
img_tensors_in = [
self.load_image(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
bg_color, cond_alpha, return_type='pt', camera_type=load_cam_type, read_depth=read_depth).permute(
2, 0, 1)
] * self.num_views
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
# if self.read_color:
# img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
# if self.read_normal:
# normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
# if self.read_depth:
# depth_tensors_out = torch.stack(depth_tensors_out, dim=0).float() # (Nv, 3, H, W)
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
# if self.pred_ortho and self.pred_persp:
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
color_class = torch.tensor([0, 1]).float()
color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
depth_class = torch.tensor([1, 1]).float()
depth_task_embeddings = torch.stack([depth_class]*self.num_views, dim=0)
if read_color:
task_embeddings = color_task_embeddings
# img_tensors_out = depth_tensors_out
elif read_normal:
task_embeddings = normal_task_embeddings
# img_tensors_out = normal_tensors_out
elif read_depth:
task_embeddings = depth_task_embeddings
# img_tensors_out = depth_tensors_out
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'normals_out': normal_tensors_out,
'depth_out': depth_tensors_out,
'camera_embeddings': camera_embeddings,
'task_embeddings': task_embeddings,
}
def __getitem_image_normal_mixed__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index%len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
# get crop size for each mv instance:
center_crop_size = 0
for view in view_types:
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
img = Image.open(img_path)
img = img.resize([512,512])
img = np.array(img).astype(np.float32) / 255. # [0, 1]
max_w_h = self.cal_single_view_crop(img)
center_crop_size = max(center_crop_size, max_w_h)
center_crop_size = center_crop_size * 4. / 3.
center_crop_size = center_crop_size + (random.random()-0.5) * 10.
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
else:
cond_alpha = None
# img_tensors_in = [
# self.load_image(os.path.join(self.root_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
# ] * self.num_views
img_tensors_out = []
normal_tensors_out = []
depth_tensors_out = []
random_select = random.random()
read_color, read_normal = [random_select < 1 / 2, 1 / 2 <= random_select <= 1]
# print(read_color, read_normal, read_depth)
assert sum([read_color, read_normal]) == 1, "Only one variable should be True"
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type, read_depth=False, center_crop_size=center_crop_size)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if read_normal:
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt", camera_type=load_cam_type, center_crop_size=center_crop_size).permute(2, 0, 1)
img_tensors_out.append(normal_tensor)
# if read_depth:
# if alpha is None:
# alpha = self.load_mask_from_rgba(img_path, camera_type=load_cam_type)
# depth_tensor = self.load_depth(depth_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type).permute(2, 0, 1)
# img_tensors_out.append(depth_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
img_tensors_in = [
self.load_image(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
bg_color, cond_alpha, return_type='pt', camera_type=load_cam_type, read_depth=False, center_crop_size=center_crop_size).permute(
2, 0, 1)
] * self.num_views
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
# if self.read_color:
# img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
# if self.read_normal:
# normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
# if self.read_depth:
# depth_tensors_out = torch.stack(depth_tensors_out, dim=0).float() # (Nv, 3, H, W)
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
# if self.pred_ortho and self.pred_persp:
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
color_class = torch.tensor([0, 1]).float()
color_task_embeddings = torch.stack([color_class]*self.num_views, dim=0) # (Nv, 2)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class]*self.num_views, dim=0) # (Nv, 2)
# depth_class = torch.tensor([1, 1]).float()
# depth_task_embeddings = torch.stack([depth_class]*self.num_views, dim=0)
if read_color:
task_embeddings = color_task_embeddings
# img_tensors_out = depth_tensors_out
elif read_normal:
task_embeddings = normal_task_embeddings
# img_tensors_out = normal_tensors_out
# elif read_depth:
# task_embeddings = depth_task_embeddings
# img_tensors_out = depth_tensors_out
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'normals_out': normal_tensors_out,
'depth_out': depth_tensors_out,
'camera_embeddings': camera_embeddings,
'task_embeddings': task_embeddings,
}
def cal_single_view_crop(self, image):
assert np.shape(image)[-1] == 4 # RGBA
# Extract the alpha channel (transparency) and the object (RGB channels)
alpha_channel = image[:, :, 3]
# Find the bounding box coordinates of the object
coords = cv2.findNonZero(alpha_channel)
x, y, width, height = cv2.boundingRect(coords)
return max(width, height)
def __getitem_joint_rgb_noraml__(self, index, debug_object=None):
if debug_object is not None:
object_name = debug_object #
set_idx = random.sample(range(0, self.groups_num), 1)[0] # without replacement
else:
object_name = self.all_objects[index%len(self.all_objects)]
set_idx = 0
if self.augment_data:
cond_view = random.sample(self.view_types, k=1)[0]
else:
cond_view = 'front'
assert self.pred_ortho or self.pred_persp
if self.pred_ortho and self.pred_persp:
if random.random() < 0.5:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
else:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
elif self.pred_ortho and not self.pred_persp:
load_dir = self.root_dir_ortho
load_cam_type = 'ortho'
elif self.pred_persp and not self.pred_ortho:
load_dir = self.root_dir_persp
load_cam_type = 'persp'
view_types = self.view_types
cond_w2c = self.fix_cam_poses[cond_view]
tgt_w2cs = [self.fix_cam_poses[view] for view in view_types]
elevations = []
azimuths = []
# get the bg color
bg_color = self.get_bg_color()
if self.read_mask:
cond_alpha = self.load_mask(os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, cond_view, self.suffix)), return_type='np')
else:
cond_alpha = None
img_tensors_out = []
normal_tensors_out = []
read_color, read_normal = True, True
# print(read_color, read_normal, read_depth)
# get crop size for each mv instance:
center_crop_size = 0
for view in view_types:
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
img = Image.open(img_path)
img = img.resize([512,512])
img = np.array(img).astype(np.float32) / 255. # [0, 1]
max_w_h = self.cal_single_view_crop(img)
center_crop_size = max(center_crop_size, max_w_h)
center_crop_size = center_crop_size * 4. / 3.
center_crop_size = center_crop_size + (random.random()-0.5) * 10.
for view, tgt_w2c in zip(view_types, tgt_w2cs):
img_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "rgb_%03d_%s.%s" % (set_idx, view, self.suffix))
mask_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "mask_%03d_%s.%s" % (set_idx, view, self.suffix))
depth_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "depth_%03d_%s.%s" % (set_idx, view, self.suffix))
if self.read_mask:
alpha = self.load_mask(mask_path, return_type='np')
else:
alpha = None
if read_color:
img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt", camera_type=load_cam_type, read_depth=False, center_crop_size=center_crop_size)
img_tensor = img_tensor.permute(2, 0, 1)
img_tensors_out.append(img_tensor)
if read_normal:
normal_path = os.path.join(load_dir, object_name[:self.subscene_tag], object_name, "normals_%03d_%s.%s" % (set_idx, view, self.suffix))
normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt", camera_type=load_cam_type, center_crop_size=center_crop_size).permute(2, 0, 1)
normal_tensors_out.append(normal_tensor)
# evelations, azimuths
elevation, azimuth = self.get_T(tgt_w2c, cond_w2c)
elevations.append(elevation)
azimuths.append(azimuth)
if self.load_switcher: # rgb input, use domain switcher to control the output type
img_tensors_in = [
self.load_image(os.path.join(load_dir, object_name[:self.subscene_tag], object_name,
"rgb_%03d_%s.%s" % (set_idx, cond_view, self.suffix)),
bg_color, cond_alpha, return_type='pt', camera_type=load_cam_type,
read_depth=False, center_crop_size=center_crop_size).permute(
2, 0, 1)
] * self.num_views
color_class = torch.tensor([0, 1]).float()
color_task_embeddings = torch.stack([color_class] * self.num_views, dim=0) # (Nv, 2)
normal_class = torch.tensor([1, 0]).float()
normal_task_embeddings = torch.stack([normal_class] * self.num_views, dim=0) # (Nv, 2)
if read_color:
task_embeddings = color_task_embeddings
# img_tensors_out = depth_tensors_out
elif read_normal:
task_embeddings = normal_task_embeddings
# img_tensors_out = normal_tensors_out
else: # for stage 1 training, the input and the output are in the same domain
img_tensors_in = [img_tensors_out[0]] * self.num_views
empty_class = torch.tensor([0, 0]).float() # empty task
empty_task_embeddings = torch.stack([empty_class] * self.num_views, dim=0)
task_embeddings = empty_task_embeddings
img_tensors_in = torch.stack(img_tensors_in, dim=0).float() # (Nv, 3, H, W)
img_tensors_out = torch.stack(img_tensors_out, dim=0).float() # (Nv, 3, H, W)
normal_tensors_out = torch.stack(normal_tensors_out, dim=0).float() # (Nv, 3, H, W)
elevations = torch.as_tensor(elevations).float().squeeze(1)
azimuths = torch.as_tensor(azimuths).float().squeeze(1)
elevations_cond = torch.as_tensor([0] * self.num_views).float() # fixed only use 4 views to train
if load_cam_type == 'ortho':
cam_type_emb = torch.tensor([0, 1]).expand(self.num_views, -1)
else:
cam_type_emb = torch.tensor([1, 0]).expand(self.num_views, -1)
camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1)
if self.load_cam_type:
camera_embeddings = torch.cat((camera_embeddings, cam_type_emb), dim=-1) # (Nv, 5)
return {
'elevations_cond': elevations_cond,
'elevations_cond_deg': torch.rad2deg(elevations_cond),
'elevations': elevations,
'azimuths': azimuths,
'elevations_deg': torch.rad2deg(elevations),
'azimuths_deg': torch.rad2deg(azimuths),
'imgs_in': img_tensors_in,
'imgs_out': img_tensors_out,
'normals_out': normal_tensors_out,
'camera_embeddings': camera_embeddings,
'color_task_embeddings': color_task_embeddings,
'normal_task_embeddings': normal_task_embeddings
}
def __getitem__(self, index):
try:
if self.pred_type == 'color':
data = self.backup_data = self.__getitem_color__(index)
elif self.pred_type == 'normal_depth':
data = self.backup_data = self.__getitem_normal_depth__(index)
elif self.pred_type == 'mixed_rgb_normal_depth':
data = self.backup_data = self.__getitem_mixed__(index)
elif self.pred_type == 'mixed_color_normal':
data = self.backup_data = self.__getitem_image_normal_mixed__(index)
elif self.pred_type == 'mixed_rgb_noraml_mask':
data = self.backup_data = self.__getitem_mixed_rgb_noraml_mask__(index)
elif self.pred_type == 'joint_color_normal':
data = self.backup_data = self.__getitem_joint_rgb_noraml__(index)
return data
except:
print("load error ", self.all_objects[index%len(self.all_objects)])
return self.backup_data
class ConcatDataset(torch.utils.data.Dataset):
def __init__(self, datasets, weights):
self.datasets = datasets
self.weights = weights
self.num_datasets = len(datasets)
def __getitem__(self, i):
chosen = random.choices(self.datasets, self.weights, k=1)[0]
return chosen[i]
def __len__(self):
return max(len(d) for d in self.datasets)
if __name__ == "__main__":
train_dataset = ObjaverseDataset(
root_dir="/ghome/l5/xxlong/.objaverse/hf-objaverse-v1/renderings",
size=(128, 128),
ext="hdf5",
default_trans=torch.zeros(3),
return_paths=False,
total_view=8,
validation=False,
object_list=None,
views_mode='fourviews'
)
data0 = train_dataset[0]
data1 = train_dataset[50]
# print(data)