File size: 2,500 Bytes
98bebfc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import os
import numpy as np
import torch
from PIL import Image
import os
from .scripts.proj_commands import projection as isomer_projection
from .data.utils import simple_remove_bkg_normal
# mesh_address,
def projection(
meshes,
masks,
images,
azimuths,
elevations,
weights,
fov,
radius,
save_dir,
save_glb_addr=None,
remove_background=False,
auto_center=False,
projection_type="perspective",
below_confidence_strategy="smooth",
complete_unseen=True,
mesh_scale_factor=1.0,
rm_bkg_with_rembg=True,
):
if save_glb_addr is None:
os.makedirs(save_dir, exist_ok=True)
save_glb_addr=os.path.join(save_dir, "rgb_projected.glb")
bs = len(images)
assert len(azimuths) == bs, f'len(azimuths) ({len(azimuths)} != batchsize ({bs}))'
assert len(elevations) == bs, f'len(elevations) ({len(elevations)} != batchsize ({bs}))'
assert len(weights) == bs, f'len(weights) ({len(weights)} != batchsize ({bs}))'
image_rgba = torch.cat([images[:,:,:,:3], masks.unsqueeze(-1)], dim=-1)
assert image_rgba.shape[-1] == 4, f'image_rgba.shape is {image_rgba.shape}'
img_list = [Image.fromarray((image.cpu()*255).numpy().astype(np.uint8)) for image in image_rgba]
if remove_background:
if rm_bkg_with_rembg:
os.environ["OMP_NUM_THREADS"] = '8'
img_list = simple_remove_bkg_normal(img_list, rm_bkg_with_rembg, return_Image=True)
resolution = img_list[0].size[0]
new_img_list = []
for i in range(len(img_list)):
new_img = img_list[i].resize((resolution,resolution))
path_dir = os.path.join(save_dir, f'projection_images')
os.makedirs(path_dir, exist_ok=True)
path_ = os.path.join(path_dir, f'ProjectionImg{i}.png')
new_img.save(path_)
new_img_list.append(new_img)
img_list = new_img_list
isomer_projection(meshes,
img_list=img_list,
weights=weights,
azimuths=azimuths,
elevations=elevations,
projection_type=projection_type,
auto_center=auto_center,
resolution=resolution,
fovy=fov,
radius=radius,
scale_factor=mesh_scale_factor,
save_glb_addr=save_glb_addr,
scale_verts=True,
complete_unseen=complete_unseen,
below_confidence_strategy=below_confidence_strategy
)
return save_glb_addr
|