203 lines
6.0 KiB
Python
203 lines
6.0 KiB
Python
import os
|
|
import sys
|
|
import torch
|
|
import numpy as np
|
|
|
|
from deodr import Camera, Scene3D
|
|
from deodr.triangulated_mesh import TriMesh
|
|
from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
|
|
from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
from cloth3d.util import proj
|
|
|
|
|
|
class RenderCore:
|
|
def __init__(self, max_depth=0, depth_scale=1):
|
|
self.mesh = None
|
|
self.depthScale = depth_scale
|
|
self.scene = Scene3D()
|
|
self.set_max_depth(max_depth)
|
|
|
|
def set_mesh(self, vertices, faces):
|
|
self.mesh = TriMesh(
|
|
faces, vertices
|
|
) # we do a copy to avoid negative stride not support by Tensorflow
|
|
|
|
object_center = vertices.mean(axis=0)
|
|
object_radius = np.max(np.std(vertices, axis=0))
|
|
self.camera_center = object_center + np.array([-0.5, 0, 5]) * object_radius
|
|
|
|
self.scene.set_mesh(self.mesh)
|
|
|
|
def set_max_depth(self, max_depth):
|
|
self.scene.max_depth = max_depth
|
|
self.scene.set_background_color([max_depth])
|
|
|
|
def set_depth_scale(self, depth_scale):
|
|
self.depthScale = depth_scale
|
|
|
|
def set_image(self, width, height, intrinsic, extrinsic):
|
|
self.width = width
|
|
self.height = height
|
|
|
|
# focal = None
|
|
# if focal is None:
|
|
# focal = 2 * self.width
|
|
|
|
# rot = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
|
|
# trans = -rot.T.dot(self.camera_center)
|
|
# intrinsic = np.array(
|
|
# [[focal, 0, self.width / 2], [0, focal, self.height / 2], [0, 0, 1]]
|
|
# )
|
|
# extrinsic = np.column_stack((rot, trans))
|
|
|
|
self.camera = Camera(
|
|
extrinsic=extrinsic,
|
|
intrinsic=intrinsic,
|
|
width=self.width,
|
|
height=self.height,
|
|
distortion=None,
|
|
)
|
|
|
|
def render(self):
|
|
depth_scale = 1 * self.depthScale
|
|
depth = self.scene.render_depth(
|
|
self.camera,
|
|
depth_scale=depth_scale,
|
|
)
|
|
#depth = tf.clip_by_value(depth, 0, self.scene.max_depth)
|
|
return depth
|
|
|
|
|
|
""" CAMERA """
|
|
def intrinsic():
|
|
RES_X = 640
|
|
RES_Y = 480
|
|
f_mm = 50 # blender default
|
|
sensor_w_mm = 36 # blender default
|
|
sensor_h_mm = sensor_w_mm * RES_Y / RES_X
|
|
|
|
fx_px = f_mm * RES_X / sensor_w_mm
|
|
fy_px = f_mm * RES_Y / sensor_h_mm
|
|
|
|
u = RES_X / 2
|
|
v = RES_Y / 2
|
|
|
|
return np.array([[fx_px, 0, u],
|
|
[0, fy_px, v],
|
|
[0, 0, 1]], np.float32)
|
|
|
|
|
|
def extrinsic(camLoc):
|
|
R_w2bc = np.array([[0, 1, 0],
|
|
[0, 0, 1],
|
|
[1, 0, 0]], np.float32)
|
|
|
|
T_w2bc = -1 * R_w2bc.dot(camLoc)
|
|
|
|
R_bc2cv = np.array([[1, 0, 0],
|
|
[0, -1, 0],
|
|
[0, 0, -1]], np.float32)
|
|
|
|
R_w2cv = R_bc2cv.dot(R_w2bc)
|
|
T_w2cv = R_bc2cv.dot(T_w2bc)
|
|
|
|
return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
|
|
|
|
def proj(camLoc):
|
|
return intrinsic().dot(extrinsic(camLoc))
|
|
|
|
|
|
class Render:
|
|
def __init__(self, max_depth=0, depth_scale=1, size=(640, 480)):
|
|
self.render_core = RenderCore(max_depth=max_depth)
|
|
self.size = size
|
|
self.max_depth = max_depth
|
|
|
|
def render(self, mesh, faces, cam):
|
|
self.render_core.set_mesh(mesh, faces)
|
|
self.render_core.set_image(self.size[0], self.size[1], intrinsic(), extrinsic(cam))
|
|
depth = self.render_core.render()
|
|
|
|
thresh = self.max_depth - 1
|
|
mask = depth<thresh
|
|
# dmin = depth.min()
|
|
dmean = depth[mask].mean()
|
|
# dmax = depth[mask].max()
|
|
depth[mask] = depth[mask] - dmean
|
|
depth[~mask] = 0.0
|
|
mask = mask.astype(np.float32)
|
|
|
|
return depth
|
|
|
|
|
|
from torch import nn
|
|
from pytorch3d.renderer import (
|
|
look_at_view_transform,
|
|
FoVPerspectiveCameras,
|
|
PerspectiveCameras,
|
|
PointLights,
|
|
DirectionalLights,
|
|
Materials,
|
|
RasterizationSettings,
|
|
MeshRenderer,
|
|
MeshRasterizer,
|
|
SoftPhongShader,
|
|
TexturesUV,
|
|
TexturesVertex
|
|
)
|
|
from pytorch3d.structures import Meshes
|
|
from pytorch3d.utils import cameras_from_opencv_projection
|
|
|
|
|
|
class DepthRenderer(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
|
|
|
|
def forward(self, verts, faces, cam):
|
|
# R, T = look_at_view_transform(2.7, 0, 0)
|
|
# cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T)
|
|
|
|
ext = torch.tensor(extrinsic(cam).astype(np.float32), device='cuda')
|
|
# ext = proj(cam)
|
|
# fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32)
|
|
img_size = np.array([480, 640], dtype=np.float32)
|
|
# K = np.eye(4, dtype=np.float32)
|
|
# K[:3,:3] = intrinsic()
|
|
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
|
# cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=1,)# principal_point=fp[None, :], image_size=img_size[None, :])
|
|
|
|
cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(img_size[None, :], device='cuda'))
|
|
|
|
raster_settings = RasterizationSettings(
|
|
image_size=(480, 640),
|
|
blur_radius=0.0,
|
|
faces_per_pixel=1,
|
|
bin_size=0
|
|
)
|
|
|
|
rasterizer = MeshRasterizer(
|
|
cameras=cameras,
|
|
raster_settings=raster_settings
|
|
)
|
|
output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:]))
|
|
v2i = cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
|
|
depth = output.zbuf[0,:,:,0]
|
|
pix_to_face = output.pix_to_face[0,:,:,0]
|
|
|
|
mask = depth != -1
|
|
# dmin = depth.min()
|
|
dmean = depth[mask].mean()
|
|
# dmax = depth[mask].max()
|
|
# depth[mask] = depth[mask] - dmean
|
|
# depth[~mask] = 0.0
|
|
# mask = mask.astype(np.float32)
|
|
|
|
# import plotly.express as ex
|
|
# ex.imshow(depth.cpu().numpy()).show()
|
|
|
|
return depth, mask, v2i, K, cameras, pix_to_face
|