src/data/objaverse.py

import os, sys
import math
import json
import glm
from pathlib import Path

import random
import numpy as np
from PIL import Image
import webdataset as wds
import pytorch_lightning as pl
import sys
from src.utils import obj, render_utils
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from torch.utils.data.distributed import DistributedSampler
import random
import itertools
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,
)
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
import re

def spherical_camera_pose(azimuths: np.ndarray, elevations: np.ndarray, radius=2.5):
    azimuths = np.deg2rad(azimuths)
    elevations = np.deg2rad(elevations)

    xs = radius * np.cos(elevations) * np.cos(azimuths)
    ys = radius * np.cos(elevations) * np.sin(azimuths)
    zs = radius * np.sin(elevations)

    cam_locations = np.stack([xs, ys, zs], axis=-1)
    cam_locations = torch.from_numpy(cam_locations).float()

    c2ws = center_looking_at_camera_pose(cam_locations)
    return c2ws

def find_matching_files(base_path, idx):
    formatted_idx = '%03d' % idx
    pattern = re.compile(r'^%s_\d+\.png$' % formatted_idx)
    matching_files = []
    
    if os.path.exists(base_path):
        for filename in os.listdir(base_path):
            if pattern.match(filename):
                matching_files.append(filename)
                
    return os.path.join(base_path, matching_files[0])

def load_mipmap(env_path):
    diffuse_path = os.path.join(env_path, "diffuse.pth")
    diffuse = torch.load(diffuse_path, map_location=torch.device('cpu'))

    specular = []
    for i in range(6):
        specular_path = os.path.join(env_path, f"specular_{i}.pth")
        specular_tensor = torch.load(specular_path, map_location=torch.device('cpu'))
        specular.append(specular_tensor)
    return [specular, diffuse]

def convert_to_white_bg(image, write_bg=True):
    alpha = image[:, :, 3:]
    if write_bg:
        return image[:, :, :3] * alpha + 1. * (1 - alpha)
    else:
        return image[:, :, :3] * alpha
    
def load_obj(path, return_attributes=False, scale_factor=1.0):
    return obj.load_obj(path, clear_ks=True, mtl_override=None, return_attributes=return_attributes, scale_factor=scale_factor)

def custom_collate_fn(batch):
    return batch


def collate_fn_wrapper(batch):
    return custom_collate_fn(batch)

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 custom_collate_fn(self, batch):
        collated_batch = {}
        for key in batch[0].keys():
            if key == 'input_env' or key == 'target_env':
                collated_batch[key] = [d[key] for d in batch]
            else:
                collated_batch[key] = torch.stack([d[key] for d in batch], dim=0)
        return collated_batch
    
    def convert_to_white_bg(self, image):
        alpha = image[:, :, 3:]
        return image[:, :, :3] * alpha + 1. * (1 - alpha)
    
    def load_obj(self, path):
        return obj.load_obj(path, clear_ks=True, mtl_override=None)

    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, collate_fn=collate_fn_wrapper)

    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, collate_fn=collate_fn_wrapper)

    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_highQuality',
        light_dir= 'env_mipmap',
        input_view_num=6,
        target_view_num=4,
        total_view_n=18,
        distance=3.5,
        fov=50,
        camera_random=False,
        validation=False,
    ):
        self.root_dir = Path(root_dir)
        self.light_dir = light_dir
        self.all_env_name = []
        for temp_dir in os.listdir(light_dir):
            if os.listdir(os.path.join(self.light_dir, temp_dir)):
                self.all_env_name.append(temp_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_random = camera_random
        
        self.train_res = [512, 512]
        self.cam_near_far = [0.1, 1000.0]
        self.fov_rad = np.deg2rad(fov)
        self.fov_deg = fov
        self.spp = 1
        self.cam_radius = distance
        self.layers = 1
        
        numbers = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
        self.combinations = list(itertools.product(numbers, repeat=2))
        
        self.paths = os.listdir(self.root_dir)
        
        # with open("BJ_Mesh_list.json", 'r') as file:
        #     self.paths = json.load(file)

        print('total training object num:', len(self.paths))

        self.depth_scale = 6.0
            
        total_objects = len(self.paths)
        print('============= length of dataset %d =============' % total_objects)

    def __len__(self):
        return len(self.paths)
    
    def load_obj(self, path):
        return obj.load_obj(path, clear_ks=True, mtl_override=None)
    
    def sample_spherical(self, phi, theta, cam_radius):
        theta = np.deg2rad(theta)
        phi = np.deg2rad(phi)   

        z = cam_radius * np.cos(phi) * np.sin(theta)
        x = cam_radius * np.sin(phi) * np.sin(theta)
        y = cam_radius * np.cos(theta)
 
        return x, y, z
    
    def _random_scene(self, cam_radius, fov_rad):
        iter_res = self.train_res
        proj_mtx = render_utils.perspective(fov_rad, iter_res[1] / iter_res[0], self.cam_near_far[0], self.cam_near_far[1])

        azimuths = random.uniform(0, 360)
        elevations = random.uniform(30, 150)
        mv_embedding = spherical_camera_pose(azimuths, 90-elevations, cam_radius)
        x, y, z = self.sample_spherical(azimuths, elevations, cam_radius)
        eye = glm.vec3(x, y, z)
        at = glm.vec3(0.0, 0.0, 0.0)
        up = glm.vec3(0.0, 1.0, 0.0)
        view_matrix = glm.lookAt(eye, at, up)
        mv = torch.from_numpy(np.array(view_matrix))
        mvp    = proj_mtx @ (mv)  #w2c
        campos = torch.linalg.inv(mv)[:3, 3]
        return mv[None, ...], mvp[None, ...], campos[None, ...], mv_embedding[None, ...], iter_res, self.spp # Add batch dimension
        
    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 load_albedo(self, path, color, mask):
        '''
        replace background pixel with random color in rendering
        '''
        pil_img = Image.open(path)

        image = np.asarray(pil_img, dtype=np.float32) / 255.
        image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float()

        color = torch.ones_like(image)
        image = image * mask + color * (1 - mask)
        return image
    
    def convert_to_white_bg(self, image):
        alpha = image[:, :, 3:]
        return image[:, :, :3] * alpha + 1. * (1 - alpha)

    def calculate_fov(self, initial_distance, initial_fov, new_distance):
        initial_fov_rad = math.radians(initial_fov)
        
        height = 2 * initial_distance * math.tan(initial_fov_rad / 2)
        
        new_fov_rad = 2 * math.atan(height / (2 * new_distance))
        
        new_fov = math.degrees(new_fov_rad)
        
        return new_fov
    
    def __getitem__(self, index):
        obj_path = os.path.join(self.root_dir, self.paths[index])
        mesh_attributes = torch.load(obj_path, map_location=torch.device('cpu'))
        pose_list = []
        env_list = []
        material_list = []
        camera_pos = []
        c2w_list = []
        camera_embedding_list = []
        random_env = False
        random_mr = False
        if random.random() > 0.5:
            random_env = True
        if random.random() > 0.5:
            random_mr = True
        selected_env = random.randint(0, len(self.all_env_name)-1)
        materials = random.choice(self.combinations)
        if self.camera_random:
            random_perturbation = random.uniform(-1.5, 1.5)
            cam_radius = self.cam_radius + random_perturbation
            fov_deg = self.calculate_fov(initial_distance=self.cam_radius, initial_fov=self.fov_deg, new_distance=cam_radius)
            fov_rad = np.deg2rad(fov_deg)
        else:
            cam_radius = self.cam_radius
            fov_rad = self.fov_rad
            fov_deg = self.fov_deg

        if len(self.input_view_num) >= 1:
            input_view_num = random.choice(self.input_view_num)
        else:
            input_view_num = self.input_view_num
        for _ in range(input_view_num + self.target_view_num):
            mv, mvp, campos, mv_mebedding, iter_res, iter_spp = self._random_scene(cam_radius, fov_rad)
            if random_env:
                selected_env = random.randint(0, len(self.all_env_name)-1)
            env_path = os.path.join(self.light_dir, self.all_env_name[selected_env])
            env = load_mipmap(env_path)
            if random_mr:
                materials = random.choice(self.combinations)
            pose_list.append(mvp)
            camera_pos.append(campos)
            c2w_list.append(mv)
            env_list.append(env)
            material_list.append(materials)
            camera_embedding_list.append(mv_mebedding)
        data = {
            'mesh_attributes': mesh_attributes,
            'input_view_num': input_view_num,
            'target_view_num': self.target_view_num,
            'obj_path': obj_path,
            'pose_list': pose_list,
            'camera_pos': camera_pos,
            'c2w_list': c2w_list,
            'env_list': env_list,
            'material_list': material_list,
            'camera_embedding_list': camera_embedding_list,
            'fov_deg':fov_deg,
            'raduis': cam_radius
        }
        
        return data

class ValidationData(Dataset):
    def __init__(self,
        root_dir='objaverse/',
        input_view_num=6,
        input_image_size=320,
        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.light_dir = 'env_mipmap'

        # with open('Mesh_list.json') as f:
        #     filtered_dict = json.load(f)
        
        self.paths = os.listdir(self.root_dir)
            
        # self.paths = filtered_dict
        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 load_mat(self, path, color):
        '''
        replace background pixel with random color in rendering
        '''
        pil_img = Image.open(path)
        pil_img = pil_img.resize((384,384), 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 load_albedo(self, path, color, mask):
        '''
        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.
        image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float()

        color = torch.ones_like(image)
        image = image * mask + color * (1 - mask)
        return image
    
    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 = []
        albedo_list = []
        alpha_list = []
        specular_list = []
        diffuse_list = []
        metallic_list = []
        roughness_list = []
        
        exist_comb_list = []
        for subfolder in os.listdir(input_image_path):
            found_numeric_subfolder=False
            subfolder_path = os.path.join(input_image_path, subfolder)
            if os.path.isdir(subfolder_path) and '_' in subfolder and 'specular' not in subfolder and 'diffuse' not in subfolder:
                try:
                    parts = subfolder.split('_')
                    float(parts[0])  # 尝试将分隔符前后的字符串转换为浮点数
                    float(parts[1])
                    found_numeric_subfolder = True
                except ValueError:
                    continue
            if found_numeric_subfolder:
                exist_comb_list.append(subfolder)
                
        selected_one_comb = random.choice(exist_comb_list)


        for idx in range(self.input_view_num):
            img_path = find_matching_files(os.path.join(input_image_path, selected_one_comb, 'rgb'), idx)
            albedo_path = img_path.replace('rgb', 'albedo')
            metallic_path = img_path.replace('rgb', 'metallic')
            roughness_path = img_path.replace('rgb', 'roughness')
            
            image, alpha = self.load_im(img_path, bkg_color)
            albedo = self.load_albedo(albedo_path, bkg_color, alpha)
            metallic,_ = self.load_mat(metallic_path, bkg_color)
            roughness,_ = self.load_mat(roughness_path, bkg_color)
            
            light_num = os.path.basename(img_path).split('_')[1].split('.')[0]
            light_path = os.path.join(self.light_dir, str(int(light_num)+1))
            
            specular, diffuse = load_mipmap(light_path)
            
            image_list.append(image)
            alpha_list.append(alpha)
            albedo_list.append(albedo)
            metallic_list.append(metallic)
            roughness_list.append(roughness)
            specular_list.append(specular)
            diffuse_list.append(diffuse)
        
        images = torch.stack(image_list, dim=0).float()
        alphas = torch.stack(alpha_list, dim=0).float()
        albedo = torch.stack(albedo_list, dim=0).float()    
        metallic = torch.stack(metallic_list, dim=0).float()    
        roughness = torch.stack(roughness_list, dim=0).float() 

        data = {
            'input_images': images,
            'input_alphas': alphas,
            'input_c2ws': self.c2ws,
            'input_Ks': self.Ks,
            
            'input_albedos': albedo[:self.input_view_num], 
            'input_metallics': metallic[:self.input_view_num], 
            'input_roughness': roughness[:self.input_view_num], 
            
            'specular': specular_list[:self.input_view_num],
            'diffuse': diffuse_list[:self.input_view_num],

            'render_c2ws': self.render_c2ws,
            'render_Ks': self.render_Ks,
        }
        return data


if __name__ == '__main__':
    dataset = ObjaverseData()
    dataset.new(1)