| """ctypes binding of the C++ mesh renderer (mesh_render.cpp; build with `make` in this folder). |
| |
| A textured triangle mesh is ray-traced through a thin lens with radial distortion, a disk light, |
| soft shadows, a background plane and an optional rectangular occluder. Outputs: rgb, a sheet mask |
| and uv (texture pixel seen by every output pixel). |
| """ |
|
|
| import ctypes as C |
| import sys |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
| LIB = Path(__file__).resolve().parent / ("mesh_render.dll" if sys.platform.startswith("win") else "mesh_render.so") |
|
|
|
|
| @dataclass |
| class Optics: |
| F: float = 50.0 |
| L: float = 66.7 |
| R: float = 10.0 |
|
|
|
|
| @dataclass |
| class Distortion: |
| use: bool = False |
| k1: float = 0.0 |
| cx: float = 0.0 |
| cy: float = 0.0 |
|
|
|
|
| @dataclass |
| class Lighting: |
| use: bool = False |
| x: float = 0.0 |
| y: float = -100.0 |
| z: float = -100.0 |
| intensity: float = 1.0 |
| diameter: float = 20.0 |
| light_mix_koef: float = 0.8 |
|
|
|
|
| @dataclass |
| class BackgroundShadow: |
| use: bool = False |
| bg_z: float = 0.0 |
| bottom_shadow_koef: float = 0.0 |
|
|
|
|
| @dataclass |
| class RectOccluder: |
| use: bool = False |
| cx: float = 0.0 |
| cy: float = 0.0 |
| cz: float = 30.0 |
| w: float = 20.0 |
| h: float = 20.0 |
| yaw: float = 0.0 |
| pitch: float = 0.0 |
| roll: float = 0.0 |
| circle_segments: int = 32 |
|
|
|
|
| class Vector3d(C.Structure): |
| _fields_ = [("x", C.c_float), ("y", C.c_float), ("z", C.c_float)] |
|
|
|
|
| class Vector2d(C.Structure): |
| _fields_ = [("x", C.c_float), ("y", C.c_float)] |
|
|
|
|
| class Triangle(C.Structure): |
| _fields_ = [("v0", C.POINTER(Vector3d)), ("v1", C.POINTER(Vector3d)), ("v2", C.POINTER(Vector3d)), |
| ("imagePoint0", C.POINTER(Vector2d)), ("imagePoint1", C.POINTER(Vector2d)), |
| ("imagePoint2", C.POINTER(Vector2d))] |
|
|
|
|
| class Mesh(C.Structure): |
| _fields_ = [ |
| ("triangles", C.POINTER(Triangle)), ("points3d", C.POINTER(Vector3d)), ("points2d", C.POINTER(Vector2d)), |
| ("triangleCount", C.c_int), ("pointCount", C.c_int), |
| ("use_light_info", C.c_bool), ("light_x", C.c_float), ("light_y", C.c_float), ("light_z", C.c_float), |
| ("light_intensivity", C.c_float), ("use_shadow_info", C.c_bool), ("light_diameter", C.c_float), |
| ("use_bg_shadow", C.c_bool), ("bg_z", C.c_float), |
| ("occ_cx", C.c_float), ("occ_cy", C.c_float), ("occ_cz", C.c_float), ("occ_w", C.c_float), |
| ("occ_h", C.c_float), ("occ_yaw", C.c_float), ("occ_pitch", C.c_float), ("occ_roll", C.c_float), |
| ("occ_circle_segments", C.c_int), |
| ("camera_tilt_x_rad", C.c_float), ("use_distortion_k1", C.c_bool), ("k1", C.c_float), |
| ("dist_norm", C.c_float), ("cx", C.c_float), ("cy", C.c_float), |
| ("bottom_shadow_koef", C.c_float), ("light_mix_koef", C.c_float), |
| ] |
|
|
|
|
| def _load(): |
| dll = C.CDLL(str(LIB)) |
| dll.create_mesh.argtypes = [C.c_int, C.c_int, C.c_int, C.c_int] |
| dll.create_mesh.restype = C.POINTER(Mesh) |
| dll.render.argtypes = [C.c_void_p, C.c_int, C.c_int, C.c_int, C.POINTER(Mesh), C.c_void_p, C.c_int, C.c_int, |
| C.c_int, C.c_bool, C.c_void_p, C.c_int, C.c_bool, C.c_void_p, |
| C.c_float, C.c_float, C.c_float] |
| dll.render.restype = C.c_int |
| dll.delete_mesh.argtypes = [C.POINTER(Mesh)] |
| dll.delete_mesh.restype = None |
| dll.mesh_set_rect_occluder.argtypes = [C.POINTER(Mesh)] + [C.c_float] * 8 + [C.c_int] |
| dll.mesh_set_rect_occluder.restype = None |
| return dll |
|
|
|
|
| class MeshRenderer: |
| """A (grid_w + 1) x (grid_h + 1) node mesh over an input_width x input_height texture. |
| points2d / points3d: numpy views of the node texture coordinates and 3-D positions (camera frame).""" |
|
|
| def __init__(self, input_width, input_height, grid_w, grid_h): |
| self.dll = _load() |
| self.mesh = self.dll.create_mesh(int(input_width), int(input_height), int(grid_w), int(grid_h)) |
| if not self.mesh: |
| raise RuntimeError("create_mesh returned NULL") |
| m = self.mesh.contents |
| self.points2d = np.ctypeslib.as_array(C.cast(m.points2d, C.POINTER(C.c_float)), shape=(m.pointCount, 2)) |
| self.points3d = np.ctypeslib.as_array(C.cast(m.points3d, C.POINTER(C.c_float)), shape=(m.pointCount, 3)) |
|
|
| def configure(self, optics, distortion, lighting, bg_shadow, occluder): |
| self.optics = optics |
| m = self.mesh.contents |
| m.use_distortion_k1 = bool(distortion.use) |
| m.k1 = float(distortion.k1) |
| m.dist_norm = float(optics.L) |
| m.cx, m.cy = float(distortion.cx), float(distortion.cy) |
| m.use_light_info = bool(lighting.use) |
| m.light_x, m.light_y, m.light_z = float(lighting.x), float(lighting.y), float(lighting.z) |
| m.light_intensivity = float(lighting.intensity) |
| m.light_diameter = float(lighting.diameter) |
| m.light_mix_koef = float(lighting.light_mix_koef) |
| m.use_bg_shadow = bool(bg_shadow.use) |
| m.bg_z = float(bg_shadow.bg_z) |
| m.bottom_shadow_koef = float(bg_shadow.bottom_shadow_koef) |
| m.use_shadow_info = bool(bg_shadow.use) |
| m.camera_tilt_x_rad = 0.0 |
| o = occluder |
| self.dll.mesh_set_rect_occluder(self.mesh, float(o.cx), float(o.cy), float(o.cz), float(o.w), float(o.h), |
| float(o.yaw), float(o.pitch), float(o.roll), int(o.circle_segments)) |
|
|
| def render(self, texture, out_size, background): |
| """texture uint8 (H, W, 3), background BGR uint8 of out_size (W, H). |
| -> rgb uint8 (H, W, 3), mask uint8 (H, W) 0/255, uv float32 (H, W, 2) in texture px.""" |
| w, h = int(out_size[0]), int(out_size[1]) |
| rgba = np.zeros((h, w * 4), np.uint8) |
| rgba.reshape(h, w, 4)[...] = cv2.cvtColor(background, cv2.COLOR_BGR2RGBA) |
| mask = np.zeros((h, w), np.uint8) |
| uv = np.zeros((h, w, 2), np.float32) |
| ret = self.dll.render(texture.ctypes.data_as(C.c_void_p), int(texture.shape[1]), int(texture.strides[0]), |
| int(texture.shape[0]), self.mesh, rgba.ctypes.data_as(C.c_void_p), w, w * 4, h, |
| True, mask.ctypes.data_as(C.c_void_p), int(mask.strides[0] // mask.itemsize), |
| True, uv.ctypes.data_as(C.c_void_p), |
| C.c_float(self.optics.R), C.c_float(self.optics.L), C.c_float(self.optics.F)) |
| if ret != 0: |
| raise RuntimeError("render() failed with code {}".format(ret)) |
| rgb = cv2.flip(np.ascontiguousarray(rgba.reshape(h, w, 4)[..., :3]), -1) |
| return rgb, mask[::-1, ::-1], cv2.flip(uv, -1) |
|
|
| def close(self): |
| if self.mesh: |
| self.dll.delete_mesh(self.mesh) |
| self.mesh = None |
|
|
