|
|
import os |
|
|
import imageio |
|
|
import torch |
|
|
from modules.part_synthesis.utils import render_utils, postprocessing_utils |
|
|
from modules.part_synthesis.representations.gaussian.gaussian_model import Gaussian |
|
|
|
|
|
|
|
|
def save_parts_outputs(outputs, output_dir, simplify_ratio, save_video=True, save_glb=True, textured=True): |
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
|
|
|
|
|
|
num_parts = min(len(outputs['gaussian']), len(outputs['mesh'])) |
|
|
gs_list = [] |
|
|
|
|
|
for i in range(num_parts): |
|
|
if i == 0: |
|
|
continue |
|
|
if save_video: |
|
|
video = render_utils.render_video(outputs['gaussian'][i])['color'] |
|
|
gaussian_video_path = f"{output_dir}/part{i}_gs_text.mp4" |
|
|
if os.path.exists(gaussian_video_path): |
|
|
os.remove(gaussian_video_path) |
|
|
imageio.mimsave(gaussian_video_path, video, fps=30) |
|
|
|
|
|
video = render_utils.render_video(outputs['radiance_field'][i])['color'] |
|
|
rf_video_path = f"{output_dir}/part{i}_rf_text.mp4" |
|
|
if os.path.exists(rf_video_path): |
|
|
os.remove(rf_video_path) |
|
|
imageio.mimsave(rf_video_path, video, fps=30) |
|
|
|
|
|
video = render_utils.render_video(outputs['mesh'][i])['normal'] |
|
|
mesh_video_path = f"{output_dir}/part{i}_mesh_text.mp4" |
|
|
if os.path.exists(mesh_video_path): |
|
|
os.remove(mesh_video_path) |
|
|
imageio.mimsave(mesh_video_path, video, fps=30) |
|
|
|
|
|
if save_glb: |
|
|
glb = postprocessing_utils.to_glb( |
|
|
outputs['gaussian'][i], |
|
|
outputs['mesh'][i], |
|
|
simplify=simplify_ratio, |
|
|
texture_size=1024, |
|
|
textured=textured, |
|
|
) |
|
|
if glb is None: |
|
|
continue |
|
|
glb_path = f"{output_dir}/part{i}.glb" |
|
|
if os.path.exists(glb_path): |
|
|
os.remove(glb_path) |
|
|
glb.export(glb_path) |
|
|
|
|
|
if i == 0: |
|
|
ply_path = f"{output_dir}/part{i}_gs.ply" |
|
|
if os.path.exists(ply_path): |
|
|
os.remove(ply_path) |
|
|
outputs['gaussian'][i].save_ply(ply_path) |
|
|
else: |
|
|
gs_list.append(outputs['gaussian'][i]) |
|
|
|
|
|
merged_gaussian = merge_gaussians(gs_list) |
|
|
merged_gaussian.save_ply(f"{output_dir}/merged_gs.ply") |
|
|
|
|
|
exploded_gs = exploded_gaussians(gs_list, explosion_scale=0.3) |
|
|
exploded_gs.save_ply(f"{output_dir}/exploded_gs.ply") |
|
|
|
|
|
|
|
|
def merge_gaussians(gaussians_list): |
|
|
if not gaussians_list: |
|
|
raise ValueError("gaussians_list is empty") |
|
|
|
|
|
first_gaussian = gaussians_list[0] |
|
|
merged_gaussian = Gaussian(**first_gaussian.init_params, device=first_gaussian.device) |
|
|
|
|
|
xyz_list = [] |
|
|
features_dc_list = [] |
|
|
features_rest_list = [] |
|
|
scaling_list = [] |
|
|
rotation_list = [] |
|
|
opacity_list = [] |
|
|
|
|
|
for gaussian in gaussians_list: |
|
|
if (gaussian.sh_degree != first_gaussian.sh_degree or |
|
|
not torch.allclose(gaussian.aabb, first_gaussian.aabb)): |
|
|
raise ValueError("All Gaussian objects must have the same sh_degree and aabb parameters") |
|
|
|
|
|
if gaussian._xyz is not None: |
|
|
xyz_list.append(gaussian._xyz) |
|
|
if gaussian._features_dc is not None: |
|
|
features_dc_list.append(gaussian._features_dc) |
|
|
if gaussian._features_rest is not None: |
|
|
features_rest_list.append(gaussian._features_rest) |
|
|
if gaussian._scaling is not None: |
|
|
scaling_list.append(gaussian._scaling) |
|
|
if gaussian._rotation is not None: |
|
|
rotation_list.append(gaussian._rotation) |
|
|
if gaussian._opacity is not None: |
|
|
opacity_list.append(gaussian._opacity) |
|
|
|
|
|
if xyz_list: |
|
|
merged_gaussian._xyz = torch.cat(xyz_list, dim=0) |
|
|
if features_dc_list: |
|
|
merged_gaussian._features_dc = torch.cat(features_dc_list, dim=0) |
|
|
if features_rest_list: |
|
|
merged_gaussian._features_rest = torch.cat(features_rest_list, dim=0) |
|
|
else: |
|
|
merged_gaussian._features_rest = None |
|
|
if scaling_list: |
|
|
merged_gaussian._scaling = torch.cat(scaling_list, dim=0) |
|
|
if rotation_list: |
|
|
merged_gaussian._rotation = torch.cat(rotation_list, dim=0) |
|
|
if opacity_list: |
|
|
merged_gaussian._opacity = torch.cat(opacity_list, dim=0) |
|
|
|
|
|
return merged_gaussian |
|
|
|
|
|
|
|
|
def exploded_gaussians(gaussians_list, explosion_scale=0.4): |
|
|
|
|
|
if not gaussians_list: |
|
|
raise ValueError("gaussians_list is empty") |
|
|
|
|
|
first_gaussian = gaussians_list[0] |
|
|
merged_gaussian = Gaussian(**first_gaussian.init_params, device=first_gaussian.device) |
|
|
|
|
|
xyz_list = [] |
|
|
features_dc_list = [] |
|
|
features_rest_list = [] |
|
|
scaling_list = [] |
|
|
rotation_list = [] |
|
|
opacity_list = [] |
|
|
|
|
|
all_centers = [] |
|
|
for gaussian in gaussians_list: |
|
|
if gaussian._xyz is not None: |
|
|
center = gaussian.get_xyz.mean(dim=0) |
|
|
all_centers.append(center) |
|
|
|
|
|
if not all_centers: |
|
|
raise ValueError("No valid gaussians with xyz data found") |
|
|
|
|
|
all_centers = torch.stack(all_centers) |
|
|
global_center = all_centers.mean(dim=0) |
|
|
|
|
|
for i, gaussian in enumerate(gaussians_list): |
|
|
if (gaussian.sh_degree != first_gaussian.sh_degree or |
|
|
not torch.allclose(gaussian.aabb, first_gaussian.aabb)): |
|
|
raise ValueError("All Gaussian objects must have the same sh_degree and aabb parameters") |
|
|
|
|
|
if i < len(all_centers): |
|
|
part_center = all_centers[i] |
|
|
direction = part_center - global_center |
|
|
direction_norm = torch.norm(direction) |
|
|
if direction_norm > 1e-6: |
|
|
direction = direction / direction_norm |
|
|
else: |
|
|
direction = torch.randn(3, device=gaussian.device) |
|
|
direction = direction / torch.norm(direction) |
|
|
|
|
|
offset = direction * explosion_scale |
|
|
else: |
|
|
offset = torch.zeros(3, device=gaussian.device) |
|
|
|
|
|
if gaussian._xyz is not None: |
|
|
original_xyz = gaussian.get_xyz |
|
|
exploded_xyz = original_xyz + offset |
|
|
exploded_xyz_normalized = (exploded_xyz - gaussian.aabb[None, :3]) / gaussian.aabb[None, 3:] |
|
|
xyz_list.append(exploded_xyz_normalized) |
|
|
|
|
|
if gaussian._features_dc is not None: |
|
|
features_dc_list.append(gaussian._features_dc) |
|
|
if gaussian._features_rest is not None: |
|
|
features_rest_list.append(gaussian._features_rest) |
|
|
if gaussian._scaling is not None: |
|
|
scaling_list.append(gaussian._scaling) |
|
|
if gaussian._rotation is not None: |
|
|
rotation_list.append(gaussian._rotation) |
|
|
if gaussian._opacity is not None: |
|
|
opacity_list.append(gaussian._opacity) |
|
|
|
|
|
if xyz_list: |
|
|
merged_gaussian._xyz = torch.cat(xyz_list, dim=0) |
|
|
if features_dc_list: |
|
|
merged_gaussian._features_dc = torch.cat(features_dc_list, dim=0) |
|
|
if features_rest_list: |
|
|
merged_gaussian._features_rest = torch.cat(features_rest_list, dim=0) |
|
|
else: |
|
|
merged_gaussian._features_rest = None |
|
|
if scaling_list: |
|
|
merged_gaussian._scaling = torch.cat(scaling_list, dim=0) |
|
|
if rotation_list: |
|
|
merged_gaussian._rotation = torch.cat(rotation_list, dim=0) |
|
|
if opacity_list: |
|
|
merged_gaussian._opacity = torch.cat(opacity_list, dim=0) |
|
|
|
|
|
return merged_gaussian |
