Spaces:
Sleeping
Sleeping
File size: 1,130 Bytes
d7e58f0 |
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 |
import torch
def vis_z_buffer(z, percentile=1, vis_pad=0.2):
z = z[:, :, 0]
mask = z > 1e-5
if torch.sum(mask) == 0:
z[...] = 0
else:
vmin = torch.quantile(z[mask], percentile / 100)
vmax = torch.quantile(z[mask], 1 - percentile / 100)
pad = (vmax - vmin) * vis_pad
vmin_padded = vmin - pad
vmax_padded = vmax + pad
z[mask] = vmin + vmax - z[mask]
z = (z - vmin_padded) / (vmax_padded - vmin_padded)
z = torch.clip(torch.round(z * 255), 0, 255)
z_cpu = z.to(dtype=torch.uint8).detach().cpu().numpy()
return z_cpu
def vis_normals(coords, normals, vis_pad=0.2):
mask = coords[:, :, 2] > 0
coords_masked = -coords[mask]
normals_masked = normals[mask]
coords_len = torch.sqrt(torch.sum(coords_masked**2, dim=1))
dot = torch.sum(coords_masked * normals_masked, dim=1) / coords_len
h, w = normals.shape[:2]
vis = torch.zeros((h, w), dtype=coords.dtype, device=coords.device)
vis[mask] = torch.clamp(dot, 0, 1) * (1 - 2 * vis_pad) + vis_pad
vis = (vis * 255).to(dtype=torch.uint8)
return vis
|