souging/TRELLIS_TextTo3D

1import os2from PIL import Image3import json4import numpy as np5import pandas as pd6import torch7import utils3d.torch8from ..modules.sparse.basic import SparseTensor9from .components import StandardDatasetBase10 11 12class SparseFeat2Render(StandardDatasetBase):13    """14    SparseFeat2Render dataset.15    16    Args:17        roots (str): paths to the dataset18        image_size (int): size of the image19        model (str): model name20        resolution (int): resolution of the data21        min_aesthetic_score (float): minimum aesthetic score22        max_num_voxels (int): maximum number of voxels23    """24    def __init__(25        self,26        roots: str,27        image_size: int,28        model: str = 'dinov2_vitl14_reg',29        resolution: int = 64,30        min_aesthetic_score: float = 5.0,31        max_num_voxels: int = 32768,32    ):33        self.image_size = image_size34        self.model = model35        self.resolution = resolution36        self.min_aesthetic_score = min_aesthetic_score37        self.max_num_voxels = max_num_voxels38        self.value_range = (0, 1)39        40        super().__init__(roots)41        42    def filter_metadata(self, metadata):43        stats = {}44        metadata = metadata[metadata[f'feature_{self.model}']]45        stats['With features'] = len(metadata)46        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]47        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)48        metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]49        stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)50        return metadata, stats51 52    def _get_image(self, root, instance):53        with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:54            metadata = json.load(f)55        n_views = len(metadata['frames'])56        view = np.random.randint(n_views)57        metadata = metadata['frames'][view]58        fov = metadata['camera_angle_x']59        intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))60        c2w = torch.tensor(metadata['transform_matrix'])61        c2w[:3, 1:3] *= -162        extrinsics = torch.inverse(c2w)63 64        image_path = os.path.join(root, 'renders', instance, metadata['file_path'])65        image = Image.open(image_path)66        alpha = image.getchannel(3)67        image = image.convert('RGB')68        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)69        alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)70        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.071        alpha = torch.tensor(np.array(alpha)).float() / 255.072        73        return {74            'image': image,75            'alpha': alpha,76            'extrinsics': extrinsics,77            'intrinsics': intrinsics,78        }79    80    def _get_feat(self, root, instance):81        DATA_RESOLUTION = 6482        feats_path = os.path.join(root, 'features', self.model, f'{instance}.npz')83        feats = np.load(feats_path, allow_pickle=True)84        coords = torch.tensor(feats['indices']).int()85        feats = torch.tensor(feats['patchtokens']).float()86        87        if self.resolution != DATA_RESOLUTION:88            factor = DATA_RESOLUTION // self.resolution89            coords = coords // factor90            coords, idx = coords.unique(return_inverse=True, dim=0)91            feats = torch.scatter_reduce(92                torch.zeros(coords.shape[0], feats.shape[1], device=feats.device),93                dim=0,94                index=idx.unsqueeze(-1).expand(-1, feats.shape[1]),95                src=feats,96                reduce='mean'97            )98        99        return {100            'coords': coords,101            'feats': feats,102        }103 104    @torch.no_grad()105    def visualize_sample(self, sample: dict):106        return sample['image']107 108    @staticmethod109    def collate_fn(batch):110        pack = {}111        coords = []112        for i, b in enumerate(batch):113            coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))114        coords = torch.cat(coords)115        feats = torch.cat([b['feats'] for b in batch])116        pack['feats'] = SparseTensor(117            coords=coords,118            feats=feats,119        )120        121        pack['image'] = torch.stack([b['image'] for b in batch])122        pack['alpha'] = torch.stack([b['alpha'] for b in batch])123        pack['extrinsics'] = torch.stack([b['extrinsics'] for b in batch])124        pack['intrinsics'] = torch.stack([b['intrinsics'] for b in batch])125 126        return pack127 128    def get_instance(self, root, instance):129        image = self._get_image(root, instance)130        feat = self._get_feat(root, instance)131        return {132            **image,133            **feat,134        }135