souging/TRELLIS_TextTo3D

1import os2from PIL import Image3import json4import numpy as np5import torch6import utils3d.torch7from ..modules.sparse.basic import SparseTensor8from .components import StandardDatasetBase9 10 11class SLat2Render(StandardDatasetBase):12    """13    Dataset for Structured Latent and rendered images.14    15    Args:16        roots (str): paths to the dataset17        image_size (int): size of the image18        latent_model (str): latent model name19        min_aesthetic_score (float): minimum aesthetic score20        max_num_voxels (int): maximum number of voxels21    """22    def __init__(23        self,24        roots: str,25        image_size: int,26        latent_model: str,27        min_aesthetic_score: float = 5.0,28        max_num_voxels: int = 32768,29    ):30        self.image_size = image_size31        self.latent_model = latent_model32        self.min_aesthetic_score = min_aesthetic_score33        self.max_num_voxels = max_num_voxels34        self.value_range = (0, 1)35        36        super().__init__(roots)37        38    def filter_metadata(self, metadata):39        stats = {}40        metadata = metadata[metadata[f'latent_{self.latent_model}']]41        stats['With latent'] = len(metadata)42        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]43        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)44        metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]45        stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)46        return metadata, stats47 48    def _get_image(self, root, instance):49        with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:50            metadata = json.load(f)51        n_views = len(metadata['frames'])52        view = np.random.randint(n_views)53        metadata = metadata['frames'][view]54        fov = metadata['camera_angle_x']55        intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))56        c2w = torch.tensor(metadata['transform_matrix'])57        c2w[:3, 1:3] *= -158        extrinsics = torch.inverse(c2w)59 60        image_path = os.path.join(root, 'renders', instance, metadata['file_path'])61        image = Image.open(image_path)62        alpha = image.getchannel(3)63        image = image.convert('RGB')64        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)65        alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)66        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.067        alpha = torch.tensor(np.array(alpha)).float() / 255.068        69        return {70            'image': image,71            'alpha': alpha,72            'extrinsics': extrinsics,73            'intrinsics': intrinsics,74        }75    76    def _get_latent(self, root, instance):77        data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz'))78        coords = torch.tensor(data['coords']).int()79        feats = torch.tensor(data['feats']).float()80        return {81            'coords': coords,82            'feats': feats,83        }84 85    @torch.no_grad()86    def visualize_sample(self, sample: dict):87        return sample['image']88 89    @staticmethod90    def collate_fn(batch):91        pack = {}92        coords = []93        for i, b in enumerate(batch):94            coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))95        coords = torch.cat(coords)96        feats = torch.cat([b['feats'] for b in batch])97        pack['latents'] = SparseTensor(98            coords=coords,99            feats=feats,100        )101        102        # collate other data103        keys = [k for k in batch[0].keys() if k not in ['coords', 'feats']]104        for k in keys:105            if isinstance(batch[0][k], torch.Tensor):106                pack[k] = torch.stack([b[k] for b in batch])107            elif isinstance(batch[0][k], list):108                pack[k] = sum([b[k] for b in batch], [])109            else:110                pack[k] = [b[k] for b in batch]111 112        return pack113 114    def get_instance(self, root, instance):115        image = self._get_image(root, instance)116        latent = self._get_latent(root, instance)117        return {118            **image,119            **latent,120        }121        122 123class Slat2RenderGeo(SLat2Render):124    def __init__(125        self,126        roots: str,127        image_size: int,128        latent_model: str,129        min_aesthetic_score: float = 5.0,130        max_num_voxels: int = 32768,131    ):132        super().__init__(133            roots,134            image_size,135            latent_model,136            min_aesthetic_score,137            max_num_voxels,138        )139        140    def _get_geo(self, root, instance):141        verts, face = utils3d.io.read_ply(os.path.join(root, 'renders', instance, 'mesh.ply'))142        mesh = {143            "vertices" : torch.from_numpy(verts),144            "faces" : torch.from_numpy(face),145        }146        return  {147            "mesh" : mesh,148        }149        150    def get_instance(self, root, instance):151        image = self._get_image(root, instance)152        latent = self._get_latent(root, instance)153        geo = self._get_geo(root, instance)154        return {155            **image,156            **latent,157            **geo,158        }159        160