sgl-project--sglang/python/sglang/multimodal_gen/runtime/utils/mesh3d_utils.py

"""Adapted from Hunyuan3D-2: https://github.com/Tencent/Hunyuan3D-2"""

from __future__ import annotations

import math
from typing import Any, List, Optional, Tuple, Union

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import trimesh
from einops import rearrange, repeat
from PIL import Image

from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)

# Import C++ mesh processor extension
from sglang.multimodal_gen.csrc.render.mesh_processor import meshVerticeInpaint


def transform_pos(
    mtx: Union[np.ndarray, torch.Tensor],
    pos: torch.Tensor,
    keepdim: bool = False,
) -> torch.Tensor:
    """Transform positions by a matrix."""
    t_mtx = torch.from_numpy(mtx).to(pos.device) if isinstance(mtx, np.ndarray) else mtx

    if pos.shape[-1] == 3:
        posw = torch.cat([pos, torch.ones([pos.shape[0], 1]).to(pos.device)], axis=1)
    else:
        posw = pos

    if keepdim:
        return torch.matmul(posw, t_mtx.t())[...]
    else:
        return torch.matmul(posw, t_mtx.t())[None, ...]


def get_mv_matrix(
    elev: float,
    azim: float,
    camera_distance: float,
    center: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Compute model-view matrix from camera parameters."""
    elev = -elev
    azim += 90

    elev_rad = math.radians(elev)
    azim_rad = math.radians(azim)

    camera_position = np.array(
        [
            camera_distance * math.cos(elev_rad) * math.cos(azim_rad),
            camera_distance * math.cos(elev_rad) * math.sin(azim_rad),
            camera_distance * math.sin(elev_rad),
        ]
    )

    if center is None:
        center = np.array([0, 0, 0])
    else:
        center = np.array(center)

    lookat = center - camera_position
    lookat = lookat / np.linalg.norm(lookat)

    up = np.array([0, 0, 1.0])
    right = np.cross(lookat, up)
    right = right / np.linalg.norm(right)
    up = np.cross(right, lookat)
    up = up / np.linalg.norm(up)

    c2w = np.concatenate(
        [np.stack([right, up, -lookat], axis=-1), camera_position[:, None]], axis=-1
    )

    w2c = np.zeros((4, 4))
    w2c[:3, :3] = np.transpose(c2w[:3, :3], (1, 0))
    w2c[:3, 3:] = -np.matmul(np.transpose(c2w[:3, :3], (1, 0)), c2w[:3, 3:])
    w2c[3, 3] = 1.0

    return w2c.astype(np.float32)


def get_orthographic_projection_matrix(
    left: float = -1,
    right: float = 1,
    bottom: float = -1,
    top: float = 1,
    near: float = 0,
    far: float = 2,
) -> np.ndarray:
    """Compute orthographic projection matrix."""
    ortho_matrix = np.eye(4, dtype=np.float32)
    ortho_matrix[0, 0] = 2 / (right - left)
    ortho_matrix[1, 1] = 2 / (top - bottom)
    ortho_matrix[2, 2] = -2 / (far - near)
    ortho_matrix[0, 3] = -(right + left) / (right - left)
    ortho_matrix[1, 3] = -(top + bottom) / (top - bottom)
    ortho_matrix[2, 3] = -(far + near) / (far - near)
    return ortho_matrix


def get_perspective_projection_matrix(
    fovy: float,
    aspect_wh: float,
    near: float,
    far: float,
) -> np.ndarray:
    """Compute perspective projection matrix."""
    fovy_rad = math.radians(fovy)
    return np.array(
        [
            [1.0 / (math.tan(fovy_rad / 2.0) * aspect_wh), 0, 0, 0],
            [0, 1.0 / math.tan(fovy_rad / 2.0), 0, 0],
            [0, 0, -(far + near) / (far - near), -2.0 * far * near / (far - near)],
            [0, 0, -1, 0],
        ]
    ).astype(np.float32)


def export_to_trimesh(mesh_output: Any) -> Any:
    """Convert mesh output to trimesh format."""
    if isinstance(mesh_output, list):
        outputs = []
        for mesh in mesh_output:
            if mesh is None:
                outputs.append(None)
            else:
                # Reverse face winding
                mesh.mesh_f = mesh.mesh_f[:, ::-1]
                mesh_obj = trimesh.Trimesh(mesh.mesh_v, mesh.mesh_f)
                outputs.append(mesh_obj)
        return outputs
    else:
        mesh_output.mesh_f = mesh_output.mesh_f[:, ::-1]
        return trimesh.Trimesh(mesh_output.mesh_v, mesh_output.mesh_f)


def mesh_uv_wrap(mesh: Any) -> Any:
    """Apply UV unwrapping to mesh. In-place like native Hunyuan3D-2 for same layout."""
    try:
        import xatlas
    except ImportError:
        logger.warning("xatlas not available, skipping UV unwrap")
        return mesh

    if isinstance(mesh, trimesh.Scene):
        mesh = mesh.dump(concatenate=True)

    if len(mesh.faces) > 500000000:
        raise ValueError(
            "The mesh has more than 500,000,000 faces, which is not supported."
        )

    vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)

    mesh.vertices = mesh.vertices[vmapping]
    mesh.faces = indices
    if not hasattr(mesh.visual, "uv"):
        mesh.visual = trimesh.visual.TextureVisuals(
            uv=uvs, material=trimesh.visual.material.SimpleMaterial()
        )
    else:
        mesh.visual.uv = uvs

    return mesh


def stride_from_shape(shape: Tuple[int, ...]) -> List[int]:
    """Compute stride from shape for scatter operations."""
    stride = [1]
    for x in reversed(shape[1:]):
        stride.append(stride[-1] * x)
    return list(reversed(stride))


def scatter_add_nd_with_count(
    input: torch.Tensor,
    count: torch.Tensor,
    indices: torch.Tensor,
    values: torch.Tensor,
    weights: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Scatter add with counting for texture baking."""
    D = indices.shape[-1]
    C = input.shape[-1]
    size = input.shape[:-1]
    stride = stride_from_shape(size)

    assert len(size) == D

    input = input.view(-1, C)
    count = count.view(-1, 1)

    flatten_indices = (
        indices * torch.tensor(stride, dtype=torch.long, device=indices.device)
    ).sum(-1)

    if weights is None:
        weights = torch.ones_like(values[..., :1])

    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
    count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)

    return input.view(*size, C), count.view(*size, 1)


def linear_grid_put_2d(
    H: int,
    W: int,
    coords: torch.Tensor,
    values: torch.Tensor,
    return_count: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
    """Put values into a 2D grid using linear interpolation."""
    C = values.shape[-1]

    indices = coords * torch.tensor(
        [H - 1, W - 1], dtype=torch.float32, device=coords.device
    )
    indices_00 = indices.floor().long()
    indices_00[:, 0].clamp_(0, H - 2)
    indices_00[:, 1].clamp_(0, W - 2)
    indices_01 = indices_00 + torch.tensor(
        [0, 1], dtype=torch.long, device=indices.device
    )
    indices_10 = indices_00 + torch.tensor(
        [1, 0], dtype=torch.long, device=indices.device
    )
    indices_11 = indices_00 + torch.tensor(
        [1, 1], dtype=torch.long, device=indices.device
    )

    h = indices[..., 0] - indices_00[..., 0].float()
    w = indices[..., 1] - indices_00[..., 1].float()
    w_00 = (1 - h) * (1 - w)
    w_01 = (1 - h) * w
    w_10 = h * (1 - w)
    w_11 = h * w

    result = torch.zeros(H, W, C, device=values.device, dtype=values.dtype)
    count = torch.zeros(H, W, 1, device=values.device, dtype=values.dtype)
    weights = torch.ones_like(values[..., :1])

    result, count = scatter_add_nd_with_count(
        result,
        count,
        indices_00,
        values * w_00.unsqueeze(1),
        weights * w_00.unsqueeze(1),
    )
    result, count = scatter_add_nd_with_count(
        result,
        count,
        indices_01,
        values * w_01.unsqueeze(1),
        weights * w_01.unsqueeze(1),
    )
    result, count = scatter_add_nd_with_count(
        result,
        count,
        indices_10,
        values * w_10.unsqueeze(1),
        weights * w_10.unsqueeze(1),
    )
    result, count = scatter_add_nd_with_count(
        result,
        count,
        indices_11,
        values * w_11.unsqueeze(1),
        weights * w_11.unsqueeze(1),
    )

    if return_count:
        return result, count

    mask = count.squeeze(-1) > 0
    result[mask] = result[mask] / count[mask].repeat(1, C)

    return result


class MeshRender:
    """Mesh renderer using CUDA rasterization for texture generation."""

    def __init__(
        self,
        camera_distance: float = 1.45,
        camera_type: str = "orth",
        default_resolution: int = 1024,
        texture_size: int = 1024,
        bake_mode: str = "linear",
        device: str = "cuda",
    ):
        """Initialize the mesh renderer."""
        self.device = device

        self.set_default_render_resolution(default_resolution)
        self.set_default_texture_resolution(texture_size)

        self.camera_distance = camera_distance
        self.camera_type = camera_type
        self.bake_angle_thres = 75
        self.bake_unreliable_kernel_size = int(
            (2 / 512) * max(self.default_resolution[0], self.default_resolution[1])
        )
        self.bake_mode = bake_mode

        # Set up camera projection matrix
        if camera_type == "orth":
            self.ortho_scale = 1.2
            self.camera_proj_mat = get_orthographic_projection_matrix(
                left=-self.ortho_scale * 0.5,
                right=self.ortho_scale * 0.5,
                bottom=-self.ortho_scale * 0.5,
                top=self.ortho_scale * 0.5,
                near=0.1,
                far=100,
            )
        elif camera_type == "perspective":
            self.camera_proj_mat = get_perspective_projection_matrix(
                49.13,
                self.default_resolution[1] / self.default_resolution[0],
                0.01,
                100.0,
            )
        else:
            raise ValueError(f"Unknown camera type: {camera_type}")

        # Mesh data
        self.vtx_pos = None
        self.pos_idx = None
        self.vtx_uv = None
        self.uv_idx = None
        self.tex = None
        self.mesh_copy = None
        self.scale_factor = 1.0

    def set_default_render_resolution(
        self, default_resolution: Union[int, Tuple[int, int]]
    ):
        """Set default rendering resolution."""
        if isinstance(default_resolution, int):
            default_resolution = (default_resolution, default_resolution)
        self.default_resolution = default_resolution

    def set_default_texture_resolution(self, texture_size: Union[int, Tuple[int, int]]):
        """Set default texture resolution."""
        if isinstance(texture_size, int):
            texture_size = (texture_size, texture_size)
        self.texture_size = texture_size

    def _rasterize(
        self,
        pos_clip: torch.Tensor,
        tri: torch.Tensor,
        resolution: Tuple[int, int],
    ) -> torch.Tensor:
        """Rasterize using CUDA rasterizer."""
        from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import rasterize

        if pos_clip.dim() == 2:
            pos_clip = pos_clip.unsqueeze(0)

        findices, barycentric = rasterize(pos_clip, tri, resolution)
        rast_out = torch.cat((barycentric, findices.unsqueeze(-1).float()), dim=-1)
        rast_out = rast_out.unsqueeze(0)
        return rast_out

    def _interpolate(
        self,
        attr: torch.Tensor,
        rast_out: torch.Tensor,
        tri: torch.Tensor,
    ) -> torch.Tensor:
        """Interpolate vertex attributes."""
        from sglang.multimodal_gen.csrc.render.hunyuan3d_rasterizer import interpolate

        barycentric = rast_out[0, ..., :-1]
        findices = rast_out[0, ..., -1].int()

        if attr.dim() == 2:
            attr = attr.unsqueeze(0)

        result = interpolate(attr, findices, barycentric, tri)
        return result

    def load_mesh(
        self,
        mesh: Union[trimesh.Trimesh, trimesh.Scene],
        scale_factor: float = 1.15,
        auto_center: bool = True,
    ):
        """Load a mesh for rendering."""
        if isinstance(mesh, trimesh.Scene):
            mesh = mesh.dump(concatenate=True)

        self.mesh_copy = mesh.copy()

        vtx_pos = mesh.vertices.astype(np.float32)
        pos_idx = mesh.faces.astype(np.int32)

        # Get UV coordinates if available
        if hasattr(mesh.visual, "uv") and mesh.visual.uv is not None:
            vtx_uv = mesh.visual.uv.astype(np.float32)
            uv_idx = pos_idx.copy()
        else:
            vtx_uv = None
            uv_idx = None

        self.vtx_pos = torch.from_numpy(vtx_pos).to(self.device).float()
        self.pos_idx = torch.from_numpy(pos_idx).to(self.device).to(torch.int32)

        if vtx_uv is not None and uv_idx is not None:
            self.vtx_uv = torch.from_numpy(vtx_uv).to(self.device).float()
            self.uv_idx = torch.from_numpy(uv_idx).to(self.device).to(torch.int32)
        else:
            self.vtx_uv = None
            self.uv_idx = None

        # Coordinate transformation (Y-up to Z-up)
        self.vtx_pos[:, [0, 1]] = -self.vtx_pos[:, [0, 1]]
        self.vtx_pos[:, [1, 2]] = self.vtx_pos[:, [2, 1]]
        if self.vtx_uv is not None:
            self.vtx_uv[:, 1] = 1.0 - self.vtx_uv[:, 1]

        if auto_center:
            max_bb = (self.vtx_pos - 0).max(0)[0]
            min_bb = (self.vtx_pos - 0).min(0)[0]
            center = (max_bb + min_bb) / 2
            scale = torch.norm(self.vtx_pos - center, dim=1).max() * 2.0
            self.vtx_pos = (self.vtx_pos - center) * (scale_factor / float(scale))
            self.scale_factor = scale_factor

    def save_mesh(self) -> trimesh.Trimesh:
        """Save mesh with current texture, reusing the original mesh object."""
        texture_data = self.get_texture()
        texture_img = Image.fromarray((texture_data * 255).astype(np.uint8))

        material = trimesh.visual.material.SimpleMaterial(
            image=texture_img, diffuse=(255, 255, 255)
        )
        self.mesh_copy.visual = trimesh.visual.TextureVisuals(
            uv=self.mesh_copy.visual.uv, image=texture_img, material=material
        )
        return self.mesh_copy

    def get_mesh(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Get mesh data with inverse coordinate transformation."""
        vtx_pos = self.vtx_pos.cpu().numpy().copy()
        pos_idx = self.pos_idx.cpu().numpy()

        # Inverse coordinate transformation
        vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]]
        vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]]

        if self.vtx_uv is not None:
            vtx_uv = self.vtx_uv.cpu().numpy().copy()
            vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1]
            uv_idx = self.uv_idx.cpu().numpy()
        else:
            vtx_uv = None
            uv_idx = None

        return vtx_pos, pos_idx, vtx_uv, uv_idx

    def set_texture(self, tex: Union[np.ndarray, torch.Tensor, Image.Image]):
        """Set texture for the mesh."""
        if isinstance(tex, np.ndarray):
            if tex.max() <= 1.0:
                tex = (tex * 255).astype(np.uint8)
            tex = Image.fromarray(tex.astype(np.uint8))
        elif isinstance(tex, torch.Tensor):
            tex_np = tex.cpu().numpy()
            if tex_np.max() <= 1.0:
                tex_np = (tex_np * 255).astype(np.uint8)
            tex = Image.fromarray(tex_np.astype(np.uint8))

        tex = tex.resize(self.texture_size).convert("RGB")
        tex = np.array(tex) / 255.0
        self.tex = torch.from_numpy(tex).to(self.device).float()

    def get_texture(self) -> np.ndarray:
        """Get current texture as numpy array."""
        if self.tex is None:
            return np.ones((*self.texture_size, 3), dtype=np.float32)
        return self.tex.cpu().numpy()

    def _get_pos_from_mvp(
        self,
        elev: float,
        azim: float,
        camera_distance: Optional[float] = None,
        center: Optional[np.ndarray] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Get camera-space and clip-space positions."""
        proj = self.camera_proj_mat
        r_mv = get_mv_matrix(
            elev=elev,
            azim=azim,
            camera_distance=(
                self.camera_distance if camera_distance is None else camera_distance
            ),
            center=center,
        )

        pos_camera = transform_pos(r_mv, self.vtx_pos, keepdim=True)
        pos_clip = transform_pos(proj, pos_camera)

        return pos_camera, pos_clip

    def render_normal(
        self,
        elev: float,
        azim: float,
        camera_distance: Optional[float] = None,
        center: Optional[np.ndarray] = None,
        resolution: Optional[Tuple[int, int]] = None,
        bg_color: List[float] = [1, 1, 1],
        use_abs_coor: bool = False,
        normalize_rgb: bool = True,
        return_type: str = "th",
    ) -> Union[torch.Tensor, np.ndarray, Image.Image]:
        """Render normal map from a viewpoint."""
        pos_camera, pos_clip = self._get_pos_from_mvp(
            elev, azim, camera_distance, center
        )

        if resolution is None:
            resolution = self.default_resolution
        if isinstance(resolution, (int, float)):
            resolution = (int(resolution), int(resolution))

        rast_out = self._rasterize(pos_clip, self.pos_idx, resolution)

        # Compute face normals
        if use_abs_coor:
            mesh_triangles = self.vtx_pos[self.pos_idx[:, :3].long(), :]
        else:
            pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4]
            mesh_triangles = pos_camera_3d[self.pos_idx[:, :3].long(), :]

        face_normals = F.normalize(
            torch.cross(
                mesh_triangles[:, 1, :] - mesh_triangles[:, 0, :],
                mesh_triangles[:, 2, :] - mesh_triangles[:, 0, :],
                dim=-1,
            ),
            dim=-1,
        )

        # Compute vertex normals
        vertex_normals = trimesh.geometry.mean_vertex_normals(
            vertex_count=self.vtx_pos.shape[0],
            faces=self.pos_idx.cpu().numpy(),
            face_normals=face_normals.cpu().numpy(),
        )
        vertex_normals = (
            torch.from_numpy(vertex_normals).float().to(self.device).contiguous()
        )

        # Interpolate normals
        normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx)

        # Apply visibility mask
        visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)
        bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device)
        normal = normal * visible_mask + bg_tensor * (1 - visible_mask)

        if normalize_rgb:
            normal = (normal + 1) * 0.5

        image = normal[0, ...]

        if return_type == "np":
            image = image.cpu().numpy()
        elif return_type == "pl":
            image = image.cpu().numpy() * 255
            image = Image.fromarray(image.astype(np.uint8))

        return image

    def render_position(
        self,
        elev: float,
        azim: float,
        camera_distance: Optional[float] = None,
        center: Optional[np.ndarray] = None,
        resolution: Optional[Tuple[int, int]] = None,
        bg_color: List[float] = [1, 1, 1],
        return_type: str = "th",
    ) -> Union[torch.Tensor, np.ndarray, Image.Image]:
        """Render position map from a viewpoint."""
        pos_camera, pos_clip = self._get_pos_from_mvp(
            elev, azim, camera_distance, center
        )

        if resolution is None:
            resolution = self.default_resolution
        if isinstance(resolution, (int, float)):
            resolution = (int(resolution), int(resolution))

        rast_out = self._rasterize(pos_clip, self.pos_idx, resolution)

        # Position colors (normalized vertex positions)
        tex_position = 0.5 - self.vtx_pos[:, :3] / self.scale_factor
        tex_position = tex_position.contiguous()

        # Interpolate positions
        position = self._interpolate(tex_position[None, ...], rast_out, self.pos_idx)

        # Apply visibility mask
        visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)
        bg_tensor = torch.tensor(bg_color, dtype=torch.float32, device=self.device)
        position = position * visible_mask + bg_tensor * (1 - visible_mask)

        image = position[0, ...]

        if return_type == "np":
            image = image.cpu().numpy()
        elif return_type == "pl":
            image = image.cpu().numpy() * 255
            image = Image.fromarray(image.astype(np.uint8))

        return image

    def render_normal_multiview(
        self,
        camera_elevs: List[float],
        camera_azims: List[float],
        use_abs_coor: bool = True,
    ) -> List[Image.Image]:
        """Render normal maps from multiple viewpoints."""
        normal_maps = []
        for elev, azim in zip(camera_elevs, camera_azims):
            normal_map = self.render_normal(
                elev, azim, use_abs_coor=use_abs_coor, return_type="pl"
            )
            normal_maps.append(normal_map)
        return normal_maps

    def render_position_multiview(
        self,
        camera_elevs: List[float],
        camera_azims: List[float],
    ) -> List[Image.Image]:
        """Render position maps from multiple viewpoints."""
        position_maps = []
        for elev, azim in zip(camera_elevs, camera_azims):
            position_map = self.render_position(elev, azim, return_type="pl")
            position_maps.append(position_map)
        return position_maps

    def _render_sketch_from_depth(self, depth_image: torch.Tensor) -> torch.Tensor:
        """Render sketch from depth using edge detection."""
        depth_image_np = depth_image.cpu().numpy()
        depth_image_np = (depth_image_np * 255).astype(np.uint8)
        depth_edges = cv2.Canny(depth_image_np, 30, 80)
        sketch_image = (
            torch.from_numpy(depth_edges).to(depth_image.device).float() / 255.0
        )
        sketch_image = sketch_image.unsqueeze(-1)
        return sketch_image

    def back_project(
        self,
        image: Union[Image.Image, np.ndarray, torch.Tensor],
        elev: float,
        azim: float,
        camera_distance: Optional[float] = None,
        center: Optional[np.ndarray] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Back-project an image onto mesh UV space."""
        if isinstance(image, Image.Image):
            image = torch.tensor(np.array(image) / 255.0)
        elif isinstance(image, np.ndarray):
            image = torch.tensor(image)
        if image.dim() == 2:
            image = image.unsqueeze(-1)
        image = image.float().to(self.device)
        resolution = image.shape[:2]
        channel = image.shape[-1]

        pos_camera, pos_clip = self._get_pos_from_mvp(
            elev, azim, camera_distance, center
        )

        rast_out = self._rasterize(pos_clip, self.pos_idx, resolution)
        visible_mask = torch.clamp(rast_out[..., -1:], 0, 1)[0, ...]

        # Compute vertex normals for angle-based weighting
        pos_camera_3d = pos_camera[:, :3] / pos_camera[:, 3:4]
        v0 = pos_camera_3d[self.pos_idx[:, 0].long(), :]
        v1 = pos_camera_3d[self.pos_idx[:, 1].long(), :]
        v2 = pos_camera_3d[self.pos_idx[:, 2].long(), :]
        face_normals = F.normalize(torch.cross(v1 - v0, v2 - v0, dim=-1), dim=-1)

        vertex_normals = trimesh.geometry.mean_vertex_normals(
            vertex_count=self.vtx_pos.shape[0],
            faces=self.pos_idx.cpu().numpy(),
            face_normals=face_normals.cpu().numpy(),
        )
        vertex_normals = (
            torch.from_numpy(vertex_normals).float().to(self.device).contiguous()
        )

        # Interpolate normals and UVs
        normal = self._interpolate(vertex_normals[None, ...], rast_out, self.pos_idx)
        normal = normal[0, ...]

        if self.vtx_uv is not None:
            uv = self._interpolate(self.vtx_uv[None, ...], rast_out, self.uv_idx)
        else:
            # No UV coordinates
            texture = torch.zeros(
                self.texture_size[1], self.texture_size[0], channel, device=self.device
            )
            cos_map = torch.zeros(
                self.texture_size[1], self.texture_size[0], 1, device=self.device
            )
            boundary_map = torch.zeros_like(cos_map)
            return texture, cos_map, boundary_map

        # Compute depth for sketch
        tex_depth = pos_camera_3d[:, 2].reshape(1, -1, 1).contiguous()
        depth = self._interpolate(tex_depth, rast_out, self.pos_idx)[0, ...]
        depth_masked = depth[visible_mask > 0]
        if depth_masked.numel() > 0:
            depth_max, depth_min = depth_masked.max(), depth_masked.min()
            depth_normalized = (depth - depth_min) / (depth_max - depth_min + 1e-8)
        else:
            depth_normalized = depth
        depth_image = depth_normalized * visible_mask

        sketch_image = self._render_sketch_from_depth(depth_image)

        # Cosine weighting
        lookat = torch.tensor([[0, 0, -1]], device=self.device)
        cos_image = torch.nn.functional.cosine_similarity(lookat, normal.view(-1, 3))
        cos_image = cos_image.view(normal.shape[0], normal.shape[1], 1)

        cos_thres = np.cos(self.bake_angle_thres / 180 * np.pi)
        cos_image[cos_image < cos_thres] = 0

        # Shrink visible mask
        kernel_size = self.bake_unreliable_kernel_size * 2 + 1
        kernel = torch.ones((1, 1, kernel_size, kernel_size), dtype=torch.float32).to(
            sketch_image.device
        )

        visible_mask_proc = visible_mask.permute(2, 0, 1).unsqueeze(0).float()
        visible_mask_proc = F.conv2d(
            1.0 - visible_mask_proc, kernel, padding=kernel_size // 2
        )
        visible_mask_proc = 1.0 - (visible_mask_proc > 0).float()
        visible_mask_proc = visible_mask_proc.squeeze(0).permute(1, 2, 0)

        sketch_proc = sketch_image.permute(2, 0, 1).unsqueeze(0)
        sketch_proc = F.conv2d(sketch_proc, kernel, padding=kernel_size // 2)
        sketch_proc = (sketch_proc > 0).float()
        sketch_proc = sketch_proc.squeeze(0).permute(1, 2, 0)
        visible_mask_proc = visible_mask_proc * (sketch_proc < 0.5)

        cos_image[visible_mask_proc == 0] = 0

        # Linear baking
        proj_mask = (visible_mask_proc != 0).view(-1)
        uv_flat = uv.squeeze(0).contiguous().view(-1, 2)[proj_mask]
        image_flat = image.squeeze(0).contiguous().view(-1, channel)[proj_mask]
        cos_flat = cos_image.contiguous().view(-1, 1)[proj_mask]
        sketch_flat = sketch_image.contiguous().view(-1, 1)[proj_mask]

        texture = linear_grid_put_2d(
            self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], image_flat
        )
        cos_map = linear_grid_put_2d(
            self.texture_size[1], self.texture_size[0], uv_flat[..., [1, 0]], cos_flat
        )
        boundary_map = linear_grid_put_2d(
            self.texture_size[1],
            self.texture_size[0],
            uv_flat[..., [1, 0]],
            sketch_flat,
        )

        return texture, cos_map, boundary_map

    def bake_from_multiview(
        self,
        views: List[Image.Image],
        camera_elevs: List[float],
        camera_azims: List[float],
        view_weights: List[float],
        method: str = "fast",
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Bake texture from multiple views."""
        project_textures, project_weighted_cos_maps = [], []
        bake_exp = 4

        for view, camera_elev, camera_azim, weight in zip(
            views, camera_elevs, camera_azims, view_weights
        ):
            project_texture, project_cos_map, _ = self.back_project(
                view, camera_elev, camera_azim
            )
            project_cos_map = weight * (project_cos_map**bake_exp)
            project_textures.append(project_texture)
            project_weighted_cos_maps.append(project_cos_map)

        if method == "fast":
            texture, ori_trust_map = self.fast_bake_texture(
                project_textures, project_weighted_cos_maps
            )
        else:
            raise ValueError(f"Unknown bake method: {method}")

        return texture, ori_trust_map > 1e-8

    @torch.no_grad()
    def fast_bake_texture(
        self,
        textures: List[torch.Tensor],
        cos_maps: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Fast texture baking by weighted averaging."""
        channel = textures[0].shape[-1]
        texture_merge = torch.zeros(self.texture_size + (channel,)).to(self.device)
        trust_map_merge = torch.zeros(self.texture_size + (1,)).to(self.device)

        for texture, cos_map in zip(textures, cos_maps):
            view_sum = (cos_map > 0).sum()
            painted_sum = ((cos_map > 0) * (trust_map_merge > 0)).sum()
            if view_sum > 0 and painted_sum / view_sum > 0.99:
                continue
            texture_merge += texture * cos_map
            trust_map_merge += cos_map

        texture_merge = texture_merge / torch.clamp(trust_map_merge, min=1e-8)
        texture_merge = texture_merge.clamp(0.0, 1.0)

        return texture_merge, trust_map_merge > 1e-8

    def texture_inpaint(
        self,
        texture: torch.Tensor,
        mask: Union[torch.Tensor, np.ndarray],
    ) -> torch.Tensor:
        """Inpaint missing regions in UV texture using mesh-aware method."""
        if isinstance(texture, torch.Tensor):
            texture_np = texture.cpu().numpy()
        else:
            texture_np = texture

        if isinstance(mask, torch.Tensor):
            mask_np = mask.cpu().numpy()
        else:
            mask_np = mask

        # Ensure proper format
        if texture_np.max() <= 1.0:
            texture_np = texture_np.astype(np.float32)
        else:
            texture_np = (texture_np / 255.0).astype(np.float32)

        if mask_np.ndim == 3:
            mask_np = mask_np.squeeze(-1)
        if mask_np.dtype == np.uint8:
            mask_uint8 = mask_np
        else:
            mask_uint8 = ((mask_np > 0) * 255).astype(np.uint8)

        # Get mesh data for mesh-aware inpainting
        vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh()

        if vtx_uv is not None and uv_idx is not None:
            texture_np, mask_uint8 = meshVerticeInpaint(
                texture_np, mask_uint8, vtx_pos, vtx_uv, pos_idx, uv_idx
            )

        # Final OpenCV inpainting for remaining holes
        texture_uint8 = (texture_np * 255).astype(np.uint8)
        inpaint_mask = 255 - mask_uint8
        texture_inpainted = cv2.inpaint(texture_uint8, inpaint_mask, 3, cv2.INPAINT_NS)

        return torch.from_numpy(texture_inpainted / 255.0).float().to(self.device)

    # Alias for compatibility
    uv_inpaint = texture_inpaint


def array_to_tensor(np_array):
    """Convert numpy array to normalized tensor."""
    image_pt = torch.tensor(np_array).float()
    image_pt = image_pt / 255 * 2 - 1
    image_pt = rearrange(image_pt, "h w c -> c h w")
    image_pts = repeat(image_pt, "c h w -> b c h w", b=1)
    return image_pts


def recenter_image(image, border_ratio=0.2):
    """Recenter a PIL image, cropping to non-transparent content with a border."""
    from PIL import Image as PILImage

    if image.mode == "RGB":
        return image
    elif image.mode == "L":
        return image.convert("RGB")
    if image.mode != "RGBA":
        image = image.convert("RGBA")

    alpha_channel = np.array(image)[:, :, 3]
    non_zero_indices = np.argwhere(alpha_channel > 0)
    if non_zero_indices.size == 0:
        raise ValueError("Image is fully transparent")

    min_row, min_col = non_zero_indices.min(axis=0)
    max_row, max_col = non_zero_indices.max(axis=0)

    cropped_image = image.crop((min_col, min_row, max_col + 1, max_row + 1))

    width, height = cropped_image.size
    border_width = int(width * border_ratio)
    border_height = int(height * border_ratio)

    new_width = width + 2 * border_width
    new_height = height + 2 * border_height
    square_size = max(new_width, new_height)

    new_image = PILImage.new("RGBA", (square_size, square_size), (255, 255, 255, 0))

    paste_x = (square_size - new_width) // 2 + border_width
    paste_y = (square_size - new_height) // 2 + border_height
    new_image.paste(cropped_image, (paste_x, paste_y))
    return new_image


class ImageProcessorV2:
    """Image processor for Hunyuan3D single-view input."""

    # External module path aliases for compatibility with Hunyuan3D configs
    _aliases = [
        "hy3dshape.preprocessors.ImageProcessorV2",
        "hy3dgen.shapegen.preprocessors.ImageProcessorV2",
    ]

    def __init__(self, size=512, border_ratio=None):
        self.size = size
        self.border_ratio = border_ratio

    @staticmethod
    def recenter(image, border_ratio: float = 0.2):
        """recenter an image to leave some empty space at the image border."""

        if image.shape[-1] == 4:
            mask = image[..., 3]
        else:
            mask = np.ones_like(image[..., 0:1]) * 255
            image = np.concatenate([image, mask], axis=-1)
            mask = mask[..., 0]

        height, width, channels = image.shape

        size = max(height, width)
        result = np.zeros((size, size, channels), dtype=np.uint8)

        coords = np.nonzero(mask)
        x_min, x_max = coords[0].min(), coords[0].max()
        y_min, y_max = coords[1].min(), coords[1].max()
        crop_h = x_max - x_min
        crop_w = y_max - y_min
        if crop_h == 0 or crop_w == 0:
            raise ValueError("input image is empty")
        desired_size = int(size * (1 - border_ratio))
        scale = desired_size / max(crop_h, crop_w)
        scaled_h = int(crop_h * scale)
        scaled_w = int(crop_w * scale)
        x2_min = (size - scaled_h) // 2
        x2_max = x2_min + scaled_h

        y2_min = (size - scaled_w) // 2
        y2_max = y2_min + scaled_w

        result[x2_min:x2_max, y2_min:y2_max] = cv2.resize(
            image[x_min:x_max, y_min:y_max],
            (scaled_w, scaled_h),
            interpolation=cv2.INTER_AREA,
        )

        bg = np.ones((result.shape[0], result.shape[1], 3), dtype=np.uint8) * 255

        mask = result[..., 3:].astype(np.float32) / 255
        result = result[..., :3] * mask + bg * (1 - mask)

        mask = mask * 255
        result = result.clip(0, 255).astype(np.uint8)
        mask = mask.clip(0, 255).astype(np.uint8)
        return result, mask

    def load_image(self, image, border_ratio=0.15, to_tensor=True):
        if isinstance(image, str):
            image = cv2.imread(image, cv2.IMREAD_UNCHANGED)
            image, mask = self.recenter(image, border_ratio=border_ratio)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif isinstance(image, Image.Image):
            image = image.convert("RGBA")
            image = np.asarray(image)
            image, mask = self.recenter(image, border_ratio=border_ratio)

        image = cv2.resize(image, (self.size, self.size), interpolation=cv2.INTER_CUBIC)
        mask = cv2.resize(mask, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
        mask = mask[..., np.newaxis]

        if to_tensor:
            image = array_to_tensor(image)
            mask = array_to_tensor(mask)
        return image, mask

    def __call__(self, image, border_ratio=0.15, to_tensor=True, **kwargs):
        if self.border_ratio is not None:
            border_ratio = self.border_ratio
        image, mask = self.load_image(
            image, border_ratio=border_ratio, to_tensor=to_tensor
        )
        outputs = {"image": image, "mask": mask}
        return outputs


class MVImageProcessorV2(ImageProcessorV2):
    """Multi-view image processor for Hunyuan3D."""

    # External module path aliases for compatibility with Hunyuan3D configs
    _aliases = [
        "hy3dshape.preprocessors.MVImageProcessorV2",
    ]

    return_view_idx = True

    def __init__(self, size=512, border_ratio=None):
        super().__init__(size, border_ratio)
        self.view2idx = {"front": 0, "left": 1, "back": 2, "right": 3}

    def __call__(self, image_dict, border_ratio=0.15, to_tensor=True, **kwargs):
        if self.border_ratio is not None:
            border_ratio = self.border_ratio

        images = []
        masks = []
        view_idxs = []
        for view_tag, image in image_dict.items():
            view_idxs.append(self.view2idx[view_tag])
            image, mask = self.load_image(
                image, border_ratio=border_ratio, to_tensor=to_tensor
            )
            images.append(image)
            masks.append(mask)

        zipped_lists = zip(view_idxs, images, masks)
        sorted_zipped_lists = sorted(zipped_lists)
        view_idxs, images, masks = zip(*sorted_zipped_lists)

        image = torch.cat(images, 0).unsqueeze(0)
        mask = torch.cat(masks, 0).unsqueeze(0)
        outputs = {"image": image, "mask": mask, "view_idxs": view_idxs}
        return outputs


# All tool classes available in this module for resolution
TOOL_CLASSES = (
    ImageProcessorV2,
    MVImageProcessorV2,
)


def resolve_hunyuan3d_tool(target: str):
    """Resolve a Hunyuan3D tool class by target string."""
    # First, try to match against _aliases
    for cls in TOOL_CLASSES:
        aliases = getattr(cls, "_aliases", [])
        if target in aliases:
            return cls

    # Then, try to match against class names
    for cls in TOOL_CLASSES:
        if cls.__name__ == target:
            return cls

    return None


__all__ = [
    "transform_pos",
    "get_mv_matrix",
    "get_orthographic_projection_matrix",
    "get_perspective_projection_matrix",
    "export_to_trimesh",
    "mesh_uv_wrap",
    "meshVerticeInpaint",
    "stride_from_shape",
    "scatter_add_nd_with_count",
    "linear_grid_put_2d",
    "MeshRender",
    "recenter_image",
    "array_to_tensor",
    "ImageProcessorV2",
    "MVImageProcessorV2",
    "TOOL_CLASSES",
    "resolve_hunyuan3d_tool",
]