|
import os |
|
import numpy as np |
|
import torch |
|
from PIL import Image |
|
import os |
|
from .model.inference_pipeline import reconstruction_pipe |
|
|
|
def reconstruction( |
|
normal_pils, |
|
masks, |
|
weights, |
|
fov, |
|
radius, |
|
camera_angles_azi, |
|
camera_angles_ele, |
|
expansion_weight_stage1=0.1, |
|
init_type="ball", |
|
init_verts=None, |
|
init_faces=None, |
|
init_mesh_from_file="", |
|
stage1_steps=200, |
|
stage2_steps=200, |
|
projection_type="perspective", |
|
need_normal_rotation=False, |
|
rotation_angles_azi=None, |
|
rotation_angles_ele=None, |
|
normal_rotation_R=None, |
|
rm_bkg=False, |
|
rm_bkg_with_rembg=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, |
|
expansion_weight_stage2=0.0, |
|
): |
|
|
|
if init_type == "file": |
|
assert ((init_verts is not None and init_faces is not None) or init_mesh_from_file), f'init_mesh_from_file or (init_verts and init_faces) must be provided if init_type=="file"' |
|
|
|
if not need_normal_rotation: |
|
rotation_angles_azi = None |
|
rotation_angles_ele = None |
|
normal_rotation_R = None |
|
|
|
bs = len(normal_pils) |
|
|
|
assert len(camera_angles_azi) == bs, f'len(camera_angles_azi) ({len(camera_angles_azi)} != batchsize ({bs}))' |
|
assert len(camera_angles_ele) == bs, f'len(camera_angles_ele) ({len(camera_angles_ele)} != batchsize ({bs}))' |
|
|
|
normal_pils_rgba = torch.cat([normal_pils[:,:,:,:3], masks.unsqueeze(-1)], dim=-1) |
|
|
|
assert normal_pils_rgba.shape[-1] == 4, f'normal_pils_rgba.shape is {normal_pils_rgba.shape}' |
|
|
|
|
|
normal_pils = [Image.fromarray((normal_pil.cpu()*255).numpy().astype(np.uint8)) for normal_pil in normal_pils_rgba] |
|
|
|
|
|
meshes = reconstruction_pipe( |
|
normal_pils=normal_pils, |
|
rotation_angles_azi=rotation_angles_azi, |
|
rotation_angles_ele=rotation_angles_ele, |
|
weights=weights, |
|
expansion_weight=expansion_weight_stage1, |
|
init_type=init_type, |
|
stage1_steps=stage1_steps, |
|
stage2_steps=stage2_steps, |
|
projection_type=projection_type, |
|
fovy=fov, |
|
radius=radius, |
|
camera_angles_azi=camera_angles_azi, |
|
camera_angles_ele=camera_angles_ele, |
|
rm_bkg=rm_bkg, rm_bkg_with_rembg=rm_bkg_with_rembg, |
|
normal_rotation_R=normal_rotation_R, |
|
init_mesh_from_file=init_mesh_from_file, |
|
start_edge_len_stage1=start_edge_len_stage1, |
|
end_edge_len_stage1=end_edge_len_stage1, |
|
start_edge_len_stage2=start_edge_len_stage2, |
|
end_edge_len_stage2=end_edge_len_stage2, |
|
expansion_weight_stage2=expansion_weight_stage2, |
|
init_verts=init_verts, |
|
init_faces=init_faces, |
|
|
|
) |
|
|
|
|
|
return meshes |
|
|
|
|
|
|
