from pathlib import Path import shutil import os import numpy as np from py_diff_stokes_flow.common.common import ndarray, create_folder from py_diff_stokes_flow.common.project_path import root_path # This class assumes z is pointing up. class PbrtRenderer(object): def __init__(self, options=None): self.__temporary_folder = Path('.tmp') create_folder(self.__temporary_folder) if options is None: options = {} # Image metadata. file_name = options['file_name'] if 'file_name' in options else 'output.exr' file_name = str(file_name) assert file_name.endswith('.png') or file_name.endswith('.exr') file_name_only = file_name[:-4] self.__file_name_only = file_name_only resolution = options['resolution'] if 'resolution' in options else (800, 800) resolution = tuple(resolution) assert len(resolution) == 2 resolution = [int(r) for r in resolution] self.__resolution = tuple(resolution) sample = options['sample'] if 'sample' in options else 4 sample = int(sample) assert sample > 0 self.__sample = sample max_depth = options['max_depth'] if 'max_depth' in options else 4 max_depth = int(max_depth) assert max_depth > 0 self.__max_depth = max_depth # Camera metadata. camera_pos = options['camera_pos'] if 'camera_pos' in options else (2, -2.2, 2) camera_pos = ndarray(camera_pos).ravel() assert camera_pos.size == 3 self.__camera_pos = camera_pos camera_lookat = options['camera_lookat'] if 'camera_lookat' in options else (0.5, 0.5, 0.5) camera_lookat = ndarray(camera_lookat).ravel() assert camera_lookat.size == 3 self.__camera_lookat = camera_lookat camera_up = options['camera_up'] if 'camera_up' in options else (0, 0, 1) camera_up = ndarray(camera_up).ravel() assert camera_up.size == 3 self.__camera_up = camera_up fov = options['fov'] if 'fov' in options else 33 fov = float(fov) assert 0 < fov < 90 self.__fov = fov # Lighting. lightmap = options['light_map'] if 'light_map' in options else 'lightmap.exr' lightmap = Path(root_path) / 'asset/texture/{}'.format(lightmap) self.__lightmap = lightmap lightmap_scale = options['light_map_scale'] if 'light_map_scale' in options else 1.0 lightmap_scale = float(lightmap_scale) self.__lightmap_scale = lightmap_scale # A list of objects. self.__tri_objects = [] self.__shape_objects = [] # - tri_mesh_file: an obj file name. # # - transforms is a list of rotation, translation, and scaling applied to the mesh applied in the order of # their occurances in transforms. # transforms = [rotation, translation, scaling, ...] # rotation = ('r', (radians, unit_axis.x, unit_axis.y, unit_axis.z)) # translation = ('t', (tx, ty, tz)) # scaling = ('s', s) # Note that we use right-handed coordinate systems in the project but pbrt uses a left-handed system. # As a result, we will take care of transforming the coordinate system in this function. # # - color: a 3D vector between 0 and 1 or a string of 6 letters in hex. If render_voxel_edge is False, # we will use a simple material. Otherwise we will generate a texture material from color and texture_img. # # - texture_img: either a texture image name (assumed to be in asset/texture) or 'chkbd_[]_{}' where an integer in [] # indicates the number of grids in the checkerboard and a floating point number between 0 and 1 in {} specifies the # darker color in the checkerboard. def add_tri_mesh(self, tri_mesh_file, transforms=None, color=(.5, .5, .5), texture_img=None): tri_num = len(self.__tri_objects) tri_pbrt_short_name = 'tri_{:08d}.pbrt'.format(tri_num) tri_pbrt_name = self.__temporary_folder / tri_pbrt_short_name tmp_error_name = self.__temporary_folder / '.tmp.error' os.system('{} {} {} 2>{}'.format(str(Path(root_path) / 'external/pbrt_build/obj2pbrt'), tri_mesh_file, tri_pbrt_name, tmp_error_name)) lines = ['AttributeBegin\n',] # Material. if isinstance(color, str): assert len(color) == 6 r = int(color[:2], 16) / 255.0 g = int(color[2:4], 16) / 255.0 b = int(color[4:], 16) / 255.0 color = (r, g, b) color = ndarray(color).ravel() assert color.size == 3 for c in color: assert 0 <= c <= 1 r, g, b = color if texture_img is None: lines.append('Material "plastic" "color Kd" [{} {} {}] "color Ks" [{} {} {}] "float roughness" .3\n'.format( r, g, b, r, g, b)) else: if 'chkbd' in texture_img: _, square_num, square_color = texture_img.split('_') square_num = int(square_num) square_color = np.clip(float(square_color), 0, 1) lines.append('Texture "checks" "spectrum" "checkerboard"\n') lines.append(' "float uscale" [{:d}] "float vscale" [{:d}]\n'.format(square_num, square_num)) lines.append(' "rgb tex1" [{:f} {:f} {:f}] "rgb tex2" [{:f} {:f} {:f}]\n'.format( r, g, b, square_color * r, square_color * g, square_color * b )) lines.append('Material "matte" "texture Kd" "checks"\n') else: texture_img = Path(root_path) / 'asset/texture/{}'.format(texture_img) lines.append('Texture "grid" "color" "imagemap" "string filename" ["{}"]\n'.format(str(texture_img))) lines.append('Texture "sgrid" "color" "scale" "texture tex1" "grid" "color tex2" [{} {} {}]\n'.format(r, g, b)) lines.append('Material "matte" "texture Kd" "sgrid"\n') # Transforms. # Flipped y because pbrt uses a left-handed system. lines.append('Scale 1 -1 1\n') if transforms is not None: for key, vals in reversed(transforms): if key == 's': lines.append('Scale {:f} {:f} {:f}\n'.format(vals, vals, vals)) elif key == 'r': deg = np.rad2deg(vals[0]) ax = vals[1:4] ax /= np.linalg.norm(ax) lines.append('Rotate {:f} {:f} {:f} {:f}\n'.format(deg, ax[0], ax[1], ax[2])) elif key == 't': lines.append('Translate {:f} {:f} {:f}\n'.format(vals[0], vals[1], vals[2])) # Original shape. with open(tri_pbrt_name, 'r') as f: lines += f.readlines() lines.append('AttributeEnd\n') # Write back script. with open(tri_pbrt_name, 'w') as f: for l in lines: f.write(l) self.__tri_objects.append(tri_pbrt_short_name) # - shape_info: a dictionary. def add_shape_mesh(self, shape_info, transforms=None, color=(.5, .5, .5)): shape_num = len(self.__shape_objects) shape_pbrt_short_name = 'shape_{:08d}.pbrt'.format(shape_num) shape_pbrt_name = self.__temporary_folder / shape_pbrt_short_name lines = ['AttributeBegin\n',] # Material. if isinstance(color, str): assert len(color) == 6 r = int(color[:2], 16) / 255.0 g = int(color[2:4], 16) / 255.0 b = int(color[4:], 16) / 255.0 color = (r, g, b) color = ndarray(color).ravel() assert color.size == 3 for c in color: assert 0 <= c <= 1 r, g, b = color lines.append('Material "plastic" "color Kd" [{} {} {}] "color Ks" [{} {} {}] "float roughness" .3\n'.format( r, g, b, r, g, b)) # Transforms. # Flipped y because pbrt uses a left-handed system. lines.append('Scale 1 -1 1\n') if transforms is not None: for key, vals in reversed(transforms): if key == 's': lines.append('Scale {:f} {:f} {:f}\n'.format(vals, vals, vals)) elif key == 'r': deg = np.rad2deg(vals[0]) ax = vals[1:4] ax /= np.linalg.norm(ax) lines.append('Rotate {:f} {:f} {:f} {:f}\n'.format(deg, ax[0], ax[1], ax[2])) elif key == 't': lines.append('Translate {:f} {:f} {:f}\n'.format(vals[0], vals[1], vals[2])) # Original shape. shape_name = shape_info['name'] if shape_name == 'curve': # This is the only viable option for now. points = ndarray(shape_info['point']).ravel() assert points.size == 12 type_info = '"string type" "flat"' if 'type' in shape_info: type_info = '"string type" "{}"'.format(shape_info['type']) width_info = '"float width" [1.0]' if 'width' in shape_info: width_info = '"float width" [{}]'.format(float(shape_info['width'])) lines.append('Shape "curve" "point P" [' + ' '.join([str(v) for v in points]) + '] {} {}\n'.format(type_info, width_info)) elif shape_name == 'sphere': radius = float(shape_info['radius']) center = ndarray(shape_info['center']).ravel() assert center.size == 3 lines.append('Translate {:f} {:f} {:f}\n'.format(center[0], center[1], center[2])) lines.append('Shape "sphere" "float radius" [{:f}]'.format(radius)) else: raise NotImplementedError lines.append('AttributeEnd\n') # Write back script. with open(shape_pbrt_name, 'w') as f: for l in lines: f.write(l) self.__shape_objects.append(shape_pbrt_short_name) # Call this function after you have set up add_tri_mesh. def render(self, verbose=False, nproc=None): scene_pbrt_name = self.__temporary_folder / 'scene.pbrt' with open(scene_pbrt_name, 'w') as f: x_res, y_res = self.__resolution f.write('Film "image" "integer xresolution" [{:d}] "integer yresolution" [{:d}]\n'.format(x_res, y_res)) f.write(' "string filename" "{:s}.exr"\n'.format(self.__file_name_only)) f.write('\n') f.write('Sampler "halton" "integer pixelsamples" [{:d}]\n'.format(self.__sample)) f.write('Integrator "path" "integer maxdepth" {:d}\n'.format(self.__max_depth)) f.write('\n') # Flipped y because pbrt uses a left-handed coordinate system. cpx, cpy, cpz = self.__camera_pos clx, cly, clz = self.__camera_lookat cux, cuy, cuz = self.__camera_up f.write('LookAt {:f} {:f} {:f} {:f} {:f} {:f} {:f} {:f} {:f}\n'.format( cpx, -cpy, cpz, clx, -cly, clz, cux, -cuy, cuz)) f.write('Camera "perspective" "float fov" [{:f}]\n'.format(self.__fov)) f.write('\n') f.write('WorldBegin\n') f.write('\n') f.write('AttributeBegin\n') f.write('LightSource "infinite" "string mapname" "{}" "color scale" [{:f}, {:f}, {:f}]\n'.format( str(self.__lightmap), self.__lightmap_scale, self.__lightmap_scale, self.__lightmap_scale)) f.write('AttributeEnd\n') f.write('\n') for tri_pbrt_name in self.__tri_objects: f.write('Include "{}"\n'.format(tri_pbrt_name)) for shape_pbrt_name in self.__shape_objects: f.write('Include "{}"\n'.format(shape_pbrt_name)) f.write('\n') f.write('WorldEnd\n') verbose_flag = ' ' if verbose else '--quiet' thread_flag = ' ' if nproc is None else '--nthreads {:d}'.format(int(nproc)) os.system('{} {} {} {}'.format(str(Path(root_path) / 'external/pbrt_build/pbrt'), verbose_flag, thread_flag, scene_pbrt_name)) os.system('convert {}.exr {}.png'.format(self.__file_name_only, self.__file_name_only)) os.remove('{}.exr'.format(self.__file_name_only)) # Cleanup data. shutil.rmtree(self.__temporary_folder)