135 lines
4.7 KiB
Python
Executable File
135 lines
4.7 KiB
Python
Executable File
import os
|
|
from PIL import Image
|
|
import json
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import utils3d.torch
|
|
from ..modules.sparse.basic import SparseTensor
|
|
from .components import StandardDatasetBase
|
|
|
|
|
|
class SparseFeat2Render(StandardDatasetBase):
|
|
"""
|
|
SparseFeat2Render dataset.
|
|
|
|
Args:
|
|
roots (str): paths to the dataset
|
|
image_size (int): size of the image
|
|
model (str): model name
|
|
resolution (int): resolution of the data
|
|
min_aesthetic_score (float): minimum aesthetic score
|
|
max_num_voxels (int): maximum number of voxels
|
|
"""
|
|
def __init__(
|
|
self,
|
|
roots: str,
|
|
image_size: int,
|
|
model: str = 'dinov2_vitl14_reg',
|
|
resolution: int = 64,
|
|
min_aesthetic_score: float = 5.0,
|
|
max_num_voxels: int = 32768,
|
|
):
|
|
self.image_size = image_size
|
|
self.model = model
|
|
self.resolution = resolution
|
|
self.min_aesthetic_score = min_aesthetic_score
|
|
self.max_num_voxels = max_num_voxels
|
|
self.value_range = (0, 1)
|
|
|
|
super().__init__(roots)
|
|
|
|
def filter_metadata(self, metadata):
|
|
stats = {}
|
|
metadata = metadata[metadata[f'feature_{self.model}']]
|
|
stats['With features'] = len(metadata)
|
|
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
|
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
|
metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]
|
|
stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
|
|
return metadata, stats
|
|
|
|
def _get_image(self, root, instance):
|
|
with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:
|
|
metadata = json.load(f)
|
|
n_views = len(metadata['frames'])
|
|
view = np.random.randint(n_views)
|
|
metadata = metadata['frames'][view]
|
|
fov = metadata['camera_angle_x']
|
|
intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
|
|
c2w = torch.tensor(metadata['transform_matrix'])
|
|
c2w[:3, 1:3] *= -1
|
|
extrinsics = torch.inverse(c2w)
|
|
|
|
image_path = os.path.join(root, 'renders', instance, metadata['file_path'])
|
|
image = Image.open(image_path)
|
|
alpha = image.getchannel(3)
|
|
image = image.convert('RGB')
|
|
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
|
alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
|
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
|
|
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
|
|
|
return {
|
|
'image': image,
|
|
'alpha': alpha,
|
|
'extrinsics': extrinsics,
|
|
'intrinsics': intrinsics,
|
|
}
|
|
|
|
def _get_feat(self, root, instance):
|
|
DATA_RESOLUTION = 64
|
|
feats_path = os.path.join(root, 'features', self.model, f'{instance}.npz')
|
|
feats = np.load(feats_path, allow_pickle=True)
|
|
coords = torch.tensor(feats['indices']).int()
|
|
feats = torch.tensor(feats['patchtokens']).float()
|
|
|
|
if self.resolution != DATA_RESOLUTION:
|
|
factor = DATA_RESOLUTION // self.resolution
|
|
coords = coords // factor
|
|
coords, idx = coords.unique(return_inverse=True, dim=0)
|
|
feats = torch.scatter_reduce(
|
|
torch.zeros(coords.shape[0], feats.shape[1], device=feats.device),
|
|
dim=0,
|
|
index=idx.unsqueeze(-1).expand(-1, feats.shape[1]),
|
|
src=feats,
|
|
reduce='mean'
|
|
)
|
|
|
|
return {
|
|
'coords': coords,
|
|
'feats': feats,
|
|
}
|
|
|
|
@torch.no_grad()
|
|
def visualize_sample(self, sample: dict):
|
|
return sample['image']
|
|
|
|
@staticmethod
|
|
def collate_fn(batch):
|
|
pack = {}
|
|
coords = []
|
|
for i, b in enumerate(batch):
|
|
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
|
|
coords = torch.cat(coords)
|
|
feats = torch.cat([b['feats'] for b in batch])
|
|
pack['feats'] = SparseTensor(
|
|
coords=coords,
|
|
feats=feats,
|
|
)
|
|
|
|
pack['image'] = torch.stack([b['image'] for b in batch])
|
|
pack['alpha'] = torch.stack([b['alpha'] for b in batch])
|
|
pack['extrinsics'] = torch.stack([b['extrinsics'] for b in batch])
|
|
pack['intrinsics'] = torch.stack([b['intrinsics'] for b in batch])
|
|
|
|
return pack
|
|
|
|
def get_instance(self, root, instance):
|
|
image = self._get_image(root, instance)
|
|
feat = self._get_feat(root, instance)
|
|
return {
|
|
**image,
|
|
**feat,
|
|
}
|