| |
| |
| |
| |
| |
| |
| import os |
|
|
| import numpy as np |
| import torch |
|
|
| from core.models.rendering.base_gs_render import BaseGSRender |
| from core.models.rendering.gaussian_decoder.mlp_decoder import GSMLPDecoder |
| from core.models.rendering.utils.typing import * |
| from core.outputs.output import GaussianAppOutput |
|
|
|
|
| class GS3DRenderer(BaseGSRender): |
| def __init__( |
| self, |
| human_model_path, |
| subdivide_num, |
| smpl_type, |
| feat_dim, |
| query_dim, |
| use_rgb, |
| sh_degree, |
| xyz_offset_max_step, |
| mlp_network_config, |
| expr_param_dim, |
| shape_param_dim, |
| clip_scaling=0.2, |
| cano_pose_type=0, |
| decoder_mlp=False, |
| skip_decoder=False, |
| fix_opacity=False, |
| fix_rotation=False, |
| decode_with_extra_info=None, |
| gradient_checkpointing=False, |
| apply_pose_blendshape=False, |
| dense_sample_pts=40000, |
| gs_deform_scale=0.005, |
| render_features=False, |
| ): |
| """ |
| Initializes the GS3DRenderer, a subclass of BaseGSRender for 3D Gaussian Splatting rendering. |
| |
| Args: |
| human_model_path (str): Path to human model files. |
| subdivide_num (int): Subdivision number for base mesh. |
| smpl_type (str): Type of SMPL/SMPL-X/other model to use. |
| feat_dim (int): Dimension of feature embeddings. |
| query_dim (int): Dimension of query points/features. |
| use_rgb (bool): Whether to use RGB channels. |
| sh_degree (int): Spherical harmonics degree for appearance. |
| xyz_offset_max_step (float): Max offset per step for position. |
| mlp_network_config (dict or None): MLP configuration for feature mapping. |
| expr_param_dim (int): Expression parameter dimension. |
| shape_param_dim (int): Shape parameter dimension. |
| clip_scaling (float, optional): Output scaling for decoder. Default 0.2. |
| cano_pose_type (int, optional): Canonical pose type. Default 0. |
| decoder_mlp (bool, optional): Use MLP in decoder cross-attention. Default False. |
| skip_decoder (bool, optional): Whether to skip decoder and cross-attn layers. Default False. |
| fix_opacity (bool, optional): Fix opacity during training. Default False. |
| fix_rotation (bool, optional): Fix rotation during training. Default False. |
| decode_with_extra_info (dict or None, optional): Provide extra info to decoder. Default None. |
| gradient_checkpointing (bool, optional): Enable gradient checkpointing. Default False. |
| apply_pose_blendshape (bool, optional): Apply pose blendshape. Default False. |
| dense_sample_pts (int, optional): Dense sample points for mesh/voxel. Default 40000. |
| gs_deform_scale (float, optional): Deformation scale for Gaussian Splatting. Default 0.005. |
| render_features (bool, optional): Output additional features in renderer. Default False. |
| """ |
|
|
| super(GS3DRenderer, self).__init__( |
| human_model_path, |
| subdivide_num, |
| smpl_type, |
| feat_dim, |
| query_dim, |
| use_rgb, |
| sh_degree, |
| xyz_offset_max_step, |
| mlp_network_config, |
| expr_param_dim, |
| shape_param_dim, |
| clip_scaling, |
| cano_pose_type, |
| decoder_mlp, |
| skip_decoder, |
| fix_opacity, |
| fix_rotation, |
| decode_with_extra_info, |
| gradient_checkpointing, |
| apply_pose_blendshape, |
| dense_sample_pts, |
| gs_deform_scale, |
| render_features, |
| ) |
|
|
| self.gs_net = GSMLPDecoder( |
| in_channels=query_dim, |
| use_rgb=use_rgb, |
| sh_degree=self.sh_degree, |
| clip_scaling=clip_scaling, |
| init_scaling=-6.0, |
| init_density=0.1, |
| xyz_offset=True, |
| restrict_offset=True, |
| xyz_offset_max_step=xyz_offset_max_step, |
| fix_opacity=fix_opacity, |
| fix_rotation=fix_rotation, |
| use_fine_feat=( |
| True |
| if decode_with_extra_info is not None |
| and decode_with_extra_info["type"] is not None |
| else False |
| ), |
| ) |
|
|
| self.gs_deform_scale = gs_deform_scale |
|
|
| def hyper_step(self, step): |
| """using to adjust the constrain scale.""" |
|
|
| self.gs_net.hyper_step(step) |
|
|
| def forward_gs_attr( |
| self, x, query_points, smplx_data, debug=False, x_fine=None, mesh_meta=None |
| ): |
| """ |
| x: [N, C] Float[Tensor, "Np Cp"], |
| query_points: [N, 3] Float[Tensor, "Np 3"] |
| """ |
| device = x.device |
| if self.mlp_network_config is not None: |
| |
| x = self.mlp_net(x) |
| if x_fine is not None: |
| x_fine = self.mlp_net(x_fine) |
|
|
| |
|
|
| is_constrain_body = mesh_meta["is_constrain_body"].to( |
| self.smplx_model.is_constrain_body |
| ) |
| is_hands = (mesh_meta["is_rhand"] + mesh_meta["is_lhand"]).to( |
| self.smplx_model.is_rhand |
| ) |
| is_upper_body = mesh_meta["is_upper_body"].to(self.smplx_model.is_upper_body) |
| |
| |
| |
|
|
| constrain_dict = dict( |
| is_constrain_body=is_constrain_body, |
| is_hands=is_hands, |
| is_upper_body=is_upper_body, |
| ) |
|
|
| gs_attr: GaussianAppOutput = self.gs_net( |
| x, query_points, x_fine, constrain_dict |
| ) |
|
|
| return gs_attr |
|
|
|
|
| def test(): |
| import cv2 |
|
|
| human_model_path = "./pretrained_models/human_model_files" |
| smplx_data_root = "/data1/projects/ExAvatar_RELEASE/avatar/data/Custom/data/gyeongsik/smplx_optimized/smplx_params_smoothed" |
| shape_param_file = "/data1/projects/ExAvatar_RELEASE/avatar/data/Custom/data/gyeongsik/smplx_optimized/shape_param.json" |
|
|
| from core.models.rendering.smpl_x import read_smplx_param |
|
|
| batch_size = 1 |
| device = "cuda" |
| smplx_data, cam_param_list, ori_image_list = read_smplx_param( |
| smplx_data_root=smplx_data_root, shape_param_file=shape_param_file, batch_size=2 |
| ) |
| smplx_data_tmp = smplx_data |
| for k, v in smplx_data.items(): |
| smplx_data_tmp[k] = v.unsqueeze(0) |
| if (k == "betas") or (k == "face_offset") or (k == "joint_offset"): |
| smplx_data_tmp[k] = v[0].unsqueeze(0) |
| smplx_data = smplx_data_tmp |
|
|
| gs_render = GS3DRenderer( |
| human_model_path=human_model_path, |
| subdivide_num=2, |
| smpl_type="smplx", |
| feat_dim=64, |
| query_dim=64, |
| use_rgb=False, |
| sh_degree=3, |
| mlp_network_config=None, |
| xyz_offset_max_step=1.8 / 32, |
| ) |
|
|
| gs_render.to(device) |
| |
|
|
| c2w_list = [] |
| intr_list = [] |
| for cam_param in cam_param_list: |
| c2w = torch.eye(4).to(device) |
| c2w[:3, :3] = cam_param["R"] |
| c2w[:3, 3] = cam_param["t"] |
| c2w_list.append(c2w) |
| intr = torch.eye(4).to(device) |
| intr[0, 0] = cam_param["focal"][0] |
| intr[1, 1] = cam_param["focal"][1] |
| intr[0, 2] = cam_param["princpt"][0] |
| intr[1, 2] = cam_param["princpt"][1] |
| intr_list.append(intr) |
|
|
| c2w = torch.stack(c2w_list).unsqueeze(0) |
| intrinsic = torch.stack(intr_list).unsqueeze(0) |
|
|
| out = gs_render.forward( |
| gs_hidden_features=torch.zeros((batch_size, 2048, 64)).float().to(device), |
| query_points=None, |
| smplx_data=smplx_data, |
| c2w=c2w, |
| intrinsic=intrinsic, |
| height=int(cam_param_list[0]["princpt"][1]) * 2, |
| width=int(cam_param_list[0]["princpt"][0]) * 2, |
| background_color=torch.tensor([1.0, 1.0, 1.0]) |
| .float() |
| .view(1, 1, 3) |
| .repeat(batch_size, 2, 1) |
| .to(device), |
| debug=False, |
| ) |
|
|
| for k, v in out.items(): |
| if k == "comp_rgb_bg": |
| print("comp_rgb_bg", v) |
| continue |
| for b_idx in range(len(v)): |
| if k == "3dgs": |
| for v_idx in range(len(v[b_idx])): |
| v[b_idx][v_idx].save_ply(f"./debug_vis/{b_idx}_{v_idx}.ply") |
| continue |
| for v_idx in range(v.shape[1]): |
| save_path = os.path.join("./debug_vis", f"{b_idx}_{v_idx}_{k}.jpg") |
| cv2.imwrite( |
| save_path, |
| (v[b_idx, v_idx].detach().cpu().numpy() * 255).astype(np.uint8), |
| ) |
|
|
