Upload 5 files

lib/renderer/camera.py

@@ -0,0 +1,226 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import cv2
+import numpy as np
+
+from .glm import ortho
+
+
+class Camera:
+    def __init__(self, width=1600, height=1200):
+        # Focal Length
+        # equivalent 50mm
+        focal = np.sqrt(width * width + height * height)
+        self.focal_x = focal
+        self.focal_y = focal
+        # Principal Point Offset
+        self.principal_x = width / 2
+        self.principal_y = height / 2
+        # Axis Skew
+        self.skew = 0
+        # Image Size
+        self.width = width
+        self.height = height
+
+        self.near = 1
+        self.far = 10
+
+        # Camera Center
+        self.center = np.array([0, 0, 1.6])
+        self.direction = np.array([0, 0, -1])
+        self.right = np.array([1, 0, 0])
+        self.up = np.array([0, 1, 0])
+
+        self.ortho_ratio = None
+
+    def sanity_check(self):
+        self.center = self.center.reshape([-1])
+        self.direction = self.direction.reshape([-1])
+        self.right = self.right.reshape([-1])
+        self.up = self.up.reshape([-1])
+
+        assert len(self.center) == 3
+        assert len(self.direction) == 3
+        assert len(self.right) == 3
+        assert len(self.up) == 3
+
+    @staticmethod
+    def normalize_vector(v):
+        v_norm = np.linalg.norm(v)
+        return v if v_norm == 0 else v / v_norm
+
+    def get_real_z_value(self, z):
+        z_near = self.near
+        z_far = self.far
+        z_n = 2.0 * z - 1.0
+        z_e = 2.0 * z_near * z_far / (z_far + z_near - z_n * (z_far - z_near))
+        return z_e
+
+    def get_rotation_matrix(self):
+        rot_mat = np.eye(3)
+        s = self.right
+        s = self.normalize_vector(s)
+        rot_mat[0, :] = s
+        u = self.up
+        u = self.normalize_vector(u)
+        rot_mat[1, :] = -u
+        rot_mat[2, :] = self.normalize_vector(self.direction)
+
+        return rot_mat
+
+    def get_translation_vector(self):
+        rot_mat = self.get_rotation_matrix()
+        trans = -np.dot(rot_mat, self.center)
+        return trans
+
+    def get_intrinsic_matrix(self):
+        int_mat = np.eye(3)
+
+        int_mat[0, 0] = self.focal_x
+        int_mat[1, 1] = self.focal_y
+        int_mat[0, 1] = self.skew
+        int_mat[0, 2] = self.principal_x
+        int_mat[1, 2] = self.principal_y
+
+        return int_mat
+
+    def get_projection_matrix(self):
+        ext_mat = self.get_extrinsic_matrix()
+        int_mat = self.get_intrinsic_matrix()
+
+        return np.matmul(int_mat, ext_mat)
+
+    def get_extrinsic_matrix(self):
+        rot_mat = self.get_rotation_matrix()
+        int_mat = self.get_intrinsic_matrix()
+        trans = self.get_translation_vector()
+
+        extrinsic = np.eye(4)
+        extrinsic[:3, :3] = rot_mat
+        extrinsic[:3, 3] = trans
+
+        return extrinsic[:3, :]
+
+    def set_rotation_matrix(self, rot_mat):
+        self.direction = rot_mat[2, :]
+        self.up = -rot_mat[1, :]
+        self.right = rot_mat[0, :]
+
+    def set_intrinsic_matrix(self, int_mat):
+        self.focal_x = int_mat[0, 0]
+        self.focal_y = int_mat[1, 1]
+        self.skew = int_mat[0, 1]
+        self.principal_x = int_mat[0, 2]
+        self.principal_y = int_mat[1, 2]
+
+    def set_projection_matrix(self, proj_mat):
+        res = cv2.decomposeProjectionMatrix(proj_mat)
+        int_mat, rot_mat, camera_center_homo = res[0], res[1], res[2]
+        camera_center = camera_center_homo[0:3] / camera_center_homo[3]
+        camera_center = camera_center.reshape(-1)
+        int_mat = int_mat / int_mat[2][2]
+
+        self.set_intrinsic_matrix(int_mat)
+        self.set_rotation_matrix(rot_mat)
+        self.center = camera_center
+
+        self.sanity_check()
+
+    def get_gl_matrix(self):
+        z_near = self.near
+        z_far = self.far
+        rot_mat = self.get_rotation_matrix()
+        int_mat = self.get_intrinsic_matrix()
+        trans = self.get_translation_vector()
+
+        extrinsic = np.eye(4)
+        extrinsic[:3, :3] = rot_mat
+        extrinsic[:3, 3] = trans
+        axis_adj = np.eye(4)
+        axis_adj[2, 2] = -1
+        axis_adj[1, 1] = -1
+        model_view = np.matmul(axis_adj, extrinsic)
+
+        projective = np.zeros([4, 4])
+        projective[:2, :2] = int_mat[:2, :2]
+        projective[:2, 2:3] = -int_mat[:2, 2:3]
+        projective[3, 2] = -1
+        projective[2, 2] = (z_near + z_far)
+        projective[2, 3] = (z_near * z_far)
+
+        if self.ortho_ratio is None:
+            ndc = ortho(0, self.width, 0, self.height, z_near, z_far)
+            perspective = np.matmul(ndc, projective)
+        else:
+            perspective = ortho(-self.width * self.ortho_ratio / 2,
+                                self.width * self.ortho_ratio / 2,
+                                -self.height * self.ortho_ratio / 2,
+                                self.height * self.ortho_ratio / 2, z_near,
+                                z_far)
+
+        return perspective, model_view
+
+
+def KRT_from_P(proj_mat, normalize_K=True):
+    res = cv2.decomposeProjectionMatrix(proj_mat)
+    K, Rot, camera_center_homog = res[0], res[1], res[2]
+    camera_center = camera_center_homog[0:3] / camera_center_homog[3]
+    trans = -Rot.dot(camera_center)
+    if normalize_K:
+        K = K / K[2][2]
+    return K, Rot, trans
+
+
+def MVP_from_P(proj_mat, width, height, near=0.1, far=10000):
+    '''
+    Convert OpenCV camera calibration matrix to OpenGL projection and model view matrix
+    :param proj_mat: OpenCV camera projeciton matrix
+    :param width: Image width
+    :param height: Image height
+    :param near: Z near value
+    :param far: Z far value
+    :return: OpenGL projection matrix and model view matrix
+    '''
+    res = cv2.decomposeProjectionMatrix(proj_mat)
+    K, Rot, camera_center_homog = res[0], res[1], res[2]
+    camera_center = camera_center_homog[0:3] / camera_center_homog[3]
+    trans = -Rot.dot(camera_center)
+    K = K / K[2][2]
+
+    extrinsic = np.eye(4)
+    extrinsic[:3, :3] = Rot
+    extrinsic[:3, 3:4] = trans
+    axis_adj = np.eye(4)
+    axis_adj[2, 2] = -1
+    axis_adj[1, 1] = -1
+    model_view = np.matmul(axis_adj, extrinsic)
+
+    zFar = far
+    zNear = near
+    projective = np.zeros([4, 4])
+    projective[:2, :2] = K[:2, :2]
+    projective[:2, 2:3] = -K[:2, 2:3]
+    projective[3, 2] = -1
+    projective[2, 2] = (zNear + zFar)
+    projective[2, 3] = (zNear * zFar)
+
+    ndc = ortho(0, width, 0, height, zNear, zFar)
+
+    perspective = np.matmul(ndc, projective)
+
+    return perspective, model_view

lib/renderer/glm.py

@@ -0,0 +1,143 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import numpy as np
+
+
+def vec3(x, y, z):
+    return np.array([x, y, z], dtype=np.float32)
+
+
+def radians(v):
+    return np.radians(v)
+
+
+def identity():
+    return np.identity(4, dtype=np.float32)
+
+
+def empty():
+    return np.zeros([4, 4], dtype=np.float32)
+
+
+def magnitude(v):
+    return np.linalg.norm(v)
+
+
+def normalize(v):
+    m = magnitude(v)
+    return v if m == 0 else v / m
+
+
+def dot(u, v):
+    return np.sum(u * v)
+
+
+def cross(u, v):
+    res = vec3(0, 0, 0)
+    res[0] = u[1] * v[2] - u[2] * v[1]
+    res[1] = u[2] * v[0] - u[0] * v[2]
+    res[2] = u[0] * v[1] - u[1] * v[0]
+    return res
+
+
+# below functions can be optimized
+
+
+def translate(m, v):
+    res = np.copy(m)
+    res[:, 3] = m[:, 0] * v[0] + m[:, 1] * v[1] + m[:, 2] * v[2] + m[:, 3]
+    return res
+
+
+def rotate(m, angle, v):
+    a = angle
+    c = np.cos(a)
+    s = np.sin(a)
+
+    axis = normalize(v)
+    temp = (1 - c) * axis
+
+    rot = empty()
+    rot[0][0] = c + temp[0] * axis[0]
+    rot[0][1] = temp[0] * axis[1] + s * axis[2]
+    rot[0][2] = temp[0] * axis[2] - s * axis[1]
+
+    rot[1][0] = temp[1] * axis[0] - s * axis[2]
+    rot[1][1] = c + temp[1] * axis[1]
+    rot[1][2] = temp[1] * axis[2] + s * axis[0]
+
+    rot[2][0] = temp[2] * axis[0] + s * axis[1]
+    rot[2][1] = temp[2] * axis[1] - s * axis[0]
+    rot[2][2] = c + temp[2] * axis[2]
+
+    res = empty()
+    res[:, 0] = m[:, 0] * rot[0][0] + m[:, 1] * rot[0][1] + m[:, 2] * rot[0][2]
+    res[:, 1] = m[:, 0] * rot[1][0] + m[:, 1] * rot[1][1] + m[:, 2] * rot[1][2]
+    res[:, 2] = m[:, 0] * rot[2][0] + m[:, 1] * rot[2][1] + m[:, 2] * rot[2][2]
+    res[:, 3] = m[:, 3]
+    return res
+
+
+def perspective(fovy, aspect, zNear, zFar):
+    tanHalfFovy = np.tan(fovy / 2)
+
+    res = empty()
+    res[0][0] = 1 / (aspect * tanHalfFovy)
+    res[1][1] = 1 / (tanHalfFovy)
+    res[2][3] = -1
+    res[2][2] = -(zFar + zNear) / (zFar - zNear)
+    res[3][2] = -(2 * zFar * zNear) / (zFar - zNear)
+
+    return res.T
+
+
+def ortho(left, right, bottom, top, zNear, zFar):
+    # res = np.ones([4, 4], dtype=np.float32)
+    res = identity()
+    res[0][0] = 2 / (right - left)
+    res[1][1] = 2 / (top - bottom)
+    res[2][2] = -2 / (zFar - zNear)
+    res[3][0] = -(right + left) / (right - left)
+    res[3][1] = -(top + bottom) / (top - bottom)
+    res[3][2] = -(zFar + zNear) / (zFar - zNear)
+    return res.T
+
+
+def lookat(eye, center, up):
+    f = normalize(center - eye)
+    s = normalize(cross(f, up))
+    u = cross(s, f)
+
+    res = identity()
+    res[0][0] = s[0]
+    res[1][0] = s[1]
+    res[2][0] = s[2]
+    res[0][1] = u[0]
+    res[1][1] = u[1]
+    res[2][1] = u[2]
+    res[0][2] = -f[0]
+    res[1][2] = -f[1]
+    res[2][2] = -f[2]
+    res[3][0] = -dot(s, eye)
+    res[3][1] = -dot(u, eye)
+    res[3][2] = -dot(f, eye)
+    return res.T
+
+
+def transform(d, m):
+    return np.dot(m, d.T).T

lib/renderer/mesh.py

@@ -0,0 +1,526 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from lib.dataset.mesh_util import SMPLX
+from lib.common.render_utils import face_vertices
+import numpy as np
+import lib.smplx as smplx
+import trimesh
+import torch
+import torch.nn.functional as F
+
+model_init_params = dict(
+    gender='male',
+    model_type='smplx',
+    model_path=SMPLX().model_dir,
+    create_global_orient=False,
+    create_body_pose=False,
+    create_betas=False,
+    create_left_hand_pose=False,
+    create_right_hand_pose=False,
+    create_expression=False,
+    create_jaw_pose=False,
+    create_leye_pose=False,
+    create_reye_pose=False,
+    create_transl=False,
+    num_pca_comps=12)
+
+
+def get_smpl_model(model_type, gender): return smplx.create(
+    **model_init_params)
+
+
+def normalization(data):
+    _range = np.max(data) - np.min(data)
+    return ((data - np.min(data)) / _range)
+
+
+def sigmoid(x):
+    z = 1 / (1 + np.exp(-x))
+    return z
+
+
+def load_fit_body(fitted_path, scale, smpl_type='smplx', smpl_gender='neutral', noise_dict=None):
+
+    param = np.load(fitted_path, allow_pickle=True)
+    for key in param.keys():
+        param[key] = torch.as_tensor(param[key])
+
+    smpl_model = get_smpl_model(smpl_type, smpl_gender)
+    model_forward_params = dict(betas=param['betas'],
+                                global_orient=param['global_orient'],
+                                body_pose=param['body_pose'],
+                                left_hand_pose=param['left_hand_pose'],
+                                right_hand_pose=param['right_hand_pose'],
+                                jaw_pose=param['jaw_pose'],
+                                leye_pose=param['leye_pose'],
+                                reye_pose=param['reye_pose'],
+                                expression=param['expression'],
+                                return_verts=True)
+
+    if noise_dict is not None:
+        model_forward_params.update(noise_dict)
+
+    smpl_out = smpl_model(**model_forward_params)
+
+    smpl_verts = (
+        (smpl_out.vertices[0] * param['scale'] + param['translation']) * scale).detach()
+    smpl_joints = (
+        (smpl_out.joints[0] * param['scale'] + param['translation']) * scale).detach()
+    smpl_mesh = trimesh.Trimesh(smpl_verts,
+                                smpl_model.faces,
+                                process=False, maintain_order=True)
+
+    return smpl_mesh, smpl_joints
+
+
+def load_ori_fit_body(fitted_path, smpl_type='smplx', smpl_gender='neutral'):
+
+    param = np.load(fitted_path, allow_pickle=True)
+    for key in param.keys():
+        param[key] = torch.as_tensor(param[key])
+
+    smpl_model = get_smpl_model(smpl_type, smpl_gender)
+    model_forward_params = dict(betas=param['betas'],
+                                global_orient=param['global_orient'],
+                                body_pose=param['body_pose'],
+                                left_hand_pose=param['left_hand_pose'],
+                                right_hand_pose=param['right_hand_pose'],
+                                jaw_pose=param['jaw_pose'],
+                                leye_pose=param['leye_pose'],
+                                reye_pose=param['reye_pose'],
+                                expression=param['expression'],
+                                return_verts=True)
+
+    smpl_out = smpl_model(**model_forward_params)
+
+    smpl_verts = smpl_out.vertices[0].detach()
+    smpl_mesh = trimesh.Trimesh(smpl_verts,
+                                smpl_model.faces,
+                                process=False, maintain_order=True)
+
+    return smpl_mesh
+
+
+def save_obj_mesh(mesh_path, verts, faces):
+    file = open(mesh_path, 'w')
+    for v in verts:
+        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
+    for f in faces:
+        f_plus = f + 1
+        file.write('f %d %d %d\n' % (f_plus[0], f_plus[1], f_plus[2]))
+    file.close()
+
+
+# https://github.com/ratcave/wavefront_reader
+def read_mtlfile(fname):
+    materials = {}
+    with open(fname) as f:
+        lines = f.read().splitlines()
+
+    for line in lines:
+        if line:
+            split_line = line.strip().split(' ', 1)
+            if len(split_line) < 2:
+                continue
+
+            prefix, data = split_line[0], split_line[1]
+            if 'newmtl' in prefix:
+                material = {}
+                materials[data] = material
+            elif materials:
+                if data:
+                    split_data = data.strip().split(' ')
+
+                    # assume texture maps are in the same level
+                    # WARNING: do not include space in your filename!!
+                    if 'map' in prefix:
+                        material[prefix] = split_data[-1].split('\\')[-1]
+                    elif len(split_data) > 1:
+                        material[prefix] = tuple(float(d) for d in split_data)
+                    else:
+                        try:
+                            material[prefix] = int(data)
+                        except ValueError:
+                            material[prefix] = float(data)
+
+    return materials
+
+
+def load_obj_mesh_mtl(mesh_file):
+    vertex_data = []
+    norm_data = []
+    uv_data = []
+
+    face_data = []
+    face_norm_data = []
+    face_uv_data = []
+
+    # face per material
+    face_data_mat = {}
+    face_norm_data_mat = {}
+    face_uv_data_mat = {}
+
+    # current material name
+    mtl_data = None
+    cur_mat = None
+
+    if isinstance(mesh_file, str):
+        f = open(mesh_file, "r")
+    else:
+        f = mesh_file
+    for line in f:
+        if isinstance(line, bytes):
+            line = line.decode("utf-8")
+        if line.startswith('#'):
+            continue
+        values = line.split()
+        if not values:
+            continue
+
+        if values[0] == 'v':
+            v = list(map(float, values[1:4]))
+            vertex_data.append(v)
+        elif values[0] == 'vn':
+            vn = list(map(float, values[1:4]))
+            norm_data.append(vn)
+        elif values[0] == 'vt':
+            vt = list(map(float, values[1:3]))
+            uv_data.append(vt)
+        elif values[0] == 'mtllib':
+            mtl_data = read_mtlfile(
+                mesh_file.replace(mesh_file.split('/')[-1], values[1]))
+        elif values[0] == 'usemtl':
+            cur_mat = values[1]
+        elif values[0] == 'f':
+            # local triangle data
+            l_face_data = []
+            l_face_uv_data = []
+            l_face_norm_data = []
+
+            # quad mesh
+            if len(values) > 4:
+                f = list(
+                    map(
+                        lambda x: int(x.split('/')[0]) if int(x.split('/')[0])
+                        < 0 else int(x.split('/')[0]) - 1, values[1:4]))
+                l_face_data.append(f)
+                f = list(
+                    map(
+                        lambda x: int(x.split('/')[0])
+                        if int(x.split('/')[0]) < 0 else int(x.split('/')[0]) -
+                        1, [values[3], values[4], values[1]]))
+                l_face_data.append(f)
+            # tri mesh
+            else:
+                f = list(
+                    map(
+                        lambda x: int(x.split('/')[0]) if int(x.split('/')[0])
+                        < 0 else int(x.split('/')[0]) - 1, values[1:4]))
+                l_face_data.append(f)
+            # deal with texture
+            if len(values[1].split('/')) >= 2:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[1])
+                            if int(x.split('/')[1]) < 0 else int(
+                                x.split('/')[1]) - 1, values[1:4]))
+                    l_face_uv_data.append(f)
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[1])
+                            if int(x.split('/')[1]) < 0 else int(
+                                x.split('/')[1]) - 1,
+                            [values[3], values[4], values[1]]))
+                    l_face_uv_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[1]) != 0:
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[1])
+                            if int(x.split('/')[1]) < 0 else int(
+                                x.split('/')[1]) - 1, values[1:4]))
+                    l_face_uv_data.append(f)
+            # deal with normal
+            if len(values[1].split('/')) == 3:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[2])
+                            if int(x.split('/')[2]) < 0 else int(
+                                x.split('/')[2]) - 1, values[1:4]))
+                    l_face_norm_data.append(f)
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[2])
+                            if int(x.split('/')[2]) < 0 else int(
+                                x.split('/')[2]) - 1,
+                            [values[3], values[4], values[1]]))
+                    l_face_norm_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[2]) != 0:
+                    f = list(
+                        map(
+                            lambda x: int(x.split('/')[2])
+                            if int(x.split('/')[2]) < 0 else int(
+                                x.split('/')[2]) - 1, values[1:4]))
+                    l_face_norm_data.append(f)
+
+            face_data += l_face_data
+            face_uv_data += l_face_uv_data
+            face_norm_data += l_face_norm_data
+
+            if cur_mat is not None:
+                if cur_mat not in face_data_mat.keys():
+                    face_data_mat[cur_mat] = []
+                if cur_mat not in face_uv_data_mat.keys():
+                    face_uv_data_mat[cur_mat] = []
+                if cur_mat not in face_norm_data_mat.keys():
+                    face_norm_data_mat[cur_mat] = []
+                face_data_mat[cur_mat] += l_face_data
+                face_uv_data_mat[cur_mat] += l_face_uv_data
+                face_norm_data_mat[cur_mat] += l_face_norm_data
+
+    vertices = np.array(vertex_data)
+    faces = np.array(face_data)
+
+    norms = np.array(norm_data)
+    norms = normalize_v3(norms)
+    face_normals = np.array(face_norm_data)
+
+    uvs = np.array(uv_data)
+    face_uvs = np.array(face_uv_data)
+
+    out_tuple = (vertices, faces, norms, face_normals, uvs, face_uvs)
+
+    if cur_mat is not None and mtl_data is not None:
+        for key in face_data_mat:
+            face_data_mat[key] = np.array(face_data_mat[key])
+            face_uv_data_mat[key] = np.array(face_uv_data_mat[key])
+            face_norm_data_mat[key] = np.array(face_norm_data_mat[key])
+
+        out_tuple += (face_data_mat, face_norm_data_mat, face_uv_data_mat,
+                      mtl_data)
+
+    return out_tuple
+
+
+def load_scan(mesh_file, with_normal=False, with_texture=False):
+    vertex_data = []
+    norm_data = []
+    uv_data = []
+
+    face_data = []
+    face_norm_data = []
+    face_uv_data = []
+
+    if isinstance(mesh_file, str):
+        f = open(mesh_file, "r")
+    else:
+        f = mesh_file
+    for line in f:
+        if isinstance(line, bytes):
+            line = line.decode("utf-8")
+        if line.startswith('#'):
+            continue
+        values = line.split()
+        if not values:
+            continue
+
+        if values[0] == 'v':
+            v = list(map(float, values[1:4]))
+            vertex_data.append(v)
+        elif values[0] == 'vn':
+            vn = list(map(float, values[1:4]))
+            norm_data.append(vn)
+        elif values[0] == 'vt':
+            vt = list(map(float, values[1:3]))
+            uv_data.append(vt)
+
+        elif values[0] == 'f':
+            # quad mesh
+            if len(values) > 4:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+                f = list(
+                    map(lambda x: int(x.split('/')[0]),
+                        [values[3], values[4], values[1]]))
+                face_data.append(f)
+            # tri mesh
+            else:
+                f = list(map(lambda x: int(x.split('/')[0]), values[1:4]))
+                face_data.append(f)
+
+            # deal with texture
+            if len(values[1].split('/')) >= 2:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
+                    face_uv_data.append(f)
+                    f = list(
+                        map(lambda x: int(x.split('/')[1]),
+                            [values[3], values[4], values[1]]))
+                    face_uv_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[1]) != 0:
+                    f = list(map(lambda x: int(x.split('/')[1]), values[1:4]))
+                    face_uv_data.append(f)
+            # deal with normal
+            if len(values[1].split('/')) == 3:
+                # quad mesh
+                if len(values) > 4:
+                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
+                    face_norm_data.append(f)
+                    f = list(
+                        map(lambda x: int(x.split('/')[2]),
+                            [values[3], values[4], values[1]]))
+                    face_norm_data.append(f)
+                # tri mesh
+                elif len(values[1].split('/')[2]) != 0:
+                    f = list(map(lambda x: int(x.split('/')[2]), values[1:4]))
+                    face_norm_data.append(f)
+
+    vertices = np.array(vertex_data)
+    faces = np.array(face_data) - 1
+
+    if with_texture and with_normal:
+        uvs = np.array(uv_data)
+        face_uvs = np.array(face_uv_data) - 1
+        norms = np.array(norm_data)
+        if norms.shape[0] == 0:
+            norms = compute_normal(vertices, faces)
+            face_normals = faces
+        else:
+            norms = normalize_v3(norms)
+            face_normals = np.array(face_norm_data) - 1
+        return vertices, faces, norms, face_normals, uvs, face_uvs
+
+    if with_texture:
+        uvs = np.array(uv_data)
+        face_uvs = np.array(face_uv_data) - 1
+        return vertices, faces, uvs, face_uvs
+
+    if with_normal:
+        norms = np.array(norm_data)
+        norms = normalize_v3(norms)
+        face_normals = np.array(face_norm_data) - 1
+        return vertices, faces, norms, face_normals
+
+    return vertices, faces
+
+
+def normalize_v3(arr):
+    ''' Normalize a numpy array of 3 component vectors shape=(n,3) '''
+    lens = np.sqrt(arr[:, 0]**2 + arr[:, 1]**2 + arr[:, 2]**2)
+    eps = 0.00000001
+    lens[lens < eps] = eps
+    arr[:, 0] /= lens
+    arr[:, 1] /= lens
+    arr[:, 2] /= lens
+    return arr
+
+
+def compute_normal(vertices, faces):
+    # Create a zeroed array with the same type and shape as our vertices i.e., per vertex normal
+    norm = np.zeros(vertices.shape, dtype=vertices.dtype)
+    # Create an indexed view into the vertex array using the array of three indices for triangles
+    tris = vertices[faces]
+    # Calculate the normal for all the triangles, by taking the cross product of the vectors v1-v0, and v2-v0 in each triangle
+    n = np.cross(tris[::, 1] - tris[::, 0], tris[::, 2] - tris[::, 0])
+    # n is now an array of normals per triangle. The length of each normal is dependent the vertices,
+    # we need to normalize these, so that our next step weights each normal equally.
+    normalize_v3(n)
+    # now we have a normalized array of normals, one per triangle, i.e., per triangle normals.
+    # But instead of one per triangle (i.e., flat shading), we add to each vertex in that triangle,
+    # the triangles' normal. Multiple triangles would then contribute to every vertex, so we need to normalize again afterwards.
+    # The cool part, we can actually add the normals through an indexed view of our (zeroed) per vertex normal array
+    norm[faces[:, 0]] += n
+    norm[faces[:, 1]] += n
+    norm[faces[:, 2]] += n
+    normalize_v3(norm)
+
+    return norm
+
+
+def compute_normal_batch(vertices, faces):
+
+    bs, nv = vertices.shape[:2]
+    bs, nf = faces.shape[:2]
+
+    vert_norm = torch.zeros(bs * nv, 3).type_as(vertices)
+    tris = face_vertices(vertices, faces)
+    face_norm = F.normalize(torch.cross(tris[:, :, 1] - tris[:, :, 0],
+                                        tris[:, :, 2] - tris[:, :, 0]),
+                            dim=-1)
+
+    faces = (faces +
+             (torch.arange(bs).type_as(faces) * nv)[:, None, None]).view(
+                 -1, 3)
+
+    vert_norm[faces[:, 0]] += face_norm.view(-1, 3)
+    vert_norm[faces[:, 1]] += face_norm.view(-1, 3)
+    vert_norm[faces[:, 2]] += face_norm.view(-1, 3)
+
+    vert_norm = F.normalize(vert_norm, dim=-1).view(bs, nv, 3)
+
+    return vert_norm
+
+
+# compute tangent and bitangent
+def compute_tangent(vertices, faces, normals, uvs, faceuvs):
+    # NOTE: this could be numerically unstable around [0,0,1]
+    # but other current solutions are pretty freaky somehow
+    c1 = np.cross(normals, np.array([0, 1, 0.0]))
+    tan = c1
+    normalize_v3(tan)
+    btan = np.cross(normals, tan)
+
+    # NOTE: traditional version is below
+
+    # pts_tris = vertices[faces]
+    # uv_tris = uvs[faceuvs]
+
+    # W = np.stack([pts_tris[::, 1] - pts_tris[::, 0], pts_tris[::, 2] - pts_tris[::, 0]],2)
+    # UV = np.stack([uv_tris[::, 1] - uv_tris[::, 0], uv_tris[::, 2] - uv_tris[::, 0]], 1)
+
+    # for i in range(W.shape[0]):
+    #     W[i,::] = W[i,::].dot(np.linalg.inv(UV[i,::]))
+
+    # tan = np.zeros(vertices.shape, dtype=vertices.dtype)
+    # tan[faces[:,0]] += W[:,:,0]
+    # tan[faces[:,1]] += W[:,:,0]
+    # tan[faces[:,2]] += W[:,:,0]
+
+    # btan = np.zeros(vertices.shape, dtype=vertices.dtype)
+    # btan[faces[:,0]] += W[:,:,1]
+    # btan[faces[:,1]] += W[:,:,1]
+    # btan[faces[:,2]] += W[:,:,1]
+
+    # normalize_v3(tan)
+
+    # ndott = np.sum(normals*tan, 1, keepdims=True)
+    # tan = tan - ndott * normals
+
+    # normalize_v3(btan)
+    # normalize_v3(tan)
+
+    # tan[np.sum(np.cross(normals, tan) * btan, 1) < 0,:] *= -1.0
+
+    return tan, btan

lib/renderer/opengl_util.py

@@ -0,0 +1,369 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+
+from lib.renderer.mesh import load_scan, compute_tangent
+from lib.renderer.camera import Camera
+import cv2
+import math
+import random
+import numpy as np
+
+
+def render_result(rndr, shader_id, path, mask=False):
+
+    cam_render = rndr.get_color(shader_id)
+    cam_render = cv2.cvtColor(cam_render, cv2.COLOR_RGBA2BGRA)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    if shader_id != 2:
+        cv2.imwrite(path, np.uint8(255.0 * cam_render))
+    else:
+        cam_render[:, :, -1] -= 0.5
+        cam_render[:, :, -1] *= 2.0
+        if not mask:
+            cv2.imwrite(path, np.uint8(255.0 / 2.0 * (cam_render + 1.0)))
+        else:
+            cv2.imwrite(path, np.uint8(-1.0 * cam_render[:, :, [3]]))
+
+
+def make_rotate(rx, ry, rz):
+    sinX = np.sin(rx)
+    sinY = np.sin(ry)
+    sinZ = np.sin(rz)
+
+    cosX = np.cos(rx)
+    cosY = np.cos(ry)
+    cosZ = np.cos(rz)
+
+    Rx = np.zeros((3, 3))
+    Rx[0, 0] = 1.0
+    Rx[1, 1] = cosX
+    Rx[1, 2] = -sinX
+    Rx[2, 1] = sinX
+    Rx[2, 2] = cosX
+
+    Ry = np.zeros((3, 3))
+    Ry[0, 0] = cosY
+    Ry[0, 2] = sinY
+    Ry[1, 1] = 1.0
+    Ry[2, 0] = -sinY
+    Ry[2, 2] = cosY
+
+    Rz = np.zeros((3, 3))
+    Rz[0, 0] = cosZ
+    Rz[0, 1] = -sinZ
+    Rz[1, 0] = sinZ
+    Rz[1, 1] = cosZ
+    Rz[2, 2] = 1.0
+
+    R = np.matmul(np.matmul(Rz, Ry), Rx)
+    return R
+
+
+def rotateSH(SH, R):
+    SHn = SH
+
+    # 1st order
+    SHn[1] = R[1, 1] * SH[1] - R[1, 2] * SH[2] + R[1, 0] * SH[3]
+    SHn[2] = -R[2, 1] * SH[1] + R[2, 2] * SH[2] - R[2, 0] * SH[3]
+    SHn[3] = R[0, 1] * SH[1] - R[0, 2] * SH[2] + R[0, 0] * SH[3]
+
+    # 2nd order
+    SHn[4:, 0] = rotateBand2(SH[4:, 0], R)
+    SHn[4:, 1] = rotateBand2(SH[4:, 1], R)
+    SHn[4:, 2] = rotateBand2(SH[4:, 2], R)
+
+    return SHn
+
+
+def rotateBand2(x, R):
+    s_c3 = 0.94617469575
+    s_c4 = -0.31539156525
+    s_c5 = 0.54627421529
+
+    s_c_scale = 1.0 / 0.91529123286551084
+    s_c_scale_inv = 0.91529123286551084
+
+    s_rc2 = 1.5853309190550713 * s_c_scale
+    s_c4_div_c3 = s_c4 / s_c3
+    s_c4_div_c3_x2 = (s_c4 / s_c3) * 2.0
+
+    s_scale_dst2 = s_c3 * s_c_scale_inv
+    s_scale_dst4 = s_c5 * s_c_scale_inv
+
+    sh0 = x[3] + x[4] + x[4] - x[1]
+    sh1 = x[0] + s_rc2 * x[2] + x[3] + x[4]
+    sh2 = x[0]
+    sh3 = -x[3]
+    sh4 = -x[1]
+
+    r2x = R[0][0] + R[0][1]
+    r2y = R[1][0] + R[1][1]
+    r2z = R[2][0] + R[2][1]
+
+    r3x = R[0][0] + R[0][2]
+    r3y = R[1][0] + R[1][2]
+    r3z = R[2][0] + R[2][2]
+
+    r4x = R[0][1] + R[0][2]
+    r4y = R[1][1] + R[1][2]
+    r4z = R[2][1] + R[2][2]
+
+    sh0_x = sh0 * R[0][0]
+    sh0_y = sh0 * R[1][0]
+    d0 = sh0_x * R[1][0]
+    d1 = sh0_y * R[2][0]
+    d2 = sh0 * (R[2][0] * R[2][0] + s_c4_div_c3)
+    d3 = sh0_x * R[2][0]
+    d4 = sh0_x * R[0][0] - sh0_y * R[1][0]
+
+    sh1_x = sh1 * R[0][2]
+    sh1_y = sh1 * R[1][2]
+    d0 += sh1_x * R[1][2]
+    d1 += sh1_y * R[2][2]
+    d2 += sh1 * (R[2][2] * R[2][2] + s_c4_div_c3)
+    d3 += sh1_x * R[2][2]
+    d4 += sh1_x * R[0][2] - sh1_y * R[1][2]
+
+    sh2_x = sh2 * r2x
+    sh2_y = sh2 * r2y
+    d0 += sh2_x * r2y
+    d1 += sh2_y * r2z
+    d2 += sh2 * (r2z * r2z + s_c4_div_c3_x2)
+    d3 += sh2_x * r2z
+    d4 += sh2_x * r2x - sh2_y * r2y
+
+    sh3_x = sh3 * r3x
+    sh3_y = sh3 * r3y
+    d0 += sh3_x * r3y
+    d1 += sh3_y * r3z
+    d2 += sh3 * (r3z * r3z + s_c4_div_c3_x2)
+    d3 += sh3_x * r3z
+    d4 += sh3_x * r3x - sh3_y * r3y
+
+    sh4_x = sh4 * r4x
+    sh4_y = sh4 * r4y
+    d0 += sh4_x * r4y
+    d1 += sh4_y * r4z
+    d2 += sh4 * (r4z * r4z + s_c4_div_c3_x2)
+    d3 += sh4_x * r4z
+    d4 += sh4_x * r4x - sh4_y * r4y
+
+    dst = x
+    dst[0] = d0
+    dst[1] = -d1
+    dst[2] = d2 * s_scale_dst2
+    dst[3] = -d3
+    dst[4] = d4 * s_scale_dst4
+
+    return dst
+
+
+def load_calib(param, render_size=512):
+    # pixel unit / world unit
+    ortho_ratio = param['ortho_ratio']
+    # world unit / model unit
+    scale = param['scale']
+    # camera center world coordinate
+    center = param['center']
+    # model rotation
+    R = param['R']
+
+    translate = -np.matmul(R, center).reshape(3, 1)
+    extrinsic = np.concatenate([R, translate], axis=1)
+    extrinsic = np.concatenate(
+        [extrinsic, np.array([0, 0, 0, 1]).reshape(1, 4)], 0)
+    # Match camera space to image pixel space
+    scale_intrinsic = np.identity(4)
+    scale_intrinsic[0, 0] = scale / ortho_ratio
+    scale_intrinsic[1, 1] = -scale / ortho_ratio
+    scale_intrinsic[2, 2] = scale / ortho_ratio
+    # Match image pixel space to image uv space
+    uv_intrinsic = np.identity(4)
+    uv_intrinsic[0, 0] = 1.0 / float(render_size // 2)
+    uv_intrinsic[1, 1] = 1.0 / float(render_size // 2)
+    uv_intrinsic[2, 2] = 1.0 / float(render_size // 2)
+
+    intrinsic = np.matmul(uv_intrinsic, scale_intrinsic)
+    calib = np.concatenate([extrinsic, intrinsic], axis=0)
+    return calib
+
+
+def render_prt_ortho(out_path,
+                     folder_name,
+                     subject_name,
+                     shs,
+                     rndr,
+                     rndr_uv,
+                     im_size,
+                     angl_step=4,
+                     n_light=1,
+                     pitch=[0]):
+    cam = Camera(width=im_size, height=im_size)
+    cam.ortho_ratio = 0.4 * (512 / im_size)
+    cam.near = -100
+    cam.far = 100
+    cam.sanity_check()
+
+    # set path for obj, prt
+    mesh_file = os.path.join(folder_name, subject_name + '_100k.obj')
+    if not os.path.exists(mesh_file):
+        print('ERROR: obj file does not exist!!', mesh_file)
+        return
+    prt_file = os.path.join(folder_name, 'bounce', 'bounce0.txt')
+    if not os.path.exists(prt_file):
+        print('ERROR: prt file does not exist!!!', prt_file)
+        return
+    face_prt_file = os.path.join(folder_name, 'bounce', 'face.npy')
+    if not os.path.exists(face_prt_file):
+        print('ERROR: face prt file does not exist!!!', prt_file)
+        return
+    text_file = os.path.join(folder_name, 'tex', subject_name + '_dif_2k.jpg')
+    if not os.path.exists(text_file):
+        print('ERROR: dif file does not exist!!', text_file)
+        return
+
+    texture_image = cv2.imread(text_file)
+    texture_image = cv2.cvtColor(texture_image, cv2.COLOR_BGR2RGB)
+
+    vertices, faces, normals, faces_normals, textures, face_textures = load_scan(
+        mesh_file, with_normal=True, with_texture=True)
+    vmin = vertices.min(0)
+    vmax = vertices.max(0)
+    up_axis = 1 if (vmax - vmin).argmax() == 1 else 2
+
+    vmed = np.median(vertices, 0)
+    vmed[up_axis] = 0.5 * (vmax[up_axis] + vmin[up_axis])
+    y_scale = 180 / (vmax[up_axis] - vmin[up_axis])
+
+    rndr.set_norm_mat(y_scale, vmed)
+    rndr_uv.set_norm_mat(y_scale, vmed)
+
+    tan, bitan = compute_tangent(vertices, faces, normals, textures,
+                                 face_textures)
+    prt = np.loadtxt(prt_file)
+    face_prt = np.load(face_prt_file)
+    rndr.set_mesh(vertices, faces, normals, faces_normals, textures,
+                  face_textures, prt, face_prt, tan, bitan)
+    rndr.set_albedo(texture_image)
+
+    rndr_uv.set_mesh(vertices, faces, normals, faces_normals, textures,
+                     face_textures, prt, face_prt, tan, bitan)
+    rndr_uv.set_albedo(texture_image)
+
+    os.makedirs(os.path.join(out_path, 'GEO', 'OBJ', subject_name),
+                exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'PARAM', subject_name), exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'RENDER', subject_name), exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'MASK', subject_name), exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'UV_RENDER', subject_name),
+                exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'UV_MASK', subject_name), exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'UV_POS', subject_name), exist_ok=True)
+    os.makedirs(os.path.join(out_path, 'UV_NORMAL', subject_name),
+                exist_ok=True)
+
+    if not os.path.exists(os.path.join(out_path, 'val.txt')):
+        f = open(os.path.join(out_path, 'val.txt'), 'w')
+        f.close()
+
+    # copy obj file
+    cmd = 'cp %s %s' % (mesh_file,
+                        os.path.join(out_path, 'GEO', 'OBJ', subject_name))
+    print(cmd)
+    os.system(cmd)
+
+    for p in pitch:
+        for y in tqdm(range(0, 360, angl_step)):
+            R = np.matmul(make_rotate(math.radians(p), 0, 0),
+                          make_rotate(0, math.radians(y), 0))
+            if up_axis == 2:
+                R = np.matmul(R, make_rotate(math.radians(90), 0, 0))
+
+            rndr.rot_matrix = R
+            rndr_uv.rot_matrix = R
+            rndr.set_camera(cam)
+            rndr_uv.set_camera(cam)
+
+            for j in range(n_light):
+                sh_id = random.randint(0, shs.shape[0] - 1)
+                sh = shs[sh_id]
+                sh_angle = 0.2 * np.pi * (random.random() - 0.5)
+                sh = rotateSH(sh, make_rotate(0, sh_angle, 0).T)
+
+                dic = {
+                    'sh': sh,
+                    'ortho_ratio': cam.ortho_ratio,
+                    'scale': y_scale,
+                    'center': vmed,
+                    'R': R
+                }
+
+                rndr.set_sh(sh)
+                rndr.analytic = False
+                rndr.use_inverse_depth = False
+                rndr.display()
+
+                out_all_f = rndr.get_color(0)
+                out_mask = out_all_f[:, :, 3]
+                out_all_f = cv2.cvtColor(out_all_f, cv2.COLOR_RGBA2BGR)
+
+                np.save(
+                    os.path.join(out_path, 'PARAM', subject_name,
+                                 '%d_%d_%02d.npy' % (y, p, j)), dic)
+                cv2.imwrite(
+                    os.path.join(out_path, 'RENDER', subject_name,
+                                 '%d_%d_%02d.jpg' % (y, p, j)),
+                    255.0 * out_all_f)
+                cv2.imwrite(
+                    os.path.join(out_path, 'MASK', subject_name,
+                                 '%d_%d_%02d.png' % (y, p, j)),
+                    255.0 * out_mask)
+
+                rndr_uv.set_sh(sh)
+                rndr_uv.analytic = False
+                rndr_uv.use_inverse_depth = False
+                rndr_uv.display()
+
+                uv_color = rndr_uv.get_color(0)
+                uv_color = cv2.cvtColor(uv_color, cv2.COLOR_RGBA2BGR)
+                cv2.imwrite(
+                    os.path.join(out_path, 'UV_RENDER', subject_name,
+                                 '%d_%d_%02d.jpg' % (y, p, j)),
+                    255.0 * uv_color)
+
+                if y == 0 and j == 0 and p == pitch[0]:
+                    uv_pos = rndr_uv.get_color(1)
+                    uv_mask = uv_pos[:, :, 3]
+                    cv2.imwrite(
+                        os.path.join(out_path, 'UV_MASK', subject_name,
+                                     '00.png'), 255.0 * uv_mask)
+
+                    data = {
+                        'default': uv_pos[:, :, :3]
+                    }  # default is a reserved name
+                    pyexr.write(
+                        os.path.join(out_path, 'UV_POS', subject_name,
+                                     '00.exr'), data)
+
+                    uv_nml = rndr_uv.get_color(2)
+                    uv_nml = cv2.cvtColor(uv_nml, cv2.COLOR_RGBA2BGR)
+                    cv2.imwrite(
+                        os.path.join(out_path, 'UV_NORMAL', subject_name,
+                                     '00.png'), 255.0 * uv_nml)

lib/renderer/prt_util.py

@@ -0,0 +1,199 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+import trimesh
+import numpy as np
+import math
+from scipy.special import sph_harm
+import argparse
+from tqdm import tqdm
+from trimesh.util import bounds_tree
+
+
+def factratio(N, D):
+    if N >= D:
+        prod = 1.0
+        for i in range(D + 1, N + 1):
+            prod *= i
+        return prod
+    else:
+        prod = 1.0
+        for i in range(N + 1, D + 1):
+            prod *= i
+        return 1.0 / prod
+
+
+def KVal(M, L):
+    return math.sqrt(((2 * L + 1) / (4 * math.pi)) * (factratio(L - M, L + M)))
+
+
+def AssociatedLegendre(M, L, x):
+    if M < 0 or M > L or np.max(np.abs(x)) > 1.0:
+        return np.zeros_like(x)
+
+    pmm = np.ones_like(x)
+    if M > 0:
+        somx2 = np.sqrt((1.0 + x) * (1.0 - x))
+        fact = 1.0
+        for i in range(1, M + 1):
+            pmm = -pmm * fact * somx2
+            fact = fact + 2
+
+    if L == M:
+        return pmm
+    else:
+        pmmp1 = x * (2 * M + 1) * pmm
+        if L == M + 1:
+            return pmmp1
+        else:
+            pll = np.zeros_like(x)
+            for i in range(M + 2, L + 1):
+                pll = (x * (2 * i - 1) * pmmp1 - (i + M - 1) * pmm) / (i - M)
+                pmm = pmmp1
+                pmmp1 = pll
+            return pll
+
+
+def SphericalHarmonic(M, L, theta, phi):
+    if M > 0:
+        return math.sqrt(2.0) * KVal(M, L) * np.cos(
+            M * phi) * AssociatedLegendre(M, L, np.cos(theta))
+    elif M < 0:
+        return math.sqrt(2.0) * KVal(-M, L) * np.sin(
+            -M * phi) * AssociatedLegendre(-M, L, np.cos(theta))
+    else:
+        return KVal(0, L) * AssociatedLegendre(0, L, np.cos(theta))
+
+
+def save_obj(mesh_path, verts):
+    file = open(mesh_path, 'w')
+    for v in verts:
+        file.write('v %.4f %.4f %.4f\n' % (v[0], v[1], v[2]))
+    file.close()
+
+
+def sampleSphericalDirections(n):
+    xv = np.random.rand(n, n)
+    yv = np.random.rand(n, n)
+    theta = np.arccos(1 - 2 * xv)
+    phi = 2.0 * math.pi * yv
+
+    phi = phi.reshape(-1)
+    theta = theta.reshape(-1)
+
+    vx = -np.sin(theta) * np.cos(phi)
+    vy = -np.sin(theta) * np.sin(phi)
+    vz = np.cos(theta)
+    return np.stack([vx, vy, vz], 1), phi, theta
+
+
+def getSHCoeffs(order, phi, theta):
+    shs = []
+    for n in range(0, order + 1):
+        for m in range(-n, n + 1):
+            s = SphericalHarmonic(m, n, theta, phi)
+            shs.append(s)
+
+    return np.stack(shs, 1)
+
+
+def computePRT(mesh_path, scale, n, order):
+
+    prt_dir = os.path.join(os.path.dirname(mesh_path), "prt")
+    bounce_path = os.path.join(prt_dir, "bounce.npy")
+    face_path = os.path.join(prt_dir, "face.npy")
+
+    os.makedirs(prt_dir, exist_ok=True)
+
+    PRT = None
+    F = None
+
+    if os.path.exists(bounce_path) and os.path.exists(face_path):
+
+        PRT = np.load(bounce_path)
+        F = np.load(face_path)
+
+    else:
+
+        mesh = trimesh.load(mesh_path,
+                            skip_materials=True,
+                            process=False,
+                            maintain_order=True)
+        mesh.vertices *= scale
+
+        vectors_orig, phi, theta = sampleSphericalDirections(n)
+        SH_orig = getSHCoeffs(order, phi, theta)
+
+        w = 4.0 * math.pi / (n * n)
+
+        origins = mesh.vertices
+        normals = mesh.vertex_normals
+        n_v = origins.shape[0]
+
+        origins = np.repeat(origins[:, None], n, axis=1).reshape(-1, 3)
+        normals = np.repeat(normals[:, None], n, axis=1).reshape(-1, 3)
+        PRT_all = None
+        for i in range(n):
+            SH = np.repeat(SH_orig[None, (i * n):((i + 1) * n)], n_v,
+                           axis=0).reshape(-1, SH_orig.shape[1])
+            vectors = np.repeat(vectors_orig[None, (i * n):((i + 1) * n)],
+                                n_v,
+                                axis=0).reshape(-1, 3)
+
+            dots = (vectors * normals).sum(1)
+            front = (dots > 0.0)
+
+            delta = 1e-3 * min(mesh.bounding_box.extents)
+
+            hits = mesh.ray.intersects_any(origins + delta * normals, vectors)
+            nohits = np.logical_and(front, np.logical_not(hits))
+
+            PRT = (nohits.astype(np.float) * dots)[:, None] * SH
+
+            if PRT_all is not None:
+                PRT_all += (PRT.reshape(-1, n, SH.shape[1]).sum(1))
+            else:
+                PRT_all = (PRT.reshape(-1, n, SH.shape[1]).sum(1))
+
+        PRT = w * PRT_all
+        F = mesh.faces
+
+        np.save(bounce_path, PRT)
+        np.save(face_path, F)
+
+    # NOTE: trimesh sometimes break the original vertex order, but topology will not change.
+    # when loading PRT in other program, use the triangle list from trimesh.
+
+    return PRT, F
+
+
+def testPRT(obj_path, n=40):
+
+    os.makedirs(os.path.join(os.path.dirname(obj_path),
+                             f'../bounce/{os.path.basename(obj_path)[:-4]}'),
+                exist_ok=True)
+
+    PRT, F = computePRT(obj_path, n, 2)
+    np.savetxt(
+        os.path.join(os.path.dirname(obj_path),
+                     f'../bounce/{os.path.basename(obj_path)[:-4]}',
+                     'bounce.npy'), PRT)
+    np.save(
+        os.path.join(os.path.dirname(obj_path),
+                     f'../bounce/{os.path.basename(obj_path)[:-4]}',
+                     'face.npy'), F)