utils/graphics.py

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,
    SoftPhongShader, MeshRasterizer, MeshRenderer,
)
from pytorch3d.transforms.transform3d import _broadcast_bmm
from pytorch3d.renderer.mesh.rasterizer import Fragments,rasterize_meshes
from utils.graphics_utils import GS_Camera
import cv2
import torch.nn.functional as F
import matplotlib.pyplot as plt

def overlay_attention_on_image(
    image_path,
    attn,
    patch_h=16,
    patch_w=12,
    alpha=0.5,
    save_path="attn_overlay.png"
):
    """
    image_path: str, path to input image
    attn: torch.Tensor [1, 8, 1, 192]
    """

    # -------------------------
    # 1. Load image
    # -------------------------
    img_bgr = cv2.imread(image_path)
    assert img_bgr is not None, f"Cannot read image: {image_path}"

    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    H_img, W_img, _ = img.shape

    # -------------------------
    # 2. Process attention
    # -------------------------
    attn_map = attn[0, :, 0, :]          # [8, 192]
    attn_mean = attn_map.mean(dim=0)     # [192]

    heat = attn_mean.view(patch_h, patch_w)  # [16, 12]

    # -------------------------
    # 3. Upsample heatmap
    # -------------------------
    heat_up = F.interpolate(
        heat[None, None],
        size=(H_img, W_img),
        mode="bilinear",
        align_corners=False
    )[0, 0]

    # normalize to [0,1]
    heat_up = (heat_up - heat_up.min()) / (heat_up.max() - heat_up.min() + 1e-6)
    heat_np = heat_up.detach().cpu().numpy()

    # -------------------------
    # 4. Apply colormap
    # -------------------------
    heat_color = cv2.applyColorMap(
        (heat_np * 255).astype(np.uint8),
        cv2.COLORMAP_JET
    )
    heat_color = cv2.cvtColor(heat_color, cv2.COLOR_BGR2RGB)

    # -------------------------
    # 5. Overlay
    # -------------------------
    overlay = (1 - alpha) * img + alpha * heat_color
    overlay = overlay.astype(np.uint8)

    # -------------------------
    # 6. Save result
    # -------------------------
    plt.figure(figsize=(6, 5))
    plt.imshow(overlay)
    plt.axis("off")
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close()

    print(f"Overlay saved to: {save_path}")


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


class GS_Camera(CamerasBase):
    #still obey pytorch 3d coordinate system
    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),),#assume to zero
        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
        self.z_near= 0.01#0.01 # only influence z in ndc 
        self.z_far=100 #100
        if self.focal_length.ndim == 1:  # (N,)
            self.focal_length = self.focal_length[:, None]  # (N, 1)  初始化为 12 ？
        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.clone()
            Tmat[:,:3,3] = T.clone()
        else:
            
            Tmat=torch.eye(4,device=R.device)[None].repeat(R.shape[0],1,1)
            Tmat[:,:3,:3] = R.clone()
            Tmat[:,:3,3] = T.clone()
            
        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)
        #points_out=torch.bmm(Tmat,points_batch.transpose(1,2)).transpose(1,2)
        #points_out=Tmat.bmm(points_batch.transpose(1,2)).transpose(1,2)
        return points_out[:,:,:3]
    
    def get_projection_transform(self,device):
        if self.proj_mats is None:
            proj_mats=[]
            if  torch.unique(self.focal_length).numel()==1: # True
                invtanfov=self.focal_length[0,0]  # 
                proj_mat=get_proj_matrix(1/invtanfov,device,z_near=self.z_near,z_far=self.z_far)  # 内参？
                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,z_near=self.z_near,z_far=self.z_far)
                    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)
            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)
        
        full_mat=torch.bmm(proj_mat,Tmat)
        #points_ndc=_broadcast_bmm(points_h,full_mat)
        points_ndc=torch.einsum('bij,bnj->bni',full_mat,points_h)
        #points_ndc=torch.bmm(proj_mat,points_view.transpose(1,2)).transpose(1,2)
        #points_ndc=full_mat.bmm(points_h.transpose(1,2)).transpose(1,2)
        
        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=torch.bmm(proj_mat,points_view.transpose(1,2)).transpose(1,2)
        #points_ndc=proj_mat.bmm(points_view.transpose(1,2)).transpose(1,2)
        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
        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_screen[...,:2]=points_ndc[...,:2]*image_size/2-image_size/2
        if with_xyflip:
            points_screen[...,:2]=points_screen[:,:,:2]*-1
        
        return points_screen
    
    def transform_points_screen(self, points, with_xyflip = True, **kwargs):
        return self.transform_points_to_screen(points, with_xyflip, **kwargs)
    
        
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.
        """
        cameras = kwargs.get("cameras", self.cameras)
        self.cameras=cameras
        # assert isinstance(cameras, GS_Camera)

        
        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
        
        perspective_correct=False
        z_clip = None
        if raster_settings.perspective_correct is not None:
            perspective_correct = raster_settings.perspective_correct
        else:
            perspective_correct = True
            
        pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
            meshes_proj,
            image_size=raster_settings.image_size,
            blur_radius=raster_settings.blur_radius,
            faces_per_pixel=raster_settings.faces_per_pixel,
            bin_size=raster_settings.bin_size,
            max_faces_per_bin=raster_settings.max_faces_per_bin,
            clip_barycentric_coords=clip_barycentric_coords,
            perspective_correct=perspective_correct,
            cull_backfaces=raster_settings.cull_backfaces,
            z_clip_value=z_clip,
            cull_to_frustum=raster_settings.cull_to_frustum,
        )

        return Fragments(
            pix_to_face=pix_to_face,
            zbuf=zbuf,
            bary_coords=bary_coords,
            dists=dists,
        )


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)
        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
        )
        
        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
        images[alpha] = 0.0
        images=torch.cat([images,1-alpha[:,:1]*1.0],dim=1)
        images = images * 255
        
        return images