Spaces:
Build error
Build error
""" | |
@author: Zhigang Jiang | |
@time: 2022/05/25 | |
@description: reference: https://github.com/sunset1995/PanoPlane360/blob/main/vis_planes.py | |
""" | |
import open3d | |
import numpy as np | |
from utils.conversion import pixel2lonlat | |
def create_3d_obj(img, depth, save_path=None, mesh=True, mesh_show_back_face=False, show=False): | |
assert img.shape[0] == depth.shape[0], "" | |
h = img.shape[0] | |
w = img.shape[1] | |
# Project to 3d | |
lon = pixel2lonlat(np.array(range(w)), w=w, axis=0)[None].repeat(h, axis=0) | |
lat = pixel2lonlat(np.array(range(h)), h=h, axis=1)[..., None].repeat(w, axis=1) | |
z = depth * np.sin(lat) | |
x = depth * np.cos(lat) * np.cos(lon) | |
y = depth * np.cos(lat) * np.sin(lon) | |
pts_xyz = np.stack([x, -z, y], -1).reshape(-1, 3) | |
pts_rgb = img.reshape(-1, 3) | |
if mesh: | |
pid = np.arange(len(pts_xyz)).reshape(h, w) | |
faces = np.concatenate([ | |
np.stack([ | |
pid[:-1, :-1], pid[1:, :-1], np.roll(pid, -1, axis=1)[:-1, :-1], | |
], -1), | |
np.stack([ | |
pid[1:, :-1], np.roll(pid, -1, axis=1)[1:, :-1], np.roll(pid, -1, axis=1)[:-1, :-1], | |
], -1) | |
]).reshape(-1, 3).tolist() | |
scene = open3d.geometry.TriangleMesh() | |
scene.vertices = open3d.utility.Vector3dVector(pts_xyz) | |
scene.vertex_colors = open3d.utility.Vector3dVector(pts_rgb) | |
scene.triangles = open3d.utility.Vector3iVector(faces) | |
else: | |
scene = open3d.geometry.PointCloud() | |
scene.points = open3d.utility.Vector3dVector(pts_xyz) | |
scene.colors = open3d.utility.Vector3dVector(pts_rgb) | |
if save_path: | |
open3d.io.write_triangle_mesh(save_path, scene, write_triangle_uvs=True) | |
if show: | |
open3d.visualization.draw_geometries([scene], mesh_show_back_face=mesh_show_back_face) | |
if __name__ == '__main__': | |
from dataset.mp3d_dataset import MP3DDataset | |
from utils.boundary import depth2boundaries, layout2depth | |
from visualization.boundary import draw_boundaries | |
mp3d_dataset = MP3DDataset(root_dir='../src/dataset/mp3d', mode='train', for_test_index=10, patch_num=1024) | |
gt = mp3d_dataset.__getitem__(3) | |
boundary_list = depth2boundaries(gt['ratio'], gt['depth'], step=None) | |
pano_img = draw_boundaries(gt['image'].transpose(1, 2, 0), boundary_list=boundary_list, show=True) | |
layout_depth = layout2depth(boundary_list, show=False) | |
create_3d_obj(gt['image'].transpose(1, 2, 0), layout_depth, save_path=f"../src/output/{gt['id']}_3d.gltf", | |
mesh=True) | |