Delete mvdiffusion/data/objaverse_dataset.py

mvdiffusion/data/objaverse_dataset.py

@@ -1,608 +0,0 @@
-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
-
-def shift_list(lst, n):
-    length = len(lst)
-    n = n % length  # Ensure n is within the range of the list length
-    return lst[-n:] + lst[:-n]
-
-
-class ObjaverseDataset(Dataset):
-    def __init__(self,
-        root_dir: str,
-        num_views: int,
-        bg_color: Any,
-        img_wh: Tuple[int, int],
-        object_list: str,
-        groups_num: int=1,
-        validation: bool = False,
-        random_views: bool = False,
-        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,
-        mix_color_normal: bool = False,
-        random_view_and_domain: bool = False
-        ) -> 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.root_dir = Path(root_dir)
-        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.mix_color_normal = mix_color_normal  # mix load color and normal maps
-        self.random_view_and_domain = random_view_and_domain # load normal or rgb of a single view
-        self.random_views = random_views
-        if not self.random_views:
-            if self.num_views == 4:
-                self.view_types  = ['front', 'right', 'back', 'left']
-            elif self.num_views == 5:
-                self.view_types  = ['front', 'front_right', 'right', 'back', 'left']
-            elif self.num_views == 6 or self.num_views==1:
-                self.view_types  = ['front', 'front_right', 'right', 'back', 'left', 'front_left']
-        else:
-            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
-
-        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.mix_color_normal:
-            self.backup_data = self.__getitem_mix__(0, "9438abf986c7453a9f4df7c34aa2e65b")
-        elif self.random_view_and_domain:
-            self.backup_data = self.__getitem_random_viewanddomain__(0, "9438abf986c7453a9f4df7c34aa2e65b")
-        else:
-            self.backup_data = self.__getitem_norm__(0, "9438abf986c7453a9f4df7c34aa2e65b") # "66b2134b7e3645b29d7c349645291f78")
-
-    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_image(self, img_path, bg_color, alpha, 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 = img.astype(np.float32) / 255. # [0, 1]
-        assert img.shape[-1] == 3 # RGB
-
-        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 load_depth(self, img_path, bg_color, alpha, 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 = img.astype(np.float32) / 65535. # [0, 1]
-
-        img[img > 0.4] = 0
-        img = img / 0.4
-        
-        assert img.ndim == 2 # depth
-        img = np.stack([img]*3, axis=-1)
-
-        if alpha.shape[-1] != 1:
-            alpha = alpha[:, :, None]
-
-        # print(np.max(img[:, :, 0]))
-
-        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_normal(self, img_path, bg_color, alpha, RT_w2c=None, RT_w2c_cond=None, 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
-        normal = np.array(Image.open(img_path).resize(self.img_wh))
-
-        assert normal.shape[-1] == 3 # RGB
-
-        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_mix__(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'
-
-        if random.random() < 0.5:
-            read_color, read_normal, read_depth = True, False, False
-        else:
-            read_color, read_normal, read_depth = False, True, True
-
-        read_normal = read_normal & self.read_normal
-        read_depth = read_depth & self.read_depth
-
-        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()
-
-        cond_alpha = self.load_mask(os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, cond_view)), return_type='np')
-        img_tensors_in = [
-            self.load_image(os.path.join(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, cond_view)), bg_color, cond_alpha, return_type='pt').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(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, view))
-            mask_path = os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, view))
-            normal_path = os.path.join(self.root_dir,  object_name[:3], object_name, "normals_%03d_%s.png" % (set_idx, view))
-            depth_path = os.path.join(self.root_dir,  object_name[:3], object_name, "depth_%03d_%s.png" % (set_idx, view))
-            alpha = self.load_mask(mask_path, return_type='np')
-
-            if read_color:                        
-                img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt")
-                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").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").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
-        camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
-
-        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
-
-        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_random_viewanddomain__(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'
-
-        if random.random() < 0.5:
-            read_color, read_normal, read_depth = True, False, False
-        else:
-            read_color, read_normal, read_depth = False, True, True
-
-        read_normal = read_normal & self.read_normal
-        read_depth = read_depth & self.read_depth
-
-        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()
-
-        cond_alpha = self.load_mask(os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, cond_view)), return_type='np')
-        img_tensors_in = [
-            self.load_image(os.path.join(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, cond_view)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
-        ] * self.num_views
-        img_tensors_out = []
-
-        random_viewidx = random.randint(0, len(view_types)-1)
-
-        for view, tgt_w2c in zip([view_types[random_viewidx]], [tgt_w2cs[random_viewidx]]):
-            img_path = os.path.join(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, view))
-            mask_path = os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, view))
-            normal_path = os.path.join(self.root_dir,  object_name[:3], object_name, "normals_%03d_%s.png" % (set_idx, view))
-            depth_path = os.path.join(self.root_dir,  object_name[:3], object_name, "depth_%03d_%s.png" % (set_idx, view))
-            alpha = self.load_mask(mask_path, return_type='np')
-
-            if read_color:                        
-                img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt")
-                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").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").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
-        camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
-
-        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
-
-        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_norm__(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'
-
-        # if self.random_views:
-        #     view_types = ['front']+random.sample(self.view_types[1:], 3)
-        # else:
-        #     view_types = self.view_types
-
-        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()
-
-        cond_alpha = self.load_mask(os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, cond_view)), return_type='np')
-        img_tensors_in = [
-            self.load_image(os.path.join(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, cond_view)), bg_color, cond_alpha, return_type='pt').permute(2, 0, 1)
-        ] * self.num_views
-        img_tensors_out = []
-        normal_tensors_out = []
-        for view, tgt_w2c in zip(view_types, tgt_w2cs):
-            img_path = os.path.join(self.root_dir,  object_name[:3], object_name, "rgb_%03d_%s.png" % (set_idx, view))
-            mask_path = os.path.join(self.root_dir,  object_name[:3], object_name, "mask_%03d_%s.png" % (set_idx, view))
-            alpha = self.load_mask(mask_path, return_type='np')
-
-            if self.read_color:                        
-                img_tensor = self.load_image(img_path, bg_color, alpha, return_type="pt")
-                img_tensor = img_tensor.permute(2, 0, 1)
-                img_tensors_out.append(img_tensor)
-
-            if self.read_normal:
-                normal_path = os.path.join(self.root_dir,  object_name[:3], object_name, "normals_%03d_%s.png" % (set_idx, view))
-                normal_tensor = self.load_normal(normal_path, bg_color, alpha, RT_w2c=tgt_w2c, RT_w2c_cond=cond_w2c, return_type="pt").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)
-
-        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)
-
-        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
-
-        camera_embeddings = torch.stack([elevations_cond, elevations, azimuths], dim=-1) # (Nv, 3)
-
-        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)
-
-        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,
-            'normal_task_embeddings': normal_task_embeddings,
-            'color_task_embeddings': color_task_embeddings
-        }
-
-    def __getitem__(self, index):
-        
-        try:
-            if self.mix_color_normal:
-                data = self.__getitem_mix__(index)
-            elif self.random_view_and_domain:
-                data = self.__getitem_random_viewanddomain__(index)
-            else:
-                data = self.__getitem_norm__(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)