souging/TRELLIS_TextTo3D

1import os2import json3from typing import Union4import numpy as np5import pandas as pd6import torch7from torch.utils.data import Dataset8import utils3d9from .components import StandardDatasetBase10from ..representations.octree import DfsOctree as Octree11from ..renderers import OctreeRenderer12 13 14class SparseStructure(StandardDatasetBase):15    """16    Sparse structure dataset17 18    Args:19        roots (str): path to the dataset20        resolution (int): resolution of the voxel grid21        min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset22    """23 24    def __init__(self,25        roots,26        resolution: int = 64,27        min_aesthetic_score: float = 5.0,28    ):29        self.resolution = resolution30        self.min_aesthetic_score = min_aesthetic_score31        self.value_range = (0, 1)32 33        super().__init__(roots)34        35    def filter_metadata(self, metadata):36        stats = {}37        metadata = metadata[metadata[f'voxelized']]38        stats['Voxelized'] = len(metadata)39        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]40        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)41        return metadata, stats42 43    def get_instance(self, root, instance):44        position = utils3d.io.read_ply(os.path.join(root, 'voxels', f'{instance}.ply'))[0]45        coords = ((torch.tensor(position) + 0.5) * self.resolution).int().contiguous()46        ss = torch.zeros(1, self.resolution, self.resolution, self.resolution, dtype=torch.long)47        ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 148        return {'ss': ss}49 50    @torch.no_grad()51    def visualize_sample(self, ss: Union[torch.Tensor, dict]):52        ss = ss if isinstance(ss, torch.Tensor) else ss['ss']53        54        renderer = OctreeRenderer()55        renderer.rendering_options.resolution = 51256        renderer.rendering_options.near = 0.857        renderer.rendering_options.far = 1.658        renderer.rendering_options.bg_color = (0, 0, 0)59        renderer.rendering_options.ssaa = 460        renderer.pipe.primitive = 'voxel'61        62        # Build camera63        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]64        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)65        yaws = [y + yaws_offset for y in yaws]66        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]67 68        exts = []69        ints = []70        for yaw, pitch in zip(yaws, pitch):71            orig = torch.tensor([72                np.sin(yaw) * np.cos(pitch),73                np.cos(yaw) * np.cos(pitch),74                np.sin(pitch),75            ]).float().cuda() * 276            fov = torch.deg2rad(torch.tensor(30)).cuda()77            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())78            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)79            exts.append(extrinsics)80            ints.append(intrinsics)81 82        images = []83        84        # Build each representation85        ss = ss.cuda()86        for i in range(ss.shape[0]):87            representation = Octree(88                depth=10,89                aabb=[-0.5, -0.5, -0.5, 1, 1, 1],90                device='cuda',91                primitive='voxel',92                sh_degree=0,93                primitive_config={'solid': True},94            )95            coords = torch.nonzero(ss[i, 0], as_tuple=False)96            representation.position = coords.float() / self.resolution97            representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')98 99            image = torch.zeros(3, 1024, 1024).cuda()100            tile = [2, 2]101            for j, (ext, intr) in enumerate(zip(exts, ints)):102                res = renderer.render(representation, ext, intr, colors_overwrite=representation.position)103                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']104            images.append(image)105            106        return torch.stack(images)107