wehub-resource-sync
299 lines
12 KiB
Python
299 lines
12 KiB
Python
import os
|
|
import io
|
|
from typing import Union
|
|
import numpy as np
|
|
import pickle
|
|
import torch
|
|
from PIL import Image
|
|
import o_voxel
|
|
import utils3d
|
|
from .components import StandardDatasetBase
|
|
from ..modules import sparse as sp
|
|
from ..renderers import VoxelRenderer
|
|
from ..representations import Voxel
|
|
from ..representations.mesh import MeshWithPbrMaterial, TextureFilterMode, TextureWrapMode, AlphaMode, PbrMaterial, Texture
|
|
|
|
from ..utils.data_utils import load_balanced_group_indices
|
|
|
|
|
|
def is_power_of_two(n: int) -> bool:
|
|
return n > 0 and (n & (n - 1)) == 0
|
|
|
|
|
|
def nearest_power_of_two(n: int) -> int:
|
|
if n < 1:
|
|
raise ValueError("n must be >= 1")
|
|
if is_power_of_two(n):
|
|
return n
|
|
lower = 2 ** (n.bit_length() - 1)
|
|
upper = 2 ** n.bit_length()
|
|
if n - lower < upper - n:
|
|
return lower
|
|
else:
|
|
return upper
|
|
|
|
|
|
class SparseVoxelPbrVisMixin:
|
|
@torch.no_grad()
|
|
def visualize_sample(self, x: Union[sp.SparseTensor, dict]):
|
|
x = x if isinstance(x, sp.SparseTensor) else x['x']
|
|
|
|
renderer = VoxelRenderer()
|
|
renderer.rendering_options.resolution = 512
|
|
renderer.rendering_options.ssaa = 4
|
|
|
|
# Build camera
|
|
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
|
|
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
|
|
yaws = [y + yaws_offset for y in yaws]
|
|
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
|
|
|
|
exts = []
|
|
ints = []
|
|
for yaw, pitch in zip(yaws, pitch):
|
|
orig = torch.tensor([
|
|
np.sin(yaw) * np.cos(pitch),
|
|
np.cos(yaw) * np.cos(pitch),
|
|
np.sin(pitch),
|
|
]).float().cuda() * 2
|
|
fov = torch.deg2rad(torch.tensor(30)).cuda()
|
|
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
|
|
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
|
|
exts.append(extrinsics)
|
|
ints.append(intrinsics)
|
|
|
|
images = {k: [] for k in self.layout}
|
|
|
|
# Build each representation
|
|
x = x.cuda()
|
|
for i in range(x.shape[0]):
|
|
rep = Voxel(
|
|
origin=[-0.5, -0.5, -0.5],
|
|
voxel_size=1/self.resolution,
|
|
coords=x[i].coords[:, 1:].contiguous(),
|
|
attrs=None,
|
|
layout={
|
|
'color': slice(0, 3),
|
|
}
|
|
)
|
|
for k in self.layout:
|
|
image = torch.zeros(3, 1024, 1024).cuda()
|
|
tile = [2, 2]
|
|
for j, (ext, intr) in enumerate(zip(exts, ints)):
|
|
attr = x[i].feats[:, self.layout[k]].expand(-1, 3)
|
|
res = renderer.render(rep, ext, intr, colors_overwrite=attr)
|
|
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
|
|
images[k].append(image)
|
|
|
|
for k in self.layout:
|
|
images[k] = torch.stack(images[k])
|
|
|
|
return images
|
|
|
|
|
|
class SparseVoxelPbrDataset(SparseVoxelPbrVisMixin, StandardDatasetBase):
|
|
"""
|
|
Sparse Voxel PBR dataset.
|
|
|
|
Args:
|
|
roots (str): path to the dataset
|
|
resolution (int): resolution of the voxel grid
|
|
min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
roots,
|
|
resolution: int = 1024,
|
|
max_active_voxels: int = 1000000,
|
|
max_num_faces: int = None,
|
|
min_aesthetic_score: float = 5.0,
|
|
attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
|
|
with_mesh: bool = True,
|
|
):
|
|
self.resolution = resolution
|
|
self.min_aesthetic_score = min_aesthetic_score
|
|
self.max_active_voxels = max_active_voxels
|
|
self.max_num_faces = max_num_faces
|
|
self.with_mesh = with_mesh
|
|
self.value_range = (-1, 1)
|
|
self.channels = {
|
|
'base_color': 3,
|
|
'metallic': 1,
|
|
'roughness': 1,
|
|
'emissive': 3,
|
|
'alpha': 1,
|
|
}
|
|
self.layout = {}
|
|
start = 0
|
|
for attr in attrs:
|
|
self.layout[attr] = slice(start, start + self.channels[attr])
|
|
start += self.channels[attr]
|
|
|
|
super().__init__(roots)
|
|
|
|
self.loads = [self.metadata.loc[sha256, f'num_pbr_voxels'] for _, sha256, _ in self.instances]
|
|
|
|
def __str__(self):
|
|
lines = [
|
|
super().__str__(),
|
|
f' - Resolution: {self.resolution}',
|
|
f' - Attributes: {list(self.layout.keys())}',
|
|
]
|
|
return '\n'.join(lines)
|
|
|
|
def filter_metadata(self, metadata, dataset_name=None):
|
|
stats = {}
|
|
metadata = metadata[metadata['pbr_voxelized'] == True]
|
|
stats['PBR Voxelized'] = len(metadata)
|
|
if self.min_aesthetic_score is not None:
|
|
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
|
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
|
metadata = metadata[metadata['num_pbr_voxels'] <= self.max_active_voxels]
|
|
stats[f'Active voxels <= {self.max_active_voxels}'] = len(metadata)
|
|
if self.max_num_faces is not None:
|
|
metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
|
|
stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
|
|
return metadata, stats
|
|
|
|
@staticmethod
|
|
def _texture_from_dump(pack) -> Texture:
|
|
png_bytes = pack['image']
|
|
image = Image.open(io.BytesIO(png_bytes))
|
|
if image.width != image.height or not is_power_of_two(image.width):
|
|
size = nearest_power_of_two(max(image.width, image.height))
|
|
image = image.resize((size, size), Image.LANCZOS)
|
|
texture = torch.tensor(np.array(image) / 255.0, dtype=torch.float32).reshape(image.height, image.width, -1)
|
|
filter_mode = {
|
|
'Linear': TextureFilterMode.LINEAR,
|
|
'Closest': TextureFilterMode.CLOSEST,
|
|
'Cubic': TextureFilterMode.LINEAR,
|
|
'Smart': TextureFilterMode.LINEAR,
|
|
}[pack['interpolation']]
|
|
wrap_mode = {
|
|
'REPEAT': TextureWrapMode.REPEAT,
|
|
'EXTEND': TextureWrapMode.CLAMP_TO_EDGE,
|
|
'CLIP': TextureWrapMode.CLAMP_TO_EDGE,
|
|
'MIRROR': TextureWrapMode.MIRRORED_REPEAT,
|
|
}[pack['extension']]
|
|
return Texture(texture, filter_mode=filter_mode, wrap_mode=wrap_mode)
|
|
|
|
def read_mesh_with_texture(self, root, instance):
|
|
with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
|
|
dump = pickle.load(f)
|
|
|
|
# Fix dump alpha map
|
|
for mat in dump['materials']:
|
|
if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
|
|
mat['alphaMode'] = 'BLEND'
|
|
|
|
# process material
|
|
materials = []
|
|
for mat in dump['materials']:
|
|
materials.append(PbrMaterial(
|
|
base_color_texture=self._texture_from_dump(mat['baseColorTexture']) if mat['baseColorTexture'] is not None else None,
|
|
base_color_factor=mat['baseColorFactor'],
|
|
metallic_texture=self._texture_from_dump(mat['metallicTexture']) if mat['metallicTexture'] is not None else None,
|
|
metallic_factor=mat['metallicFactor'],
|
|
roughness_texture=self._texture_from_dump(mat['roughnessTexture']) if mat['roughnessTexture'] is not None else None,
|
|
roughness_factor=mat['roughnessFactor'],
|
|
alpha_texture=self._texture_from_dump(mat['alphaTexture']) if mat['alphaTexture'] is not None else None,
|
|
alpha_factor=mat['alphaFactor'],
|
|
alpha_mode={
|
|
'OPAQUE': AlphaMode.OPAQUE,
|
|
'MASK': AlphaMode.MASK,
|
|
'BLEND': AlphaMode.BLEND,
|
|
}[mat['alphaMode']],
|
|
alpha_cutoff=mat['alphaCutoff'],
|
|
))
|
|
materials.append(PbrMaterial(
|
|
base_color_factor=[0.8, 0.8, 0.8],
|
|
alpha_factor=1.0,
|
|
metallic_factor=0.0,
|
|
roughness_factor=0.5,
|
|
alpha_mode=AlphaMode.OPAQUE,
|
|
alpha_cutoff=0.5,
|
|
)) # append default material
|
|
|
|
# process mesh
|
|
start = 0
|
|
vertices = []
|
|
faces = []
|
|
material_ids = []
|
|
uv_coords = []
|
|
for obj in dump['objects']:
|
|
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
|
continue
|
|
vertices.append(obj['vertices'])
|
|
faces.append(obj['faces'] + start)
|
|
obj['mat_ids'][obj['mat_ids'] == -1] = len(materials) - 1
|
|
material_ids.append(obj['mat_ids'])
|
|
uv_coords.append(obj['uvs'] if obj['uvs'] is not None else np.zeros((obj['faces'].shape[0], 3, 2), dtype=np.float32))
|
|
start += len(obj['vertices'])
|
|
|
|
vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
|
|
faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
|
|
material_ids = torch.from_numpy(np.concatenate(material_ids, axis=0)).long()
|
|
uv_coords = torch.from_numpy(np.concatenate(uv_coords, axis=0)).float()
|
|
|
|
# Normalize vertices
|
|
vertices_min = vertices.min(dim=0)[0]
|
|
vertices_max = vertices.max(dim=0)[0]
|
|
center = (vertices_min + vertices_max) / 2
|
|
scale = 0.99999 / (vertices_max - vertices_min).max()
|
|
vertices = (vertices - center) * scale
|
|
assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
|
|
|
|
return {'mesh': [MeshWithPbrMaterial(
|
|
vertices=vertices,
|
|
faces=faces,
|
|
material_ids=material_ids,
|
|
uv_coords=uv_coords,
|
|
materials=materials,
|
|
)]}
|
|
|
|
def read_pbr_voxel(self, root, instance):
|
|
coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
|
|
feats = torch.concat([attr[k] for k in self.layout], dim=-1) / 255.0 * 2 - 1
|
|
x = sp.SparseTensor(
|
|
feats.float(),
|
|
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
|
)
|
|
return {'x': x}
|
|
|
|
def get_instance(self, root, instance):
|
|
if self.with_mesh:
|
|
mesh = self.read_mesh_with_texture(root['pbr_dump'], instance)
|
|
pbr_voxel = self.read_pbr_voxel(root['pbr_voxel'], instance)
|
|
return {**mesh, **pbr_voxel}
|
|
else:
|
|
return self.read_pbr_voxel(root['pbr_voxel'], instance)
|
|
|
|
@staticmethod
|
|
def collate_fn(batch, split_size=None):
|
|
if split_size is None:
|
|
group_idx = [list(range(len(batch)))]
|
|
else:
|
|
group_idx = load_balanced_group_indices([b['x'].feats.shape[0] for b in batch], split_size)
|
|
packs = []
|
|
for group in group_idx:
|
|
sub_batch = [batch[i] for i in group]
|
|
pack = {}
|
|
|
|
keys = [k for k in sub_batch[0].keys()]
|
|
for k in keys:
|
|
if isinstance(sub_batch[0][k], torch.Tensor):
|
|
pack[k] = torch.stack([b[k] for b in sub_batch])
|
|
elif isinstance(sub_batch[0][k], sp.SparseTensor):
|
|
pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
|
|
elif isinstance(sub_batch[0][k], list):
|
|
pack[k] = sum([b[k] for b in sub_batch], [])
|
|
else:
|
|
pack[k] = [b[k] for b in sub_batch]
|
|
|
|
packs.append(pack)
|
|
|
|
if split_size is None:
|
|
return packs[0]
|
|
return packs
|