lib/pymafx/utils/renderer.py

import imp
import json
import os
from pickle import NONE

import numpy as np
# os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
import torch
import torch.nn.functional as F
import trimesh
from core import constants, path_config
from models.smpl import get_model_faces, get_model_tpose, get_smpl_faces
# import neural_renderer as nr
from skimage.transform import resize
from torchvision.utils import make_grid
from utils.densepose_methods import DensePoseMethods
from utils.imutils import crop

from .geometry import convert_to_full_img_cam

try:
    import math

    import pyrender
    from pyrender.constants import RenderFlags
except:
    pass
try:
    from opendr.camera import ProjectPoints
    from opendr.lighting import LambertianPointLight, SphericalHarmonics
    from opendr.renderer import ColoredRenderer
except:
    pass

import logging

from pytorch3d.renderer import (
    AmbientLights,
    BlendParams,
    FoVPerspectiveCameras,
    HardFlatShader,
    HardGouraudShader,
    HardPhongShader,
    MeshRasterizer,
    MeshRenderer,
    PerspectiveCameras,
    PointLights,
    RasterizationSettings,
    SoftPhongShader,
    SoftSilhouetteShader,
    TexturesVertex,
    look_at_view_transform,
)
from pytorch3d.structures.meshes import Meshes

# from pytorch3d.renderer.mesh.renderer import MeshRendererWithFragments

logger = logging.getLogger(__name__)


class WeakPerspectiveCamera(pyrender.Camera):
    def __init__(
        self, scale, translation, znear=pyrender.camera.DEFAULT_Z_NEAR, zfar=None, name=None
    ):
        super(WeakPerspectiveCamera, self).__init__(
            znear=znear,
            zfar=zfar,
            name=name,
        )
        self.scale = scale
        self.translation = translation

    def get_projection_matrix(self, width=None, height=None):
        P = np.eye(4)
        P[0, 0] = self.scale[0]
        P[1, 1] = self.scale[1]
        P[0, 3] = self.translation[0] * self.scale[0]
        P[1, 3] = -self.translation[1] * self.scale[1]
        P[2, 2] = -1
        return P


class PyRenderer:
    def __init__(
        self, resolution=(224, 224), orig_img=False, wireframe=False, scale_ratio=1., vis_ratio=1.
    ):
        self.resolution = (resolution[0] * scale_ratio, resolution[1] * scale_ratio)
        # self.scale_ratio = scale_ratio

        self.faces = {
            'smplx': get_model_faces('smplx'),
            'smpl': get_model_faces('smpl'),
        #   'mano': get_model_faces('mano'),
        #   'flame': get_model_faces('flame'),
        }
        self.orig_img = orig_img
        self.wireframe = wireframe
        self.renderer = pyrender.OffscreenRenderer(
            viewport_width=self.resolution[0], viewport_height=self.resolution[1], point_size=1.0
        )

        self.vis_ratio = vis_ratio

        # set the scene
        self.scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0], ambient_light=(0.3, 0.3, 0.3))

        light = pyrender.PointLight(color=np.array([1.0, 1.0, 1.0]) * 0.2, intensity=1)

        yrot = np.radians(120)    # angle of lights

        light_pose = np.eye(4)
        light_pose[:3, 3] = [0, -1, 1]
        self.scene.add(light, pose=light_pose)

        light_pose[:3, 3] = [0, 1, 1]
        self.scene.add(light, pose=light_pose)

        light_pose[:3, 3] = [1, 1, 2]
        self.scene.add(light, pose=light_pose)

        spot_l = pyrender.SpotLight(
            color=np.ones(3), intensity=15.0, innerConeAngle=np.pi / 3, outerConeAngle=np.pi / 2
        )

        light_pose[:3, 3] = [1, 2, 2]
        self.scene.add(spot_l, pose=light_pose)

        light_pose[:3, 3] = [-1, 2, 2]
        self.scene.add(spot_l, pose=light_pose)

        # light_pose[:3, 3] = [-2, 2, 0]
        # self.scene.add(spot_l, pose=light_pose)

        # light_pose[:3, 3] = [-2, 2, 0]
        # self.scene.add(spot_l, pose=light_pose)

        self.colors_dict = {
            'red': np.array([0.5, 0.2, 0.2]),
            'pink': np.array([0.7, 0.5, 0.5]),
            'neutral': np.array([0.7, 0.7, 0.6]),
        # 'purple': np.array([0.5, 0.5, 0.7]),
            'purple': np.array([0.55, 0.4, 0.9]),
            'green': np.array([0.5, 0.55, 0.3]),
            'sky': np.array([0.3, 0.5, 0.55]),
            'white': np.array([1.0, 0.98, 0.94]),
        }

    def __call__(
        self,
        verts,
        faces=None,
        img=np.zeros((224, 224, 3)),
        cam=np.array([1, 0, 0]),
        focal_length=[5000, 5000],
        camera_rotation=np.eye(3),
        crop_info=None,
        angle=None,
        axis=None,
        mesh_filename=None,
        color_type=None,
        color=[1.0, 1.0, 0.9],
        iwp_mode=True,
        crop_img=True,
        mesh_type='smpl',
        scale_ratio=1.,
        rgba_mode=False
    ):

        if faces is None:
            faces = self.faces[mesh_type]
        mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False)

        Rx = trimesh.transformations.rotation_matrix(math.radians(180), [1, 0, 0])
        mesh.apply_transform(Rx)

        if mesh_filename is not None:
            mesh.export(mesh_filename)

        if angle and axis:
            R = trimesh.transformations.rotation_matrix(math.radians(angle), axis)
            mesh.apply_transform(R)

        cam = cam.copy()
        if iwp_mode:
            resolution = np.array(img.shape[:2]) * scale_ratio
            if len(cam) == 4:
                sx, sy, tx, ty = cam
                # sy = sx
                camera_translation = np.array([
                    tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9)
                ])
            elif len(cam) == 3:
                sx, tx, ty = cam
                sy = sx
                camera_translation = np.array([
                    -tx, ty, 2 * focal_length[0] / (resolution[0] * sy + 1e-9)
                ])
            render_res = resolution
            self.renderer.viewport_width = render_res[1]
            self.renderer.viewport_height = render_res[0]
        else:
            if crop_info['opt_cam_t'] is None:
                camera_translation = convert_to_full_img_cam(
                    pare_cam=cam[None],
                    bbox_height=crop_info['bbox_scale'] * 200.,
                    bbox_center=crop_info['bbox_center'],
                    img_w=crop_info['img_w'],
                    img_h=crop_info['img_h'],
                    focal_length=focal_length[0],
                )
            else:
                camera_translation = crop_info['opt_cam_t']
            if torch.is_tensor(camera_translation):
                camera_translation = camera_translation[0].cpu().numpy()
            camera_translation = camera_translation.copy()
            camera_translation[0] *= -1
            if 'img_h' in crop_info and 'img_w' in crop_info:
                render_res = (int(crop_info['img_h'][0]), int(crop_info['img_w'][0]))
            else:
                render_res = img.shape[:2] if type(img) is not list else img[0].shape[:2]
            self.renderer.viewport_width = render_res[1]
            self.renderer.viewport_height = render_res[0]
            camera_rotation = camera_rotation.T
        camera = pyrender.IntrinsicsCamera(
            fx=focal_length[0], fy=focal_length[1], cx=render_res[1] / 2., cy=render_res[0] / 2.
        )

        if color_type != None:
            color = self.colors_dict[color_type]

        material = pyrender.MetallicRoughnessMaterial(
            metallicFactor=0.2,
            roughnessFactor=0.6,
            alphaMode='OPAQUE',
            baseColorFactor=(color[0], color[1], color[2], 1.0)
        )

        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)

        mesh_node = self.scene.add(mesh, 'mesh')

        camera_pose = np.eye(4)
        camera_pose[:3, :3] = camera_rotation
        camera_pose[:3, 3] = camera_rotation @ camera_translation
        cam_node = self.scene.add(camera, pose=camera_pose)

        if self.wireframe:
            render_flags = RenderFlags.RGBA | RenderFlags.ALL_WIREFRAME | RenderFlags.SHADOWS_SPOT
        else:
            render_flags = RenderFlags.RGBA | RenderFlags.SHADOWS_SPOT

        rgb, _ = self.renderer.render(self.scene, flags=render_flags)
        if crop_info is not None and crop_img:
            crop_res = img.shape[:2]
            rgb, _, _ = crop(rgb, crop_info['bbox_center'][0], crop_info['bbox_scale'][0], crop_res)

        valid_mask = (rgb[:, :, -1] > 0)[:, :, np.newaxis]

        image_list = [img] if type(img) is not list else img

        return_img = []
        for item in image_list:
            if scale_ratio != 1:
                orig_size = item.shape[:2]
                item = resize(
                    item, (orig_size[0] * scale_ratio, orig_size[1] * scale_ratio),
                    anti_aliasing=True
                )
                item = (item * 255).astype(np.uint8)
            output_img = rgb[:, :, :-1] * valid_mask * self.vis_ratio + (
                1 - valid_mask * self.vis_ratio
            ) * item
            # output_img[valid_mask < 0.5] = item[valid_mask < 0.5]
            # if scale_ratio != 1:
            #     output_img = resize(output_img, (orig_size[0], orig_size[1]), anti_aliasing=True)
            if rgba_mode:
                output_img_rgba = np.zeros((output_img.shape[0], output_img.shape[1], 4))
                output_img_rgba[:, :, :3] = output_img
                output_img_rgba[:, :, 3][valid_mask[:, :, 0]] = 255
                output_img = output_img_rgba.astype(np.uint8)
            image = output_img.astype(np.uint8)
            return_img.append(image)
            return_img.append(item)

        if type(img) is not list:
            # if scale_ratio == 1:
            return_img = return_img[0]

        self.scene.remove_node(mesh_node)
        self.scene.remove_node(cam_node)

        return return_img


class OpenDRenderer:
    def __init__(self, resolution=(224, 224), ratio=1):
        self.resolution = (resolution[0] * ratio, resolution[1] * ratio)
        self.ratio = ratio
        self.focal_length = 5000.
        self.K = np.array([[self.focal_length, 0., self.resolution[1] / 2.],
                           [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.]])
        self.colors_dict = {
            'red': np.array([0.5, 0.2, 0.2]),
            'pink': np.array([0.7, 0.5, 0.5]),
            'neutral': np.array([0.7, 0.7, 0.6]),
            'purple': np.array([0.5, 0.5, 0.7]),
            'green': np.array([0.5, 0.55, 0.3]),
            'sky': np.array([0.3, 0.5, 0.55]),
            'white': np.array([1.0, 0.98, 0.94]),
        }
        self.renderer = ColoredRenderer()
        self.faces = get_smpl_faces()

    def reset_res(self, resolution):
        self.resolution = (resolution[0] * self.ratio, resolution[1] * self.ratio)
        self.K = np.array([[self.focal_length, 0., self.resolution[1] / 2.],
                           [0., self.focal_length, self.resolution[0] / 2.], [0., 0., 1.]])

    def __call__(
        self,
        verts,
        faces=None,
        color=None,
        color_type='white',
        R=None,
        mesh_filename=None,
        img=np.zeros((224, 224, 3)),
        cam=np.array([1, 0, 0]),
        rgba=False,
        addlight=True
    ):
        '''Render mesh using OpenDR
        verts: shape - (V, 3)
        faces: shape - (F, 3)
        img: shape - (224, 224, 3), range - [0, 255] (np.uint8)
        axis: rotate along with X/Y/Z axis (by angle)
        R: rotation matrix (used to manipulate verts) shape - [3, 3]
        Return:
            rendered img: shape - (224, 224, 3), range - [0, 255] (np.uint8)
        '''
        ## Create OpenDR renderer
        rn = self.renderer
        h, w = self.resolution
        K = self.K

        f = np.array([K[0, 0], K[1, 1]])
        c = np.array([K[0, 2], K[1, 2]])

        if faces is None:
            faces = self.faces
        if len(cam) == 4:
            t = np.array([cam[2], cam[3], 2 * K[0, 0] / (w * cam[0] + 1e-9)])
        elif len(cam) == 3:
            t = np.array([cam[1], cam[2], 2 * K[0, 0] / (w * cam[0] + 1e-9)])

        rn.camera = ProjectPoints(rt=np.array([0, 0, 0]), t=t, f=f, c=c, k=np.zeros(5))
        rn.frustum = {'near': 1., 'far': 1000., 'width': w, 'height': h}

        albedo = np.ones_like(verts) * .9

        if color is not None:
            color0 = np.array(color)
            color1 = np.array(color)
            color2 = np.array(color)
        elif color_type == 'white':
            color0 = np.array([1., 1., 1.])
            color1 = np.array([1., 1., 1.])
            color2 = np.array([0.7, 0.7, 0.7])
            color = np.ones_like(verts) * self.colors_dict[color_type][None, :]
        else:
            color0 = self.colors_dict[color_type] * 1.2
            color1 = self.colors_dict[color_type] * 1.2
            color2 = self.colors_dict[color_type] * 1.2
            color = np.ones_like(verts) * self.colors_dict[color_type][None, :]

        # render_smpl = rn.r
        if R is not None:
            assert R.shape == (3, 3), "Shape of rotation matrix should be (3, 3)"
            verts = np.dot(verts, R)

        rn.set(v=verts, f=faces, vc=color, bgcolor=np.zeros(3))

        if addlight:
            yrot = np.radians(120)    # angle of lights
            # # 1. 1. 0.7
            rn.vc = LambertianPointLight(
                f=rn.f,
                v=rn.v,
                num_verts=len(rn.v),
                light_pos=rotateY(np.array([-200, -100, -100]), yrot),
                vc=albedo,
                light_color=color0
            )

            # Construct Left Light
            rn.vc += LambertianPointLight(
                f=rn.f,
                v=rn.v,
                num_verts=len(rn.v),
                light_pos=rotateY(np.array([800, 10, 300]), yrot),
                vc=albedo,
                light_color=color1
            )

            # Construct Right Light
            rn.vc += LambertianPointLight(
                f=rn.f,
                v=rn.v,
                num_verts=len(rn.v),
                light_pos=rotateY(np.array([-500, 500, 1000]), yrot),
                vc=albedo,
                light_color=color2
            )

        rendered_image = rn.r
        visibility_image = rn.visibility_image

        image_list = [img] if type(img) is not list else img

        return_img = []
        for item in image_list:
            if self.ratio != 1:
                img_resized = resize(
                    item, (item.shape[0] * self.ratio, item.shape[1] * self.ratio),
                    anti_aliasing=True
                )
            else:
                img_resized = item / 255.

            try:
                img_resized[visibility_image != (2**32 - 1)
                           ] = rendered_image[visibility_image != (2**32 - 1)]
            except:
                logger.warning('Can not render mesh.')

            img_resized = (img_resized * 255).astype(np.uint8)
            res = img_resized

            if rgba:
                img_resized_rgba = np.zeros((img_resized.shape[0], img_resized.shape[1], 4))
                img_resized_rgba[:, :, :3] = img_resized
                img_resized_rgba[:, :, 3][visibility_image != (2**32 - 1)] = 255
                res = img_resized_rgba.astype(np.uint8)
            return_img.append(res)

        if type(img) is not list:
            return_img = return_img[0]

        return return_img


#  https://github.com/classner/up/blob/master/up_tools/camera.py
def rotateY(points, angle):
    """Rotate all points in a 2D array around the y axis."""
    ry = np.array([[np.cos(angle), 0., np.sin(angle)], [0., 1., 0.],
                   [-np.sin(angle), 0., np.cos(angle)]])
    return np.dot(points, ry)


def rotateX(points, angle):
    """Rotate all points in a 2D array around the x axis."""
    rx = np.array([[1., 0., 0.], [0., np.cos(angle), -np.sin(angle)],
                   [0., np.sin(angle), np.cos(angle)]])
    return np.dot(points, rx)


def rotateZ(points, angle):
    """Rotate all points in a 2D array around the z axis."""
    rz = np.array([[np.cos(angle), -np.sin(angle), 0.], [np.sin(angle),
                                                         np.cos(angle), 0.], [0., 0., 1.]])
    return np.dot(points, rz)


class IUV_Renderer(object):
    def __init__(
        self,
        focal_length=5000.,
        orig_size=224,
        output_size=56,
        mode='iuv',
        device=torch.device('cuda'),
        mesh_type='smpl'
    ):

        self.focal_length = focal_length
        self.orig_size = orig_size
        self.output_size = output_size

        if mode in ['iuv']:
            if mesh_type == 'smpl':
                DP = DensePoseMethods()

                vert_mapping = DP.All_vertices.astype('int64') - 1
                self.vert_mapping = torch.from_numpy(vert_mapping)

                faces = DP.FacesDensePose
                faces = faces[None, :, :]
                self.faces = torch.from_numpy(
                    faces.astype(np.int32)
                )    # [1, 13774, 3], torch.int32

                num_part = float(np.max(DP.FaceIndices))
                self.num_part = num_part

                dp_vert_pid_fname = 'data/dp_vert_pid.npy'
                if os.path.exists(dp_vert_pid_fname):
                    dp_vert_pid = list(np.load(dp_vert_pid_fname))
                else:
                    print('creating data/dp_vert_pid.npy')
                    dp_vert_pid = []
                    for v in range(len(vert_mapping)):
                        for i, f in enumerate(DP.FacesDensePose):
                            if v in f:
                                dp_vert_pid.append(DP.FaceIndices[i])
                                break
                    np.save(dp_vert_pid_fname, np.array(dp_vert_pid))

                textures_vts = np.array([(dp_vert_pid[i] / num_part, DP.U_norm[i], DP.V_norm[i])
                                         for i in range(len(vert_mapping))])
                self.textures_vts = torch.from_numpy(
                    textures_vts[None].astype(np.float32)
                )    # (1, 7829, 3)
        elif mode == 'pncc':
            self.vert_mapping = None
            self.faces = torch.from_numpy(
                get_model_faces(mesh_type)[None].astype(np.int32)
            )    #  mano: torch.Size([1, 1538, 3])
            textures_vts = get_model_tpose(mesh_type).unsqueeze(
                0
            )    # mano: torch.Size([1, 778, 3])

            texture_min = torch.min(textures_vts) - 0.001
            texture_range = torch.max(textures_vts) - texture_min + 0.001
            self.textures_vts = (textures_vts - texture_min) / texture_range
        elif mode in ['seg']:
            self.vert_mapping = None
            body_model = 'smpl'

            self.faces = torch.from_numpy(get_smpl_faces().astype(np.int32)[None])

            with open(
                os.path.join(
                    path_config.SMPL_MODEL_DIR, '{}_vert_segmentation.json'.format(body_model)
                ), 'rb'
            ) as json_file:
                smpl_part_id = json.load(json_file)

            v_id = []
            for k in smpl_part_id.keys():
                v_id.extend(smpl_part_id[k])

            v_id = torch.tensor(v_id)
            n_verts = len(torch.unique(v_id))
            num_part = len(constants.SMPL_PART_ID.keys())
            self.num_part = num_part

            seg_vert_pid = np.zeros(n_verts)
            for k in smpl_part_id.keys():
                seg_vert_pid[smpl_part_id[k]] = constants.SMPL_PART_ID[k]

            print('seg_vert_pid', seg_vert_pid.shape)
            textures_vts = seg_vert_pid[:, None].repeat(3, axis=1) / num_part
            print('textures_vts', textures_vts.shape)
            # textures_vts = np.array(
            #     [(seg_vert_pid[i] / num_part,) * 3 for i in
            #     range(n_verts)])
            self.textures_vts = torch.from_numpy(textures_vts[None].astype(np.float32))

        K = np.array([[self.focal_length, 0., self.orig_size / 2.],
                      [0., self.focal_length, self.orig_size / 2.], [0., 0., 1.]])

        R = np.array([[-1., 0., 0.], [0., -1., 0.], [0., 0., 1.]])

        t = np.array([0, 0, 5])

        if self.orig_size != 224:
            rander_scale = self.orig_size / float(224)
            K[0, 0] *= rander_scale
            K[1, 1] *= rander_scale
            K[0, 2] *= rander_scale
            K[1, 2] *= rander_scale

        self.K = torch.FloatTensor(K[None, :, :])
        self.R = torch.FloatTensor(R[None, :, :])
        self.t = torch.FloatTensor(t[None, None, :])

        camK = F.pad(self.K, (0, 1, 0, 1), "constant", 0)
        camK[:, 2, 2] = 0
        camK[:, 3, 2] = 1
        camK[:, 2, 3] = 1

        self.K = camK

        self.device = device
        lights = AmbientLights(device=self.device)

        raster_settings = RasterizationSettings(
            image_size=output_size,
            blur_radius=0,
            faces_per_pixel=1,
        )
        self.renderer = MeshRenderer(
            rasterizer=MeshRasterizer(raster_settings=raster_settings),
            shader=HardFlatShader(
                device=self.device,
                lights=lights,
                blend_params=BlendParams(background_color=[0, 0, 0], sigma=0.0, gamma=0.0)
            )
        )

    def camera_matrix(self, cam):
        batch_size = cam.size(0)

        K = self.K.repeat(batch_size, 1, 1)
        R = self.R.repeat(batch_size, 1, 1)
        t = torch.stack([
            -cam[:, 1], -cam[:, 2], 2 * self.focal_length / (self.orig_size * cam[:, 0] + 1e-9)
        ],
                        dim=-1)

        if cam.is_cuda:
            # device_id = cam.get_device()
            K = K.to(cam.device)
            R = R.to(cam.device)
            t = t.to(cam.device)

        return K, R, t

    def verts2iuvimg(self, verts, cam, iwp_mode=True):
        batch_size = verts.size(0)

        K, R, t = self.camera_matrix(cam)

        if self.vert_mapping is None:
            vertices = verts
        else:
            vertices = verts[:, self.vert_mapping, :]

        mesh = Meshes(vertices, self.faces.to(verts.device).expand(batch_size, -1, -1))
        mesh.textures = TexturesVertex(
            verts_features=self.textures_vts.to(verts.device).expand(batch_size, -1, -1)
        )

        cameras = PerspectiveCameras(
            device=verts.device,
            R=R,
            T=t,
            K=K,
            in_ndc=False,
            image_size=[(self.orig_size, self.orig_size)]
        )

        iuv_image = self.renderer(mesh, cameras=cameras)
        iuv_image = iuv_image[..., :3].permute(0, 3, 1, 2)

        return iuv_image