| | 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 |
| |
|
| | |
| | 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 |
| |
|
| |
|
| |
|
