import torch import math import numpy as np from pytorch3d.structures import ( Meshes, Pointclouds ) from pytorch3d.renderer import ( PerspectiveCameras,CamerasBase, look_at_view_transform, RasterizationSettings, PointLights, TexturesVertex, BlendParams,TexturesUV, SoftPhongShader, MeshRasterizer, MeshRenderer, ) from pytorch3d.transforms.transform3d import _broadcast_bmm from pytorch3d.renderer.mesh.rasterizer import Fragments,rasterize_meshes import torch import matplotlib.pyplot as plt import trimesh def show_image(image_tensor, name): # 转为 [H, W, C] 并归一化 image_np = image_tensor.permute(1, 2, 0).cpu().numpy() # 保存为 PNG plt.imsave(f"test_images/{name}.png", image_np) def get_view_matrix(R, t): device=R.device Rt = torch.cat((R, t.view(3,1)),1) b_row=torch.tensor([0, 0, 0, 1], dtype=torch.float32, device=device).view(1, 4) #torch.FloatTensor([0,0,0,1],device=device).view(1,4) view_matrix = torch.cat((Rt, b_row)) return view_matrix def get_batch_view_matrix(R, t): """ R: [B, 3, 3] - batch of rotation matrices t: [B, 3] or [B, 1, 3] - batch of translation vectors Returns: view_matrix: [B, 4, 4] """ device = R.device B = R.shape[0] # Ensure t is shape [B, 3, 1] if t.ndim == 2: t = t.unsqueeze(-1) # [B, 3, 1] # Concatenate R and t → [B, 3, 4] Rt = torch.cat((R, t), dim=2) # Create [0, 0, 0, 1] row for each batch → [B, 1, 4] b_row = torch.tensor([0, 0, 0, 1], dtype=torch.float32, device=device).view(1, 1, 4).expand(B, 1, 4) # Concatenate to form full view matrix → [B, 4, 4] view_matrix = torch.cat((Rt, b_row), dim=1) return view_matrix def get_proj_matrix( tanfov,device, z_near=0.01, z_far=100, z_sign=1.0,): tanHalfFovY = tanfov tanHalfFovX = tanfov top = tanHalfFovY * z_near bottom = -top right = tanHalfFovX * z_near left = -right z_sign = 1.0 proj_matrix = torch.zeros(4, 4).float().to(device) proj_matrix[0, 0] = 2.0 * z_near / (right - left) proj_matrix[1, 1] = 2.0 * z_near / (top - bottom) proj_matrix[0, 2] = (right + left) / (right - left) proj_matrix[1, 2] = (top + bottom) / (top - bottom) proj_matrix[3, 2] = z_sign proj_matrix[2, 2] = z_sign * z_far / (z_far - z_near) proj_matrix[2, 3] = -(z_far * z_near) / (z_far - z_near) return proj_matrix def get_full_proj_matrix(w2c_cam,tanfov): assert len(w2c_cam.shape)==2 view_matrix=get_view_matrix(w2c_cam[:3,:3],w2c_cam[:3,3]).transpose(0,1).contiguous() proj_matrix=get_proj_matrix(tanfov,device=w2c_cam.device,z_near=0.01, z_far=100, z_sign=1.0).transpose(0,1).contiguous() full_proj_matrix = (view_matrix.unsqueeze(0).bmm(proj_matrix.unsqueeze(0))).squeeze(0)#torch.mm(view_matrix, proj_matrix) return view_matrix,full_proj_matrix def get_batch_full_proj_matrix(w2c_cam, tanfov): """ Args: w2c_cam: [B, 4, 4] - world-to-camera matrix tanfov: float or [B] - tangent of half FOV Returns: view_matrix: [B, 4, 4] full_proj_matrix: [B, 4, 4] """ device = w2c_cam.device B = w2c_cam.shape[0] # Get view matrix from R and t R = w2c_cam[:, :3, :3] t = w2c_cam[:, :3, 3] view_matrix = get_batch_view_matrix(R, t) # [B, 4, 4] # Get projection matrix per batch proj_matrix = get_proj_matrix( tanfov, device=device, z_near=0.01, z_far=100.0, z_sign=1.0 ) # Should return [B, 4, 4] # Ensure shape: [B, 4, 4] if proj_matrix.ndim == 2: # single matrix proj_matrix = proj_matrix.unsqueeze(0).expand(B, -1, -1) # Matrix multiplication in batch full_proj_matrix = proj_matrix.bmm(view_matrix) return view_matrix, full_proj_matrix def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return torch.sum(x*y, -1, keepdim=True) def length(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor: return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN def safe_normalize(x: torch.Tensor, eps: float =1e-20) -> torch.Tensor: return x / length(x, eps) # 计算三角面片的局部坐标系（方向/姿态），即每个三角面片的三个正交轴向量 def compute_face_orientation(verts, faces, return_scale=False): i0 = faces[..., 0].long() i1 = faces[..., 1].long() i2 = faces[..., 2].long() v0 = verts[..., i0, :] v1 = verts[..., i1, :] v2 = verts[..., i2, :] a0 = safe_normalize(v1 - v0) a1 = safe_normalize(torch.cross(a0, v2 - v0, dim=-1)) a2 = -safe_normalize(torch.cross(a1, a0, dim=-1)) orientation = torch.cat([a0[..., None], a1[..., None], a2[..., None]], dim=-1) if return_scale: s0 = length(v1 - v0) s1 = dot(a2, (v2 - v0)).abs() scale = (s0 + s1) / 2 return orientation, scale class VertexPositionShader(torch.nn.Module): def __init__(self): super().__init__() def forward(self, fragments, meshes, **kwargs): """ :param fragments: Fragments of the meshes that are rasterized. :param meshes: Meshes to render. :param kwargs: Additional arguments passed by the renderer. :return: The output colors, which in this case will be the vertex positions. """ pixel_positions = fragments.pix_to_face # shape (num_pixels, 3) batch_size, H, W = pixel_positions.shape[0], pixel_positions.shape[1], pixel_positions.shape[2] bary_coords=fragments.bary_coords.squeeze(-2) alpha = (pixel_positions!=-1)*1.0 vertex_faces = meshes.faces_packed()[pixel_positions.squeeze(-1)]# # shape (num_pixels, 3) vertex_positions=(meshes.verts_packed()[vertex_faces]*bary_coords[...,None]).sum(dim=-2)#.mean(dim=-2) results=torch.cat([vertex_positions,alpha],dim=-1) extra_result={"vertex_faces":vertex_faces,"bary_coords":bary_coords} return results,extra_result class GS_Camera(CamerasBase): #still obey pytorch 3d coordinate system, adapting to gaussian splatting's projection method def __init__( self, focal_length=1.0, R: torch.Tensor = torch.eye(3)[None], T: torch.Tensor = torch.zeros(1, 3), principal_point=((0.0, 0.0),),#useless device = "cpu", in_ndc: bool = True, image_size = None, ) -> None: kwargs = {"image_size": image_size} if image_size is not None else {} super().__init__( device=device, focal_length=focal_length, principal_point=((0.0, 0.0),), R=R, T=T, K=None, _in_ndc=in_ndc, **kwargs, # pyre-ignore ) if image_size is not None: if (self.image_size < 1).any(): # pyre-ignore raise ValueError("Image_size provided has invalid values") else: self.image_size = None # When focal length is provided as one value, expand to # create (N, 2) shape tensor if self.focal_length.ndim == 1: # (N,) self.focal_length = self.focal_length[:, None] # (N, 1) self.focal_length = self.focal_length.expand(-1, 2) # (N, 2) self.proj_mats=None def transform_points_to_view(self, points, eps = None, **kwargs): #from wold to view R: torch.Tensor = kwargs.get("R", self.R) T: torch.Tensor = kwargs.get("T", self.T) self.R = R self.T = T if R.dim() == 2 : Tmat=torch.eye(4,device=R.device)[None] Tmat[:,:3,:3] = R Tmat[:,:3,3] = T else: Tmat=torch.eye(4,device=R.device)[None].repeat(R.shape[0],1,1) Tmat[:,:3,:3] = R Tmat[:,:3,3] = T points_batch = points.clone() if points_batch.dim() == 2: points_batch = points_batch[None] # (P, 3) -> (1, P, 3) if points_batch.dim() != 3: msg = "Expected points to have dim = 2 or dim = 3: got shape %r" raise ValueError(msg % repr(points.shape)) N, P, _3 = points_batch.shape ones = torch.ones(N, P, 1, dtype=points.dtype, device=points.device) points_batch = torch.cat([points_batch, ones], dim=2) # points_out=_broadcast_bmm(points_batch,Tmat) points_out=torch.einsum('bij,bnj->bni',Tmat,points_batch) return points_out[:,:,:3] def get_projection_transform(self,device): if self.proj_mats is None: # 构造内参，中心点为 00， focal 为 24 proj_mats=[] if torch.unique(self.focal_length).numel()==1: invtanfov=self.focal_length[0,0] proj_mat=get_proj_matrix(1/invtanfov,device) proj_mats=proj_mat[None].repeat(self.focal_length.shape[0],1,1) else: for invtanfov in self.focal_length: invtanfov=invtanfov[0]; assert invtanfov[0]==invtanfov[1] proj_mat=get_proj_matrix(1/invtanfov,device) proj_mats.append(proj_mat[None]) proj_mats=torch.cat(proj_mats,dim=0) self.proj_mats=proj_mats else: proj_mats=self.proj_mats return proj_mats def transform_points_to_ndc(self, points, eps = None, **kwargs): #from wold to ndc R: torch.Tensor = kwargs.get("R", self.R) T: torch.Tensor = kwargs.get("T", self.T) self.R = R self.T = T if R.dim() == 2 : Tmat=torch.eye(4,device=R.device)[None] Tmat[:,:3,:3] = R.clone() Tmat[:,:3,3] = T.clone() else: Tmat=torch.eye(4,device=R.device)[None].repeat(R.shape[0],1,1) # true Tmat[:,:3,:3] = R.clone() Tmat[:,:3,3] = T.clone() #points_view=self.transform_points_to_view(points) N, P, _3 = points.shape ones = torch.ones(N, P, 1, dtype=points.dtype, device=points.device) points_h = torch.cat([points, ones], dim=2) proj_mat=self.get_projection_transform(points.device)#[None].expand(N,-1,-1) proj_mat=proj_mat.to(R.device) B = Tmat.shape[0] if B > self.focal_length.shape[0]: # 第二阶段使用的，concat在一起 proj_mat = proj_mat.repeat(2,1,1) full_mat=torch.bmm(proj_mat[:Tmat.shape[0]],Tmat) # 防止最后一个batch points_ndc=torch.einsum('bij,bnj->bni',full_mat,points_h) points_ndc_xyz=points_ndc[:,:,:3]/(points_ndc[:,:,3:]+1e-7) # points_ndc_xyz[:,:,2]=points_ndc[:,:,3] #retain z range return points_ndc_xyz def transform_points_view_to_ndc(self, points, eps = None, **kwargs): #from view to ndc points_view=points.clone() N, P, _3 = points_view.shape ones = torch.ones(N, P, 1, dtype=points.dtype, device=points_view.device) points_view = torch.cat([points_view, ones], dim=2) proj_mat=self.get_projection_transform(points.device)#[None].expand(N,-1,-1) # points_ndc=_broadcast_bmm(points_view,proj_mat) points_ndc=torch.einsum('bij,bnj->bni',proj_mat,points_view) points_ndc_xyz=points_ndc[:,:,:3]/(points_ndc[:,:,3:]+1e-7) return points_ndc_xyz def transform_points_to_screen(self, points, with_xyflip = True, **kwargs): #from wold to screen 'with_xyflip: obey pytroch 3d coordinate' R: torch.Tensor = kwargs.get("R", self.R) T: torch.Tensor = kwargs.get("T", self.T) self.R = R self.T = T points_ndc=self.transform_points_to_ndc(points) N, P, _3 = points_ndc.shape image_size=self.image_size[:N] if not torch.is_tensor(image_size): image_size = torch.tensor(image_size, device=R.device) if image_size.dim()==2: image_size = image_size[:,None] image_size=image_size[:,:,[1,0]]#width height points_screen=points_ndc.clone() # points_ndc[...,:2] 以 (0,0) 为 光心 # points_screen[...,:2]=points_ndc[...,:2] * image_size/2-image_size/2 # x,y in [-1024,0] points_screen[...,:2]= (points_ndc[...,:2] - 1)* image_size/2 # x,y in [-1024,0] if with_xyflip: # true points_screen[...,:2]=points_screen[:,:,:2]*-1 # 转化到 [0,1024]，似乎可以前面+1，这里就不用取反了 return points_screen def transform_points_screen(self, points, with_xyflip = True, **kwargs): return self.transform_points_to_screen(points, with_xyflip, **kwargs) def perspective_projection(self, points, with_xyflip = True, **kwargs): ''' Computes the perspective projection of a set of 3D points. https://github.com/shubham-goel/4D-Humans/blob/6ec79656a23c33237c724742ca2a0ec00b398b53/hmr2/utils/geometry.py#L64-L102 ''' # from wold to screen 'with_xyflip: obey pytroch 3d coordinate' rotation: torch.Tensor = kwargs.get("R", self.R) translation: torch.Tensor = kwargs.get("T", self.T) B = points.shape[0] if rotation is None: rotation = torch.tensor([ [-1.0, 0.0, 0.0], [ 0.0, -1.0, 0.0], [ 0.0, 0.0, 1.0] ], device=points.device, dtype=points.dtype).unsqueeze(0).expand(B, -1, -1) # rotation = torch.eye(3, device=points.device, dtype=points.dtype).unsqueeze(0).expand(B, -1, -1) camera_center = torch.zeros(B, 2, device=points.device, dtype=points.dtype) # Populate intrinsic camera matrix K. K = torch.zeros([B, 3, 3], device=points.device, dtype=points.dtype) K[:, 0, 0] = 24 K[:, 1, 1] = 24 K[:, 2, 2] = 1. K[:, :-1, -1] = camera_center # Transform points points = torch.einsum('bij, bkj -> bki', rotation, points) points = points + translation.unsqueeze(1) # Apply perspective distortion projected_points = points / points[:, :, -1].unsqueeze(-1) # Apply camera intrinsics projected_points = torch.einsum('bij, bkj -> bki', K, projected_points) points_screen = projected_points.clone() # points_ndc[...,:2] 以 (0,0) 为 光心 # points_screen[...,:2]=points_ndc[...,:2] * image_size/2-image_size/2 # x,y in [-1024,0] points_screen[...,:2]= (projected_points[...,:2] - 1) * 1024 / 2 # x,y in [-1024,0] if with_xyflip: # true points_screen[...,:2] = points_screen[:,:,:2]*-1 # 转化到 [0,1024]，似乎可以前面+1，这里就不用取反了 return points_screen # [B,N,3] class GS_MeshRasterizer(MeshRasterizer): """ adapted to GS_camera This class implements methods for rasterizing a batch of heterogeneous Meshes. """ def __init__(self, cameras:GS_Camera=None, raster_settings=None) -> None: """ Args: cameras: A cameras object which has a `transform_points` method which returns the transformed points after applying the world-to-view and view-to-ndc transformations. raster_settings: the parameters for rasterization. This should be a named tuple. All these initial settings can be overridden by passing keyword arguments to the forward function. """ super().__init__() if raster_settings is None: raster_settings = RasterizationSettings() self.cameras = cameras self.raster_settings = raster_settings def to(self, device): # Manually move to device cameras as it is not a subclass of nn.Module if self.cameras is not None: self.cameras = self.cameras.to(device) return self def transform(self, meshes_world, **kwargs) -> torch.Tensor: #adapted to GS_camera cameras = kwargs.get("cameras", self.cameras) if cameras is None: msg = "Cameras must be specified either at initialization \ or in the forward pass of MeshRasterizer" raise ValueError(msg) n_cameras = len(cameras) if n_cameras != 1 and n_cameras != len(meshes_world): msg = "Wrong number (%r) of cameras for %r meshes" raise ValueError(msg % (n_cameras, len(meshes_world))) verts_world = meshes_world.verts_padded() # NOTE: Retaining view space z coordinate for now. # TODO: Revisit whether or not to transform z coordinate to [-1, 1] or # [0, 1] range. eps = kwargs.get("eps", None) verts_view = cameras.transform_points_to_view(verts_world, eps=eps,**kwargs) verts_ndc = cameras.transform_points_view_to_ndc(verts_view, eps=eps,**kwargs) verts_ndc[..., 2] = verts_view[..., 2] meshes_ndc = meshes_world.update_padded(new_verts_padded=verts_ndc) return meshes_ndc def forward(self, meshes_world, **kwargs) -> Fragments: """ Args: meshes_world: a Meshes object representing a batch of meshes with coordinates in world space. Returns: Fragments: Rasterization outputs as a named tuple. """ meshes_proj = self.transform(meshes_world, **kwargs) raster_settings = kwargs.get("raster_settings", self.raster_settings) # By default, turn on clip_barycentric_coords if blur_radius > 0. # When blur_radius > 0, a face can be matched to a pixel that is outside the # face, resulting in negative barycentric coordinates. clip_barycentric_coords = raster_settings.clip_barycentric_coords if clip_barycentric_coords is None: clip_barycentric_coords = raster_settings.blur_radius > 0.0 # If not specified, infer perspective_correct and z_clip_value from the camera cameras = kwargs.get("cameras", self.cameras) perspective_correct=False z_clip = None if raster_settings.perspective_correct is not None: perspective_correct = raster_settings.perspective_correct else: perspective_correct = True # if raster_settings.z_clip_value is not None: # z_clip = raster_settings.z_clip_value # else: # znear = cameras.get_znear() # if isinstance(znear, torch.Tensor): # znear = znear.min().item() # z_clip = None if not perspective_correct or znear is None else znear / 2 # By default, turn on clip_barycentric_coords if blur_radius > 0. # When blur_radius > 0, a face can be matched to a pixel that is outside the # face, resulting in negative barycentric coordinates. pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( meshes_proj, image_size=raster_settings.image_size, # 512 blur_radius=raster_settings.blur_radius, # 0 faces_per_pixel=raster_settings.faces_per_pixel, # 1 bin_size=raster_settings.bin_size, # None Bin size was too small in the coarse rasterization phase. max_faces_per_bin=raster_settings.max_faces_per_bin, # None clip_barycentric_coords=clip_barycentric_coords, # False perspective_correct=perspective_correct, # True cull_backfaces=raster_settings.cull_backfaces, # false z_clip_value=z_clip, # None cull_to_frustum=raster_settings.cull_to_frustum, # flase ) return Fragments( pix_to_face=pix_to_face, # 每个像素最近的面片索引（-1 表示无） zbuf=zbuf, # 每个像素对应的深度值 bary_coords=bary_coords, # 每个像素相对于三角形的重心坐标 dists=dists, # 像素距离面片的距离（用于模糊/soft raster） ) class BaseMeshRenderer(): def __init__(self, faces,image_size=512,lbs_weights=None, skin_color=[252, 224, 203], bg_color=[0, 0, 0], faces_uvs=None,verts_uvs=None,focal_length=24,inverse_light=False): super(BaseMeshRenderer, self).__init__() self.image_size = image_size self.skin_color = np.array(skin_color) self.bg_color = bg_color self.focal_length = focal_length bin_size=None # if image_size==296: # bin_size=20 self.raster_settings = RasterizationSettings(image_size=image_size, blur_radius=0.0, faces_per_pixel=1, bin_size=bin_size)#bin_size=0 max_faces_per_bin=20_000_0 if inverse_light: self.lights = PointLights( location=[[0.0, -1.0, -10.0]]) else: self.lights = PointLights( location=[[0.0, 1.0, 10.0]]) self.manual_lights = PointLights( location=((0.0, 0.0, 5.0), ), ambient_color=((0.5, 0.5, 0.5), ), diffuse_color=((0.5, 0.5, 0.5), ), specular_color=((0.01, 0.01, 0.01), ) ) self.blend = BlendParams(background_color=np.array(bg_color)/225.) self.faces = torch.nn.Parameter(faces, requires_grad=False) if faces_uvs is not None: self.faces_uvs = torch.nn.Parameter(faces_uvs, requires_grad=False) if verts_uvs is not None: self.verts_uvs = torch.nn.Parameter(verts_uvs.clone(), requires_grad=False) #SoftPhongShader will flip the v, so we need to flip it back self.verts_uvs[:,1]=1-self.verts_uvs[:,1] self.lbs_weights = None if lbs_weights is not None: self.lbs_weights = torch.nn.Parameter(lbs_weights, requires_grad=False) def _build_cameras(self, transform_matrix, focal_length): device = transform_matrix.device batch_size = transform_matrix.shape[0] screen_size = torch.tensor( [self.image_size, self.image_size], device=device ).float()[None].repeat(batch_size, 1) cameras_kwargs = { 'principal_point': torch.zeros(batch_size, 2, device=device).float(), 'focal_length': focal_length, 'image_size': screen_size, 'device': device, } #cameras = PerspectiveCameras(**cameras_kwargs, R=transform_matrix[:, :3, :3], T=transform_matrix[:, :3, 3]) cameras = GS_Camera(**cameras_kwargs, R=transform_matrix[:, :3, :3], T=transform_matrix[:, :3, 3]) return cameras def forward(self, vertices, faces=None, landmarks={}, cameras=None, transform_matrix=None, focal_length=None, ret_image=True): B, V = vertices.shape[:2] focal_length = self.focal_length if focal_length is None else focal_length if isinstance(cameras, torch.Tensor): cameras = cameras.clone() elif cameras is None: cameras = self._build_cameras(transform_matrix, focal_length) t_faces = faces[None].repeat(B, 1, 1) ret_vertices = cameras.transform_points_screen(vertices) ret_landmarks = {k: cameras.transform_points_screen(v) for k,v in landmarks.items()} images = None if ret_image: # Initialize each vertex to be white in color. verts_rgb = torch.from_numpy(self.skin_color/255).float().to(self.device)[None, None, :].repeat(B, V, 1) textures = TexturesVertex(verts_features=verts_rgb) mesh = Meshes( verts=vertices.to(self.device), faces=t_faces.to(self.device), textures=textures ) renderer = MeshRenderer(#GS_MeshRasterizer MeshRasterizer rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings), shader=SoftPhongShader(cameras=cameras, lights=self.lights.to(vertices.device), device=self.device, blend_params=self.blend) ) render_results = renderer(mesh).permute(0, 3, 1, 2) images = render_results[:, :3] alpha_images = render_results[:, 3:] images[alpha_images.expand(-1, 3, -1, -1)<0.5] = 0.0 images = images * 255 return ret_vertices, ret_landmarks, images def render_mesh(self, vertices,cameras=None,transform_matrix=None, faces=None,lights=None,reverse_camera=True): #render mesh vertices value and lbs weights device = vertices.device B, V = vertices.shape[:2] if faces is None: faces = self.faces if cameras is None: transform_matrix=transform_matrix.clone() if reverse_camera: tf_mat=torch.tensor([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]],dtype=torch.float32).to(device) tf_mat=tf_mat[None].expand(B,-1,-1) transform_matrix = torch.bmm(tf_mat,transform_matrix) cameras = self._build_cameras(transform_matrix, self.focal_length) t_faces = faces[None].repeat(B, 1, 1) # self.lights=lights # if lights is None: # self.lights = PointLights(device=device, location=[[0.0, 1.0, 10.0]]) # Initialize each vertex to be white in color. # verts_rgb = vertices.clone() # textures = TexturesVertex(verts_features=verts_rgb) # mesh = Meshes( # verts=vertices.to(device), # faces=t_faces.to(device), # textures=textures # ) # renderer = MeshRenderer( # rasterizer=MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings), # shader=SoftPhongShader(cameras=cameras, lights=self.lights, device=device, blend_params=self.blend) # ) mesh = Meshes( verts=vertices.to(device), faces=t_faces.to(device), ) shader = VertexPositionShader().to(device) rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings) #GS_MeshRasterizer MeshRasterizer renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) render_results,extra_result = renderer(mesh) render_lbs_weights=None if self.lbs_weights is not None: vertex_faces=extra_result['vertex_faces'] bary_coords=extra_result['bary_coords'] lbs_weights=self.lbs_weights[None].expand(B, -1, -1).reshape(-1,55) render_lbs_weights=(lbs_weights[vertex_faces]*bary_coords[...,None]).sum(dim=-2) # images = render_results[:, :3] # alpha_images = render_results[:, 3:] return render_results,render_lbs_weights def render_fragments(self, vertices,cameras=None,transform_matrix=None, faces=None,reverse_camera=True): device = vertices.device B, V = vertices.shape[:2] if faces is None: faces = self.faces if cameras is None: transform_matrix=transform_matrix.clone() if reverse_camera: tf_mat=torch.tensor([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]],dtype=torch.float32).to(device) tf_mat=tf_mat[None].expand(B,-1,-1) transform_matrix = torch.bmm(tf_mat,transform_matrix) cameras = self._build_cameras(transform_matrix, self.focal_length) t_faces = faces[None].repeat(B, 1, 1) mesh = Meshes( # 顶点和面构成 mesh verts=vertices.to(device), faces=t_faces.to(device), ) rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings) fragments = rasterizer(mesh) # 把 mesh 渲染到平面 #return visble faces return fragments.pix_to_face,fragments def render_textured_mesh(self,vertices,uvmap,fragments=None,faces_uvs=None,verts_uvs=None,faces=None,cameras=None,transform_matrix=None,reverse_camera=True): device = vertices.device B, V = vertices.shape[:2] if faces is None: faces = self.faces if faces_uvs is None: faces_uvs = self.faces_uvs if verts_uvs is None: verts_uvs = self.verts_uvs if cameras is None: transform_matrix=transform_matrix.clone() if reverse_camera: tf_mat=torch.tensor([[-1,0,0,0],[0,-1,0,0],[0,0,1,0],[0,0,0,1]],dtype=torch.float32).to(device) tf_mat=tf_mat[None].expand(B,-1,-1) transform_matrix = torch.bmm(tf_mat,transform_matrix) cameras = self._build_cameras(transform_matrix, self.focal_length) t_faces = faces[None].repeat(B, 1, 1) t_faces_uvs = faces_uvs[None].repeat(B, 1, 1) t_verts_uvs = verts_uvs[None].repeat(B, 1, 1) textures = TexturesUV(maps=uvmap,faces_uvs=t_faces_uvs,verts_uvs=t_verts_uvs) mesh = Meshes( verts=vertices.to(self.device), faces=t_faces.to(self.device), textures=textures ) lights = PointLights( location=[[0.0, 0.0, 1000.0]]) if fragments is None: rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings) fragments = rasterizer(mesh) shader=SoftPhongShader(cameras=cameras, lights=lights, device=device, blend_params=self.blend).to(device) images=shader(fragments, mesh) return images class GS_BaseMeshRenderer(torch.nn.Module): #RENDERING IN GS PROJECTION METHOD def __init__(self,image_size=512, skin_color=[252, 224, 203], bg_color=[0, 0, 0], focal_length=24,inverse_light=False): super(GS_BaseMeshRenderer, self).__init__() self.image_size = image_size self.skin_color = np.array(skin_color) self.bg_color = bg_color self.focal_length = focal_length self.raster_settings = RasterizationSettings(image_size=image_size, blur_radius=0.0, faces_per_pixel=1) if inverse_light: self.lights = PointLights( location=[[0.0, -1.0, -10.0]]) else: self.lights = PointLights( location=[[0.0, 1.0, 10.0]]) self.manual_lights = PointLights( location=((0.0, 0.0, 5.0), ), ambient_color=((0.5, 0.5, 0.5), ), diffuse_color=((0.5, 0.5, 0.5), ), specular_color=((0.01, 0.01, 0.01), ) ) self.blend = BlendParams(background_color=np.array(bg_color)/225.) # self.faces = torch.nn.Parameter(faces, requires_grad=False) # self.faces=None self.head_color=np.array([236,248,254]) #np.array([222,235,247]) def _build_cameras(self, transform_matrix, focal_length): device = transform_matrix.device batch_size = transform_matrix.shape[0] screen_size = torch.tensor( [self.image_size, self.image_size], device=device ).float()[None].repeat(batch_size, 1) cameras_kwargs = { 'principal_point': torch.zeros(batch_size, 2, device=device).float(), 'focal_length': focal_length, 'image_size': screen_size, 'device': device, } #cameras = PerspectiveCameras(**cameras_kwargs, R=transform_matrix[:, :3, :3], T=transform_matrix[:, :3, 3]) cameras = GS_Camera(**cameras_kwargs, R=transform_matrix[:, :3, :3], T=transform_matrix[:, :3, 3]) return cameras def forward(self, vertices, faces=None, landmarks={}, cameras=None, transform_matrix=None, focal_length=None, ret_image=True): B, V = vertices.shape[:2] device=vertices.device focal_length = self.focal_length if focal_length is None else focal_length if isinstance(cameras, torch.Tensor): cameras = cameras.clone() elif cameras is None: cameras = self._build_cameras(transform_matrix, focal_length) # 24 t_faces = faces[None].repeat(B, 1, 1) ret_vertices = cameras.transform_points_screen(vertices) ret_landmarks = {k: cameras.transform_points_screen(v) for k,v in landmarks.items()} images = None if ret_image: # Initialize each vertex to be white in color. verts_rgb = torch.from_numpy(self.skin_color/255).float().to(self.device)[None, None, :].repeat(B, V, 1) textures = TexturesVertex(verts_features=verts_rgb) mesh = Meshes( verts=vertices.to(self.device), faces=t_faces.to(self.device), textures=textures ) renderer = MeshRenderer(#GS_MeshRasterizer MeshRasterizer rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings), shader=SoftPhongShader(cameras=cameras, lights=self.lights.to(device), device=device, blend_params=self.blend) ) render_results = renderer(mesh).permute(0, 3, 1, 2) images = render_results[:, :3] alpha_images = render_results[:, 3:] images[alpha_images.expand(-1, 3, -1, -1)<0.5] = 0.0 images = images * 255 return ret_vertices, ret_landmarks, images def render_mesh(self, vertices,cameras=None,transform_matrix=None, faces=None,lights=None,skin_color=None,smplx2flame_ind=None): device = vertices.device B, V = vertices.shape[:2] if faces is None: faces = self.faces assert faces is not None if cameras is None: transform_matrix=transform_matrix.clone() cameras = self._build_cameras(transform_matrix, self.focal_length) if lights is None: self.lights = self.lights else: self.lights=lights if faces.dim() == 2: faces = faces[None] t_faces = faces.repeat(B, 1, 1) # Initialize each vertex to be white in color. if skin_color is None: skin_color=self.skin_color if isinstance(skin_color, (list, tuple)): verts_rgb = torch.from_numpy(np.array(skin_color)/255.).to(vertices.device).float()[None, None, :].repeat(B, V, 1) else: verts_rgb = torch.from_numpy(skin_color/255).to(vertices.device).float()[None, None, :].repeat(B, V, 1) # here if smplx2flame_ind is not None: head_rgb = torch.from_numpy(self.head_color/255).to(vertices.device).float()[None, None, :].repeat(B, V, 1) verts_rgb[:,smplx2flame_ind] = head_rgb[:,smplx2flame_ind] textures = TexturesVertex(verts_features=verts_rgb) mesh = Meshes( verts=vertices.to(device), faces=t_faces.to(device), textures=textures ) # rot = torch.tensor( # trimesh.transformations.rotation_matrix( # np.radians(-60), [0, 1, 0] # )[:3, :3], dtype=torch.float32, device=device # ) # verts_rot = torch.matmul(mesh.verts_packed(), rot.T)[None] # # 替换顶点 # mesh = Meshes( # verts=verts_rot, # faces=t_faces.to(device), # textures=textures # ) renderer = MeshRenderer( # GS_MeshRasterizer MeshRasterizer rasterizer=GS_MeshRasterizer(cameras=cameras, raster_settings=self.raster_settings), shader=SoftPhongShader(cameras=cameras, lights=self.lights.to(device), device=device, blend_params=self.blend) ) render_results = renderer(mesh).permute(0, 3, 1, 2) images = render_results[:, :3] alpha_images = render_results[:, 3:] alpha = alpha_images.expand(-1, 3, -1, -1) < 0.5 # White background instead of black: set background pixels to 1.0 images[alpha] = 1.0 images = torch.cat([images, 1 - alpha[:, :1] * 1.0], dim=1) images = images * 255 return images if __name__=="__main__": pass import pdb;pdb.set_trace()