visualize_mesh.py

import os
import json
import argparse

import cv2
import open3d
import trimesh
import numpy as np
from panda3d.core import Triangulator

from misc.panorama import xyz_2_coorxy
from visualize_3d import convert_lines_to_vertices

from typing import List, Tuple, Dict, Any, Union


def E2P(image, corner_i, corner_j, wall_height, camera_center, resolution=512, is_wall=True):
    """convert panorama to persepctive image
    """
    corner_i = corner_i - camera_center
    corner_j = corner_j - camera_center

    if is_wall:
        xs = np.linspace(corner_i[0], corner_j[0], resolution)[None].repeat(resolution, 0)
        ys = np.linspace(corner_i[1], corner_j[1], resolution)[None].repeat(resolution, 0)
        zs = np.linspace(-camera_center[-1], wall_height - camera_center[-1], resolution)[:, None].repeat(resolution, 1)
    else:
        xs = np.linspace(corner_i[0], corner_j[0], resolution)[None].repeat(resolution, 0)
        ys = np.linspace(corner_i[1], corner_j[1], resolution)[:, None].repeat(resolution, 1)
        zs = np.zeros_like(xs) + wall_height - camera_center[-1]

    coorx, coory = xyz_2_coorxy(xs, ys, zs)

    persp = cv2.remap(image,
                      coorx.astype(np.float32),
                      coory.astype(np.float32),
                      cv2.INTER_CUBIC,
                      borderMode=cv2.BORDER_WRAP)

    return persp


def create_plane_mesh(
    vertices,
    vertices_floor,
    textures,
    texture_floor,
    texture_ceiling,
    delta_height,
    ignore_ceiling=False,
    camera_center=None,
):
    # create mesh for 3D floorplan visualization
    triangles = []
    triangle_uvs = []

    # the number of vertical walls
    num_walls = len(vertices)

    # 1. vertical wall (always rectangle)
    num_vertices = 0
    for i in range(len(vertices)):
        # hardcode triangles for each vertical wall
        triangle = np.array([[0, 2, 1], [2, 0, 3]])
        triangles.append(triangle + num_vertices)
        num_vertices += 4

        triangle_uv = np.array([[i / (num_walls + 2), 0], [i / (num_walls + 2), 1], [(i + 1) / (num_walls + 2), 1],
                                [(i + 1) / (num_walls + 2), 0]],
                               dtype=np.float32)
        triangle_uvs.append(triangle_uv)

    # 2. floor and ceiling
    # Since the floor and ceiling may not be a rectangle, triangulate the polygon first.
    tri = Triangulator()
    for i in range(len(vertices_floor)):
        tri.add_vertex(vertices_floor[i, 0], vertices_floor[i, 1])

    for i in range(len(vertices_floor)):
        tri.add_polygon_vertex(i)

    tri.triangulate()

    # polygon triangulation
    triangle = []
    for i in range(tri.getNumTriangles()):
        triangle.append([tri.get_triangle_v0(i), tri.get_triangle_v1(i), tri.get_triangle_v2(i)])
    triangle = np.array(triangle)

    # add triangles for floor and ceiling
    triangles.append(triangle + num_vertices)
    num_vertices += len(np.unique(triangle))
    if not ignore_ceiling:
        triangles.append(triangle + num_vertices)

    # texture for floor and ceiling
    vertices_floor_min = np.min(vertices_floor[:, :2], axis=0)
    vertices_floor_max = np.max(vertices_floor[:, :2], axis=0)

    # normalize to [0, 1]
    triangle_uv = (vertices_floor[:, :2] - vertices_floor_min) / (vertices_floor_max - vertices_floor_min)
    triangle_uv[:, 0] = (triangle_uv[:, 0] + num_walls) / (num_walls + 2)

    triangle_uvs.append(triangle_uv)

    # normalize to [0, 1]
    triangle_uv = (vertices_floor[:, :2] - vertices_floor_min) / (vertices_floor_max - vertices_floor_min)
    triangle_uv[:, 0] = (triangle_uv[:, 0] + num_walls + 1) / (num_walls + 2)

    triangle_uvs.append(triangle_uv)

    # 3. Merge wall, floor, and ceiling
    vertices.append(vertices_floor)
    vertices.append(vertices_floor + delta_height)
    vertices = np.concatenate(vertices, axis=0)

    triangles = np.concatenate(triangles, axis=0)

    textures.append(texture_floor)
    textures.append(texture_ceiling)
    textures = np.concatenate(textures, axis=1)

    triangle_uvs = np.concatenate(triangle_uvs, axis=0)

    # 4. create mesh
    vertices_in_cam = (vertices - camera_center) * 0.001 if camera_center is not None else vertices * 0.001
    mesh = open3d.geometry.TriangleMesh(
        # convert vertices from mm to m
        vertices=open3d.utility.Vector3dVector(vertices_in_cam),
        triangles=open3d.utility.Vector3iVector(triangles))
    mesh.compute_vertex_normals()
    # mesh.textures = [open3d.geometry.Image(textures)]
    # mesh.triangle_uvs = np.array(triangle_uvs[triangles.reshape(-1), :], dtype=np.float64)
    return mesh


def verify_normal(corner_i, corner_j, delta_height, plane_normal):
    edge_a = corner_j + delta_height - corner_i
    edge_b = delta_height

    normal = np.cross(edge_a, edge_b)
    normal /= np.linalg.norm(normal, ord=2)

    inner_product = normal.dot(plane_normal)

    if inner_product > 1e-8:
        return False
    else:
        return True


def create_spatial_quad_polygen(quad_vertices: List, normal: np.array, camera_center: np.array):
    """create a quad polygen for spatial mesh
    """
    if camera_center is None:
        camera_center = np.array([0, 0, 0])
    quad_vertices = (quad_vertices - camera_center)
    quad_triangles = []
    triangle = np.array([[0, 2, 1], [2, 0, 3]])
    quad_triangles.append(triangle)

    quad_triangles = np.concatenate(quad_triangles, axis=0)

    mesh = trimesh.Trimesh(vertices=quad_vertices,
                           faces=quad_triangles,
                           vertex_normals=np.tile(normal, (4, 1)),
                           process=False)

    centroid = np.mean(quad_vertices, axis=0)
    # print(f'centroid: {centroid}')
    normal_point = centroid + np.array(normal) * 0.5
    # print(f'normal_point: {normal_point}')

    pcd_o3d = open3d.geometry.PointCloud()
    pcd_o3d.points = open3d.utility.Vector3dVector(np.asarray(mesh.vertices))
    pcd_o3d.points.append(normal_point)
    pcd_o3d.points.append(centroid)

    # pcd = trimesh.PointCloud(np.asarray(pcd.points))
    return mesh, pcd_o3d


def visualize_mesh(args):
    """visualize as water-tight mesh
    """

    image = cv2.imread(
        os.path.join(args.path, f"scene_{args.scene:05d}", "2D_rendering", str(args.room),
                     "panorama/full/rgb_rawlight.png"))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # load room annotations
    with open(os.path.join(args.path, f"scene_{args.scene:05d}", "annotation_3d.json")) as f:
        annos = json.load(f)

    # load camera info
    camera_center = np.loadtxt(
        os.path.join(args.path, f"scene_{args.scene:05d}", "2D_rendering", str(args.room), "panorama",
                     "camera_xyz.txt"))

    # parse corners
    junctions = np.array([item['coordinate'] for item in annos['junctions']])
    lines_holes = []
    for semantic in annos['semantics']:
        if semantic['type'] in ['window', 'door']:
            for planeID in semantic['planeID']:
                lines_holes.extend(np.where(np.array(annos['planeLineMatrix'][planeID]))[0].tolist())

    lines_holes = np.unique(lines_holes)
    _, vertices_holes = np.where(np.array(annos['lineJunctionMatrix'])[lines_holes])
    vertices_holes = np.unique(vertices_holes)

    # parse annotations
    walls = dict()
    walls_normal = dict()
    for semantic in annos['semantics']:
        if semantic['ID'] != int(args.room):
            continue

        # find junctions of ceiling and floor
        for planeID in semantic['planeID']:
            plane_anno = annos['planes'][planeID]

            if plane_anno['type'] != 'wall':
                lineIDs = np.where(np.array(annos['planeLineMatrix'][planeID]))[0]
                lineIDs = np.setdiff1d(lineIDs, lines_holes)
                junction_pairs = [
                    np.where(np.array(annos['lineJunctionMatrix'][lineID]))[0].tolist() for lineID in lineIDs
                ]
                wall = convert_lines_to_vertices(junction_pairs)
                walls[plane_anno['type']] = wall[0]

        # save normal of the vertical walls
        for planeID in semantic['planeID']:
            plane_anno = annos['planes'][planeID]

            if plane_anno['type'] == 'wall':
                lineIDs = np.where(np.array(annos['planeLineMatrix'][planeID]))[0]
                lineIDs = np.setdiff1d(lineIDs, lines_holes)
                junction_pairs = [
                    np.where(np.array(annos['lineJunctionMatrix'][lineID]))[0].tolist() for lineID in lineIDs
                ]
                wall = convert_lines_to_vertices(junction_pairs)
                walls_normal[tuple(np.intersect1d(wall, walls['floor']))] = plane_anno['normal']

    # we assume that zs of floor equals 0, then the wall height is from the ceiling
    wall_height = np.mean(junctions[walls['ceiling']], axis=0)[-1]
    delta_height = np.array([0, 0, wall_height])

    # list of corner index
    wall_floor = walls['floor']

    corners = []  # 3D coordinate for each wall
    textures = []  # texture for each wall

    # wall
    quad_wall_lst = []
    for i, j in zip(wall_floor, np.roll(wall_floor, shift=-1)):
        corner_i, corner_j = junctions[i], junctions[j]
        plane_normal = walls_normal[tuple(sorted([i, j]))]
        flip = verify_normal(corner_i, corner_j, delta_height, plane_normal)

        if flip:
            corner_j, corner_i = corner_i, corner_j

        texture = E2P(image, corner_i, corner_j, wall_height, camera_center)

        quad_corner = np.array([corner_i, corner_i + delta_height, corner_j + delta_height, corner_j])
        print(f'plane normal: {plane_normal}')
        quad_polygen_mesh, quad_polygen_pcd = create_spatial_quad_polygen(quad_corner * 0.001, plane_normal,
                                                                          camera_center * 0.001)
        quad_polygen_mesh = open3d.geometry.TriangleMesh(
            # convert vertices from mm to m
            vertices=open3d.utility.Vector3dVector(quad_polygen_mesh.vertices),
            triangles=open3d.utility.Vector3iVector(quad_polygen_mesh.faces))
        open3d.visualization.draw_geometries([quad_polygen_mesh, quad_polygen_pcd])
        quad_wall_lst.append(quad_polygen_mesh)

        corners.append(quad_corner)
        textures.append(texture)

    # floor and ceiling
    # the floor/ceiling texture is cropped by the maximum bounding box
    corner_floor = junctions[wall_floor]
    corner_min = np.min(corner_floor, axis=0)
    corner_max = np.max(corner_floor, axis=0)
    texture_floor = E2P(image, corner_min, corner_max, 0, camera_center, is_wall=False)
    texture_ceiling = E2P(image, corner_min, corner_max, wall_height, camera_center, is_wall=False)

    # create mesh
    mesh = create_plane_mesh(corners,
                             corner_floor,
                             textures,
                             texture_floor,
                             texture_ceiling,
                             delta_height,
                             ignore_ceiling=args.ignore_ceiling,
                             camera_center=camera_center)
    open3d.io.write_triangle_mesh(
        "/home/hkust/fangchuan/codes/Structured3D/sample_dataset_visualization/scene_00001_906323.ply", mesh)

    quad_meshes = open3d.geometry.TriangleMesh()

    for mesh in quad_wall_lst:
        quad_meshes += mesh
    open3d.io.write_triangle_mesh(
        "/home/hkust/fangchuan/codes/Structured3D/sample_dataset_visualization/scene_00001_906323_quadwall.ply",
        quad_meshes)
    # visualize mesh
    open3d.visualization.draw_geometries(quad_wall_lst)


def parse_args():
    parser = argparse.ArgumentParser(description="Structured3D 2D Layout Visualization")
    parser.add_argument("--path", required=True, help="dataset path", metavar="DIR")
    parser.add_argument("--scene", required=True, help="scene id", type=int)
    parser.add_argument("--room", required=True, help="room id", type=int)
    parser.add_argument("--ignore_ceiling", action='store_true', help="ignore ceiling for better visualization")
    return parser.parse_args()


def main():
    args = parse_args()

    visualize_mesh(args)


if __name__ == "__main__":
    main()