Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,872 Bytes
7f51798 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os, cv2
import matplotlib.pyplot as plt
import math
import point_cloud_utils as pcu
from pdb import set_trace as st
def depths_to_points_2(world_view_transform, tanfov, W, H, depthmap):
c2w = (world_view_transform).inverse()
# W, H = view.image_width, view.image_height
# fx = W / (2 * math.tan(view.FoVx / 2.))
# fy = H / (2 * math.tan(view.FoVy / 2.))
fx = W / (2 * tanfov)
fy = H / (2 * tanfov)
intrins = torch.tensor(
[[fx, 0., W/2.],
[0., fy, H/2.],
[0., 0., 1.0]]
).float().cuda()
grid_x, grid_y = torch.meshgrid(torch.arange(W, device='cuda').float(), torch.arange(H, device='cuda').float(), indexing='xy')
points = torch.stack([grid_x, grid_y, torch.ones_like(grid_x)], dim=-1).reshape(-1, 3)
rays_d = points @ intrins.inverse().T @ c2w[:3,:3].T
rays_o = c2w[:3,3]
points = depthmap.reshape(-1, 1) * rays_d + rays_o
return points
def depths_to_points(view, depthmap):
c2w = (view.world_view_transform.T).inverse()
W, H = view.image_width, view.image_height
fx = W / (2 * math.tan(view.FoVx / 2.))
fy = H / (2 * math.tan(view.FoVy / 2.))
intrins = torch.tensor(
[[fx, 0., W/2.],
[0., fy, H/2.],
[0., 0., 1.0]]
).float().cuda()
grid_x, grid_y = torch.meshgrid(torch.arange(W, device='cuda').float(), torch.arange(H, device='cuda').float(), indexing='xy')
points = torch.stack([grid_x, grid_y, torch.ones_like(grid_x)], dim=-1).reshape(-1, 3)
rays_d = points @ intrins.inverse().T @ c2w[:3,:3].T
rays_o = c2w[:3,3]
points = depthmap.reshape(-1, 1) * rays_d + rays_o
return points
def depth_to_normal(view, depth):
"""
view: view camera
depth: depthmap
"""
points = depths_to_points(view, depth).reshape(*depth.shape[1:], 3)
output = torch.zeros_like(points)
dx = torch.cat([points[2:, 1:-1] - points[:-2, 1:-1]], dim=0)
dy = torch.cat([points[1:-1, 2:] - points[1:-1, :-2]], dim=1)
normal_map = torch.nn.functional.normalize(torch.cross(dx, dy, dim=-1), dim=-1)
output[1:-1, 1:-1, :] = normal_map
return output
def depth_to_normal_2(world_view_transform, tanfov, W, H, depth):
"""
view: view camera
depth: depthmap
"""
points = depths_to_points_2(world_view_transform, tanfov, W, H, depth).reshape(*depth.shape[1:], 3)
# st()
# pcu.save_mesh_v( f'tmp/depth2pts.ply', points.detach().reshape(-1,3).cpu().numpy(),)
output = torch.zeros_like(points)
dx = torch.cat([points[2:, 1:-1] - points[:-2, 1:-1]], dim=0)
dy = torch.cat([points[1:-1, 2:] - points[1:-1, :-2]], dim=1)
normal_map = torch.nn.functional.normalize(torch.cross(dx, dy, dim=-1), dim=-1)
output[1:-1, 1:-1, :] = normal_map
return output |