skyreels_a1/src/renderer.py

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch3d.structures import Meshes
from pytorch3d.io import load_obj
from pytorch3d.renderer.mesh import rasterize_meshes
import pickle
import chumpy as ch
import cv2
import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from skyreels_a1.src.utils.mediapipe_utils import face_vertices, vertex_normals, batch_orth_proj
from skyreels_a1.src.media_pipe.draw_util import FaceMeshVisualizer
from mediapipe.framework.formats import landmark_pb2

def keep_vertices_and_update_faces(faces, vertices_to_keep):
    """
    Keep specified vertices in the mesh and update the faces.
    """
    if isinstance(vertices_to_keep, list) or isinstance(vertices_to_keep, np.ndarray):
        vertices_to_keep = torch.tensor(vertices_to_keep, dtype=torch.long)

    vertices_to_keep = torch.unique(vertices_to_keep)
    max_vertex_index = faces.max().long().item() + 1

    mask = torch.zeros(max_vertex_index, dtype=torch.bool)
    mask[vertices_to_keep] = True

    new_vertex_indices = torch.full((max_vertex_index,), -1, dtype=torch.long)
    new_vertex_indices[mask] = torch.arange(len(vertices_to_keep))

    valid_faces_mask = (new_vertex_indices[faces] != -1).all(dim=1)
    filtered_faces = faces[valid_faces_mask]
    updated_faces = new_vertex_indices[filtered_faces]

    return updated_faces

def predict_landmark_position(ref_points, relative_coords):
    """
    Predict the new position of the eyeball based on reference points and relative coordinates.
    """
    left_corner = ref_points[0]
    right_corner = ref_points[8]
    
    eye_center = (left_corner + right_corner) / 2
    eye_width_vector = right_corner - left_corner
    eye_width = np.linalg.norm(eye_width_vector)
    eye_direction = eye_width_vector / eye_width
    eye_vertical = np.array([-eye_direction[1], eye_direction[0]])
    
    predicted_pos = eye_center + \
                   (eye_width/2) * relative_coords[0] * eye_direction + \
                   (eye_width/2) * relative_coords[1] * eye_vertical
    
    return predicted_pos

def mesh_points_by_barycentric_coordinates(mesh_verts, mesh_faces, lmk_face_idx, lmk_b_coords):
    """
    Evaluation 3d points given mesh and landmark embedding
    """
    dif1 = ch.vstack([
        (mesh_verts[mesh_faces[lmk_face_idx], 0] * lmk_b_coords).sum(axis=1),
        (mesh_verts[mesh_faces[lmk_face_idx], 1] * lmk_b_coords).sum(axis=1),
        (mesh_verts[mesh_faces[lmk_face_idx], 2] * lmk_b_coords).sum(axis=1)
    ]).T
    return dif1

class Renderer(nn.Module):
    def __init__(self, render_full_head=False, obj_filename='pretrained_models/FLAME/head_template.obj'):
        super(Renderer, self).__init__()
        self.image_size = 224
        self.mediapipe_landmark_embedding = np.load("pretrained_models/smirk/mediapipe_landmark_embedding.npz")
        self.vis = FaceMeshVisualizer(forehead_edge=False)
        
        verts, faces, aux = load_obj(obj_filename)
        uvcoords = aux.verts_uvs[None, ...]      # (N, V, 2)
        uvfaces = faces.textures_idx[None, ...]   # (N, F, 3)
        faces = faces.verts_idx[None,...]

        self.render_full_head = render_full_head

        red_color = torch.tensor([255, 0, 0])[None, None, :].float() / 255.
        transparent_color = torch.tensor([0, 0, 0])[None, None, :].float()
        colors = transparent_color.repeat(1, 5023, 1)

        flame_masks = pickle.load(
            open('pretrained_models/FLAME/FLAME_masks.pkl', 'rb'),
            encoding='latin1')
        self.flame_masks = flame_masks

        self.register_buffer('faces', faces)

        face_colors = face_vertices(colors, faces)
        self.register_buffer('face_colors', face_colors)
        
        self.register_buffer('raw_uvcoords', uvcoords)

        uvcoords = torch.cat([uvcoords, uvcoords[:,:,0:1]*0.+1.], -1) #[bz, ntv, 3]
        uvcoords = uvcoords*2 - 1; uvcoords[...,1] = -uvcoords[...,1]
        face_uvcoords = face_vertices(uvcoords, uvfaces)
        self.register_buffer('uvcoords', uvcoords)
        self.register_buffer('uvfaces', uvfaces)
        self.register_buffer('face_uvcoords', face_uvcoords)

        pi = np.pi
        constant_factor = torch.tensor([1/np.sqrt(4*pi), ((2*pi)/3)*(np.sqrt(3/(4*pi))), ((2*pi)/3)*(np.sqrt(3/(4*pi))),\
                           ((2*pi)/3)*(np.sqrt(3/(4*pi))), (pi/4)*(3)*(np.sqrt(5/(12*pi))), (pi/4)*(3)*(np.sqrt(5/(12*pi))),\
                           (pi/4)*(3)*(np.sqrt(5/(12*pi))), (pi/4)*(3/2)*(np.sqrt(5/(12*pi))), (pi/4)*(1/2)*(np.sqrt(5/(4*pi)))]).float()
        self.register_buffer('constant_factor', constant_factor)

    def forward(self, vertices, cam_params, source_tform=None, tform_512=None, weights_468=None, weights_473=None,shape = None, **landmarks):
        transformed_vertices = batch_orth_proj(vertices, cam_params)
        transformed_vertices[:, :, 1:] = -transformed_vertices[:, :, 1:]

        transformed_landmarks = {}
        for key in landmarks.keys():
            transformed_landmarks[key] = batch_orth_proj(landmarks[key], cam_params)
            transformed_landmarks[key][:, :, 1:] = - transformed_landmarks[key][:, :, 1:]
            transformed_landmarks[key] = transformed_landmarks[key][...,:2]

        # rendered_img = self.render(vertices, transformed_vertices, source_tform, tform_512, weights_468, weights_473,shape)
        if weights_468 is None:
            rendered_img = self.render_with_pulid_in_vertices(vertices, transformed_vertices, source_tform, tform_512, shape)
        else:
            rendered_img = self.render(vertices, transformed_vertices, source_tform, tform_512, weights_468, weights_473,shape)

        outputs = {
            'rendered_img': rendered_img,
            'transformed_vertices': transformed_vertices
        }
        outputs.update(transformed_landmarks)

        return outputs
    
    def _calculate_eye_landmarks(self, landmark_list_pixlevel, weights_468, weights_473, source_tform):
        #  [np.array([x_relative, y_relative]),target_point,ref_points] 根据当前的new_landmarks，根据target_point， 利用映射变换，计算出眼部landmarks
       
       import pdb; pdb.set_trace()
       pass

    def render_with_pulid_in_vertices(self, vertices, transformed_vertices, source_tform, tform_512, shape):
        batch_size = vertices.shape[0]
        ## rasterizer near 0 far 100. move mesh so minz larger than 0
        transformed_vertices[:,:,2] = transformed_vertices[:,:,2] + 10
        # import pdb;pdb.set_trace()
        # 只使用颜色作为attributes
        colors = self.face_colors.expand(batch_size, -1, -1, -1)

        # # 加载 mediapipe_landmark_embedding 数据
        lmk_b_coords = self.mediapipe_landmark_embedding['lmk_b_coords']
        lmk_face_idx = self.mediapipe_landmark_embedding['lmk_face_idx']
        # import pdb;pdb.set_trace()
        # 计算 v_selected
        v_selected = mesh_points_by_barycentric_coordinates(transformed_vertices.detach().cpu().numpy()[0], self.faces.detach().cpu().numpy()[0], lmk_face_idx, lmk_b_coords)
        # v_selected 增加对应左眼和右眼的8个位置，序号分别是：[4051, 3997, 3965, 3933, 4020]，[4597, 4543, 4511, 4479, 4575]，得根据transformed_vertices.detach().cpu().numpy()[0]来获取
        v_selected = np.concatenate([v_selected, transformed_vertices.detach().cpu().numpy()[0][[4543, 4511, 4479, 4575]], transformed_vertices.detach().cpu().numpy()[0][[3997, 3965, 3933, 4020]]], axis=0)
        
        v_selected_tensor = torch.tensor( np.array(v_selected), dtype=torch.float32).to(transformed_vertices.device)
        new_landmarks = landmark_pb2.NormalizedLandmarkList()
        for v in v_selected_tensor:
            # 将 v 映射到图像坐标
            img_x = (v[0] + 1) * 0.5 * self.image_size
            img_y = ((v[1] + 1) * 0.5) * self.image_size
            # import pdb;pdb.set_trace()
            point = np.array([img_x.cpu().numpy(), img_y.cpu().numpy(), 1.0])
            croped_point = np.dot(source_tform.inverse.params, point)
            
            # original_point = np.dot(tform_512.inverse.params, point)
            landmark = new_landmarks.landmark.add()
            landmark.x = croped_point[0]/shape[1]
            landmark.y = croped_point[1]/shape[0]
            landmark.z = 1.0
        # 将 v 映射到图像坐标
        right_eye_x = (transformed_vertices[0,4597,0] + 1) * 0.5 * self.image_size
        right_eye_y = (transformed_vertices[0,4597,1] + 1) * 0.5 * self.image_size
        right_eye_point = np.array([right_eye_x.cpu().numpy(), right_eye_y.cpu().numpy(), 1.0])
        right_eye_original = np.dot(source_tform.inverse.params, right_eye_point)
        right_eye_landmarks = right_eye_original[:2]

        left_eye_x = (transformed_vertices[0,4051,0] + 1) * 0.5 * self.image_size
        left_eye_y = (transformed_vertices[0,4051,1] + 1) * 0.5 * self.image_size
        left_eye_point = np.array([left_eye_x.cpu().numpy(), left_eye_y.cpu().numpy(), 1.0])
        left_eye_original = np.dot(source_tform.inverse.params, left_eye_point)
        left_eye_landmarks = left_eye_original[:2]

        image_new = np.zeros([shape[0],shape[1],3], dtype=np.uint8)
        self.vis.mp_drawing.draw_landmarks(image=image_new,landmark_list=new_landmarks,connections=self.vis.face_connection_spec.keys(),landmark_drawing_spec=None,connection_drawing_spec=self.vis.face_connection_spec)
        
        # 直接设置单个像素点的颜色
        left_point = (int(left_eye_landmarks[0]), int(left_eye_landmarks[1]))
        right_point = (int(right_eye_landmarks[0]), int(right_eye_landmarks[1]))
        # import pdb;pdb.set_trace()
        # 左眼点 - 3x3 区域
        image_new[left_point[1]-1:left_point[1]+2, left_point[0]-1:left_point[0]+2] = [180, 200, 10]  # RGB格式
        # 右眼点 - 3x3 区域 
        image_new[right_point[1]-1:right_point[1]+2, right_point[0]-1:right_point[0]+2] = [10, 200, 180]

        landmark_58 = new_landmarks.landmark[57]  # 因为索引从0开始，所以57表示第58个点
        x = int(landmark_58.x * shape[1])
        y = int(landmark_58.y * shape[0])
        image_new[y-2:y+3, x-2:x+3] = [255, 255, 255]  # 设置3x3的白色区域

        return np.copy(image_new)

    def render(self, vertices, transformed_vertices, source_tform, tform_512, weights_468, weights_473, shape):
        # batch_size = vertices.shape[0]
        transformed_vertices[:,:,2] += 10  # Z-axis offset
        
        # colors = self.face_colors.expand(batch_size, -1, -1, -1)
        # rendering = self.rasterize(transformed_vertices, self.faces.expand(batch_size, -1, -1), colors)
        
        v_selected = self._calculate_landmark_points(transformed_vertices)
        v_selected_tensor = torch.tensor(v_selected, dtype=torch.float32, device=transformed_vertices.device) #torch.Size([113, 3])
        # import pdb; pdb.set_trace()
        new_landmarks, landmark_list_pixlevel = self._create_landmark_list(v_selected_tensor, source_tform, shape)
        # 基于weights_468和weights_473，计算眼部landmarks
        left_eye_point_indices = weights_468[3]
        right_eye_point_indices = weights_473[3]
        # 遍历每个索引以找到其在 index_mapping 中的位置
        left_eye_point_indices = [self.vis.index_mapping.index(idx) for idx in left_eye_point_indices]
        right_eye_point_indices = [self.vis.index_mapping.index(idx) for idx in right_eye_point_indices]

        left_eye_point = [landmark_list_pixlevel[idx] for idx in left_eye_point_indices]
        right_eye_point = [landmark_list_pixlevel[idx] for idx in right_eye_point_indices]
        # import pdb; pdb.set_trace()
        # weights_468[2].shape = (16, 2)    
        M_affine_left, _ = cv2.estimateAffine2D(np.array(weights_468[2], dtype=np.float32), np.array(left_eye_point, dtype=np.float32))
        M_affine_right, _ = cv2.estimateAffine2D(np.array(weights_473[2], dtype=np.float32), np.array(right_eye_point, dtype=np.float32))

        # 计算瞳孔点
        pupil_left_eye = cv2.transform(weights_468[1].reshape(1, 1, 2), M_affine_left).reshape(-1)
        pupil_right_eye = cv2.transform(weights_473[1].reshape(1, 1, 2), M_affine_right).reshape(-1)

        # left_eye_point, right_eye_point = self._calculate_eye_landmarks(landmark_list_pixlevel, weights_468, weights_473, source_tform)
        # left_eye_point, right_eye_point = self._process_eye_landmarks(transformed_vertices, source_tform)
        # import pdb; pdb.set_trace()
        return self._generate_final_image(new_landmarks, pupil_left_eye, pupil_right_eye, shape)
        # return self._generate_final_image(new_landmarks, left_eye_point, right_eye_point, shape)

    def _calculate_landmark_points(self, transformed_vertices):
        lmk_b_coords = self.mediapipe_landmark_embedding['lmk_b_coords']
        lmk_face_idx = self.mediapipe_landmark_embedding['lmk_face_idx']
        
        base_points = mesh_points_by_barycentric_coordinates(
            transformed_vertices.detach().cpu().numpy()[0],
            self.faces.detach().cpu().numpy()[0],
            lmk_face_idx, lmk_b_coords
        )
        
        RIGHT_EYE_INDICES = [4543, 4511, 4479, 4575]
        LEFT_EYE_INDICES = [3997, 3965, 3933, 4020]
        return np.concatenate([
            base_points,
            transformed_vertices.detach().cpu().numpy()[0][RIGHT_EYE_INDICES],
            transformed_vertices.detach().cpu().numpy()[0][LEFT_EYE_INDICES]
        ], axis=0)

    def _create_landmark_list(self, vertices, transform, shape):
        landmark_list = landmark_pb2.NormalizedLandmarkList()
        landmark_list_pixlevel = []
        for v in vertices:
            img_x = (v[0] + 1) * 0.5 * self.image_size
            img_y = (v[1] + 1) * 0.5 * self.image_size
            projected = np.dot(transform.inverse.params, [img_x.cpu().numpy(), img_y.cpu().numpy(), 1.0])
            landmark_list_pixlevel.append((projected[0], projected[1]))
            landmark = landmark_list.landmark.add()
            landmark.x = projected[0] / shape[1]
            landmark.y = projected[1] / shape[0]
            landmark.z = 1.0
        return landmark_list, landmark_list_pixlevel

    def _process_eye_landmarks(self, vertices, transform):
        def project_eye_point(vertex_idx):
            x = (vertices[0, vertex_idx, 0] + 1) * 0.5 * self.image_size
            y = (vertices[0, vertex_idx, 1] + 1) * 0.5 * self.image_size
            # import pdb; pdb.set_trace()
            projected = np.dot(transform.inverse.params, [x.cpu().numpy(), y.cpu().numpy(), 1.0])
            return (int(projected[0]), int(projected[1]))

        return (
            project_eye_point(4051),  # Left eye index
            project_eye_point(4597)   # Right eye index
        )

    def _generate_final_image(self, landmarks, left_eye, right_eye, shape):
        image = np.zeros([shape[0], shape[1], 3], dtype=np.uint8)
        
        self.vis.mp_drawing.draw_landmarks(
            image=image,
            landmark_list=landmarks,
            connections=self.vis.face_connection_spec.keys(),
            landmark_drawing_spec=None,
            connection_drawing_spec=self.vis.face_connection_spec
        )
        
        self._draw_eye_markers(image, np.array(left_eye, dtype=np.int32), np.array(right_eye, dtype=np.int32))
        self._draw_landmark_58(image, landmarks, shape)
        return np.copy(image)

    def _draw_eye_markers(self, image, left_eye, right_eye):
        y, x = left_eye[1]-1, left_eye[0]-1
        image[y:y+3, x:x+3] = [10, 200, 250]
        
        y, x = right_eye[1]-1, right_eye[0]-1
        image[y:y+3, x:x+3] = [250, 200, 10]

    def _draw_landmark_58(self, image, landmarks, shape):
        if len(landmarks.landmark) > 57:
            point = landmarks.landmark[57]
            x = int(point.x * shape[1])
            y = int(point.y * shape[0])
            image[y-2:y+3, x-2:x+3] = [255, 255, 255]

    def rasterize(self, vertices, faces, attributes=None, h=None, w=None):
        fixed_vertices = vertices.clone()
        fixed_vertices[...,:2] = -fixed_vertices[...,:2]

        if h is None and w is None:
            image_size = self.image_size
        else:
            image_size = [h, w]
            if h>w:
                fixed_vertices[..., 1] = fixed_vertices[..., 1]*h/w
            else:
                fixed_vertices[..., 0] = fixed_vertices[..., 0]*w/h
        meshes_screen = Meshes(verts=fixed_vertices.float(), faces=faces.long())
        pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
            meshes_screen,
            image_size=image_size,
            blur_radius=0.0,
            faces_per_pixel=1,
            bin_size=None,
            max_faces_per_bin=None,
            perspective_correct=False,
        )
        vismask = (pix_to_face > -1).float()
        D = attributes.shape[-1]
        attributes = attributes.clone(); attributes = attributes.view(attributes.shape[0]*attributes.shape[1], 3, attributes.shape[-1])
        N, H, W, K, _ = bary_coords.shape
        mask = pix_to_face == -1
        pix_to_face = pix_to_face.clone()
        pix_to_face[mask] = 0
        idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
        pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D)
        pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2)
        pixel_vals[mask] = 0  # Replace masked values in output.
        pixel_vals = pixel_vals[:,:,:,0].permute(0,3,1,2)
        pixel_vals = torch.cat([pixel_vals, vismask[:,:,:,0][:,None,:,:]], dim=1)
        return pixel_vals

    def add_SHlight(self, normal_images, sh_coeff):
        '''
            sh_coeff: [bz, 9, 3]
        '''
        N = normal_images
        sh = torch.stack([
                N[:,0]*0.+1., N[:,0], N[:,1], \
                N[:,2], N[:,0]*N[:,1], N[:,0]*N[:,2], 
                N[:,1]*N[:,2], N[:,0]**2 - N[:,1]**2, 3*(N[:,2]**2) - 1
                ], 
                1) # [bz, 9, h, w]
        sh = sh*self.constant_factor[None,:,None,None]
        shading = torch.sum(sh_coeff[:,:,:,None,None]*sh[:,:,None,:,:], 1) # [bz, 9, 3, h, w]  
        return shading

    def add_pointlight(self, vertices, normals, lights):
        '''
            vertices: [bz, nv, 3]
            lights: [bz, nlight, 6]
        returns:
            shading: [bz, nv, 3]
        '''
        light_positions = lights[:,:,:3]; light_intensities = lights[:,:,3:]
        directions_to_lights = F.normalize(light_positions[:,:,None,:] - vertices[:,None,:,:], dim=3)
        normals_dot_lights = (normals[:,None,:,:]*directions_to_lights).sum(dim=3)
        shading = normals_dot_lights[:,:,:,None]*light_intensities[:,:,None,:]
        return shading.mean(1)
    

    def add_directionlight(self, normals, lights):
        '''
            normals: [bz, nv, 3]
            lights: [bz, nlight, 6]
        returns:
            shading: [bz, nv, 3]
        '''
        light_direction = lights[:,:,:3]; light_intensities = lights[:,:,3:]
        directions_to_lights = F.normalize(light_direction[:,:,None,:].expand(-1,-1,normals.shape[1],-1), dim=3)
        normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.)
        shading = normals_dot_lights[:,:,:,None]*light_intensities[:,:,None,:]
        return shading.mean(1)



    def render_multiface(self, vertices, transformed_vertices, faces):
        
        batch_size = vertices.shape[0]

        light_positions = torch.tensor(
            [
            [-1,-1,-1],
            [1,-1,-1],
            [-1,+1,-1],
            [1,+1,-1],
            [0,0,-1]
            ]
        )[None,:,:].expand(batch_size, -1, -1).float()

        light_intensities = torch.ones_like(light_positions).float()*1.7
        lights = torch.cat((light_positions, light_intensities), 2).to(vertices.device)
        
        transformed_vertices[:,:,2] = transformed_vertices[:,:,2] + 10
        normals = vertex_normals(vertices, faces) 
        face_normals = face_vertices(normals, faces)
        
        colors = torch.tensor([180, 180, 180])[None, None, :].repeat(1, transformed_vertices.shape[1]+1, 1).float()/255.
        colors = colors.cuda()
        face_colors = face_vertices(colors, faces[0].unsqueeze(0))
        
        colors = face_colors.expand(batch_size, -1, -1, -1)

        attributes = torch.cat([colors,
                                face_normals], 
                                -1)
        rendering = self.rasterize(transformed_vertices, faces, attributes)
        
        albedo_images = rendering[:, :3, :, :]

        normal_images = rendering[:, 3:6, :, :]

        shading = self.add_directionlight(normal_images.permute(0,2,3,1).reshape([batch_size, -1, 3]), lights)
        shading_images = shading.reshape([batch_size, albedo_images.shape[2], albedo_images.shape[3], 3]).permute(0,3,1,2).contiguous()        
        shaded_images = albedo_images*shading_images
        
        return shaded_images