Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,259 Bytes
178f950 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import os
import json
from typing import Union
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
import utils3d
from .components import StandardDatasetBase
from ..representations.octree import DfsOctree as Octree
from ..renderers import OctreeRenderer
class SparseStructure(StandardDatasetBase):
"""
Sparse structure 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 = 64,
min_aesthetic_score: float = 5.0,
):
self.resolution = resolution
self.min_aesthetic_score = min_aesthetic_score
self.value_range = (0, 1)
super().__init__(roots)
def filter_metadata(self, metadata):
stats = {}
metadata = metadata[metadata[f'voxelized']]
stats['Voxelized'] = len(metadata)
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
return metadata, stats
def get_instance(self, root, instance):
position = utils3d.io.read_ply(os.path.join(root, 'voxels', f'{instance}.ply'))[0]
coords = ((torch.tensor(position) + 0.5) * self.resolution).int().contiguous()
ss = torch.zeros(1, self.resolution, self.resolution, self.resolution, dtype=torch.long)
ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
return {'ss': ss}
@torch.no_grad()
def visualize_sample(self, ss: Union[torch.Tensor, dict]):
ss = ss if isinstance(ss, torch.Tensor) else ss['ss']
renderer = OctreeRenderer()
renderer.rendering_options.resolution = 512
renderer.rendering_options.near = 0.8
renderer.rendering_options.far = 1.6
renderer.rendering_options.bg_color = (0, 0, 0)
renderer.rendering_options.ssaa = 4
renderer.pipe.primitive = 'voxel'
# 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 = []
# Build each representation
ss = ss.cuda()
for i in range(ss.shape[0]):
representation = Octree(
depth=10,
aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
device='cuda',
primitive='voxel',
sh_degree=0,
primitive_config={'solid': True},
)
coords = torch.nonzero(ss[i, 0], as_tuple=False)
representation.position = coords.float() / self.resolution
representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')
image = torch.zeros(3, 1024, 1024).cuda()
tile = [2, 2]
for j, (ext, intr) in enumerate(zip(exts, ints)):
res = renderer.render(representation, ext, intr, colors_overwrite=representation.position)
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
images.append(image)
return torch.stack(images)
|