import os
import numpy as np
import torch
from PIL import Image
from pytorch3d.structures import Meshes
from pytorch3d.renderer import TexturesVertex
from ..scripts.fast_geo import fast_geo, create_sphere, create_box
from ..scripts.project_mesh import get_cameras_list_azi_ele
from ..mesh_reconstruction.recon import reconstruct_stage1
from ..mesh_reconstruction.refine import run_mesh_refine
from ..mesh_reconstruction.func import make_star_cameras_orthographic, make_star_cameras_perspective
from ..data.utils import (
simple_remove_bkg_normal,
load_glb,
load_obj_with_verts_faces)
from ..scripts.utils import (
to_pyml_mesh,
simple_clean_mesh,
normal_rotation_img2img_c2w,
rotate_normal_R,
get_rotation_matrix_azi_ele,
manage_elevation_azimuth)
@torch.enable_grad()
def reconstruction_pipe(normal_pils,
rotation_angles_azi,
rotation_angles_ele,
front_index=0,
back_index=2,
side_index=1,
weights=None,
expansion_weight=0.1,
expansion_weight_stage2=0.0,
init_type="ball",
sphere_r=None, # only used if init_type=="ball"
box_width=1.0, # only used if init_type=="box"
box_length=1.0, # only used if init_type=="box"
box_height=1.0, # only used if init_type=="box"
init_verts=None,
init_faces=None,
init_mesh_from_file="",
stage1_steps=200,
stage2_steps=200,
projection_type="orthographic",
fovy=None,
radius=None,
ortho_dist=1.1,
camera_angles_azi=None,
camera_angles_ele=None,
rm_bkg=False,
rm_bkg_with_rembg=False, # only used if rm_bkg
normal_rotation_R=None,
train_stage1=True,
train_stage2=True,
use_remesh_stage1=True,
use_remesh_stage2=True,
start_edge_len_stage1=0.1,
end_edge_len_stage1=0.02,
start_edge_len_stage2=0.02,
end_edge_len_stage2=0.005,
):
assert projection_type in ['perspective', 'orthographic'], f"projection_type ({projection_type}) should be one of ['perspective', 'orthographic']"
if stage1_steps == 0:
train_stage1 = False
if stage2_steps == 0:
train_stage2 = False
if normal_rotation_R is not None:
assert normal_rotation_R.shape[-2] == 3 and normal_rotation_R.shape[-1] == 3
assert len(normal_rotation_R.shape) == 2
normal_rotation_R = normal_rotation_R.float()
camera_angles_azi = camera_angles_azi.float()
camera_angles_ele = camera_angles_ele.float()
camera_angles_ele, camera_angles_azi = manage_elevation_azimuth(camera_angles_ele, camera_angles_azi)
if init_type in ["std", "thin"]:
assert camera_angles_azi[front_index]%360==0, f"the camera_angles_azi associated with front image (index {front_index}) should be 0 not {camera_angles_azi[front_index]}"
assert camera_angles_azi[back_index]%360==180, f"the camera_angles_azi associated with back image (index {back_index}) should be 180 not {camera_angles_azi[back_index]}"
assert camera_angles_azi[side_index]%360==90, f"the camera_angles_azi associated with left side image (index {side_index}) should be 90, not {camera_angles_azi[back_index]}"
if rm_bkg:
if rm_bkg_with_rembg:
os.environ["OMP_NUM_THREADS"] = '8'
normal_pils = simple_remove_bkg_normal(normal_pils,rm_bkg_with_rembg)
if rotation_angles_azi is not None:
rotation_angles_azi = -rotation_angles_azi.float()
rotation_angles_ele = rotation_angles_ele.float()
rotation_angles_ele, rotation_angles_azi = manage_elevation_azimuth(rotation_angles_ele, rotation_angles_azi)
assert len(normal_pils) == len(rotation_angles_azi), f'len(normal_pils) ({len(normal_pils)}) != len(rotation_angles_azi) ({len(rotation_angles_azi)})'
if rotation_angles_ele is None:
rotation_angles_ele = [0] * len(normal_pils)
normal_pils_rotated = []
for i in range(len(normal_pils)):
c2w_R = get_rotation_matrix_azi_ele(rotation_angles_azi[i], rotation_angles_ele[i])
rotated_ = normal_rotation_img2img_c2w(normal_pils[i], c2w=c2w_R)
normal_pils_rotated.append(rotated_)
normal_pils = normal_pils_rotated
if normal_rotation_R is not None:
normal_pils_rotated = []
for i in range(len(normal_pils)):
rotated_ = rotate_normal_R(normal_pils[i], normal_rotation_R, save_addr="", device="cuda")
normal_pils_rotated.append(rotated_)
normal_pils = normal_pils_rotated
normal_stg1 = [img for img in normal_pils]
if init_type in ['thin', 'std']:
front_ = normal_stg1[front_index]
back_ = normal_stg1[back_index]
side_ = normal_stg1[side_index]
meshes, depth_front, depth_back, mesh_front, mesh_back = fast_geo(front_, back_, side_, init_type=init_type, return_depth_and_sep_mesh=True)
elif init_type in ["ball", "box"]:
if init_type == "ball":
assert sphere_r is not None, f"sphere_r ({sphere_r}) should not be None when init_type is 'ball'"
meshes = create_sphere(sphere_r)
if init_type == "box":
assert box_width is not None and box_length is not None and box_height is not None, f"box_width ({box_width}), box_length ({box_length}), and box_height ({box_height}) should not be None when init_type is 'box'"
meshes = create_box(width=box_width, length=box_length, height=box_height)
# add texture just in case
num_meshes = len(meshes)
num_verts_per_mesh = meshes.verts_packed().shape[0] // num_meshes
black_texture = torch.zeros((num_meshes, num_verts_per_mesh, 3), device="cuda")
textures = TexturesVertex(verts_features=black_texture)
meshes.textures = textures
elif init_type == "file":
assert init_mesh_from_file or (init_verts is not None and init_faces is not None), f"init_mesh_from_file ({init_mesh_from_file}) should not be None when init_type is 'file', else init_verts and init_faces should not be None"
if init_verts is not None and init_faces is not None:
meshes = Meshes(verts=[init_verts], faces=[init_faces]).to('cuda')
elif init_mesh_from_file.endswith('.glb'):
meshes = load_glb(init_mesh_from_file).to('cuda')
else:
meshes = load_obj_with_verts_faces(init_mesh_from_file).to('cuda')
# add texture just in case
num_meshes = len(meshes)
num_verts_per_mesh = meshes.verts_packed().shape[0] // num_meshes
black_texture = torch.zeros((num_meshes, num_verts_per_mesh, 3), device="cuda")
textures = TexturesVertex(verts_features=black_texture)
meshes.textures = textures
if projection_type == 'perspective':
assert fovy is not None and radius is not None, f"fovy ({fovy}) and radius ({radius}) should not be None when projection_type is 'perspective'"
cameras = get_cameras_list_azi_ele(camera_angles_azi, camera_angles_ele, fov_in_degrees=fovy,device="cuda", dist=radius, cam_type='fov')
elif projection_type == 'orthographic':
cameras = get_cameras_list_azi_ele(camera_angles_azi, camera_angles_ele, fov_in_degrees=fovy, device="cuda", focal=1., dist=ortho_dist, cam_type='orthographic')
vertices, faces = meshes.verts_list()[0], meshes.faces_list()[0]
render_camera_angles_azi = -camera_angles_azi
render_camera_angles_ele = camera_angles_ele
if projection_type == 'orthographic':
mv, proj = make_star_cameras_orthographic(render_camera_angles_azi, render_camera_angles_ele)
else:
mv, proj = make_star_cameras_perspective(render_camera_angles_azi, render_camera_angles_ele, distance=radius, r=radius, fov=fovy, device='cuda')
# stage 1
if train_stage1:
vertices, faces = reconstruct_stage1(normal_stg1, mv=mv, proj=proj, steps=stage1_steps, vertices=vertices, faces=faces, start_edge_len=start_edge_len_stage1, end_edge_len=end_edge_len_stage1, gain=0.05, return_mesh=False, loss_expansion_weight=expansion_weight, use_remesh=use_remesh_stage1)
# stage 2
if train_stage2:
vertices, faces = run_mesh_refine(vertices, faces, normal_pils, mv=mv, proj=proj, weights=weights, steps=stage2_steps, start_edge_len=start_edge_len_stage2, end_edge_len=end_edge_len_stage2, decay=0.99, update_normal_interval=20, update_warmup=5, return_mesh=False, process_inputs=False, process_outputs=False, cameras=cameras, use_remesh=use_remesh_stage2, loss_expansion_weight=expansion_weight_stage2)
meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=1, apply_sub_divide=True, sub_divide_threshold=0.25).to("cuda")
return meshes