data_preprocess/stru3d/PointCloudReaderPanorama.py

import os

import cv2
import numpy as np
import open3d as o3d

NUM_SECTIONS = -1


class PointCloudReaderPanorama:
    def __init__(self, path, resolution="full", random_level=0, generate_color=False, generate_normal=False):
        self.path = path
        self.random_level = random_level
        self.resolution = resolution
        self.generate_color = generate_color
        self.generate_normal = generate_normal
        sections = [p for p in os.listdir(os.path.join(path, "2D_rendering"))]
        self.depth_paths = [
            os.path.join(*[path, "2D_rendering", p, "panorama", self.resolution, "depth.png"]) for p in sections
        ]
        self.rgb_paths = [
            os.path.join(*[path, "2D_rendering", p, "panorama", self.resolution, "rgb_coldlight.png"])
            for p in sections
        ]
        self.normal_paths = [
            os.path.join(*[path, "2D_rendering", p, "panorama", self.resolution, "normal.png"]) for p in sections
        ]
        self.camera_paths = [os.path.join(*[path, "2D_rendering", p, "panorama", "camera_xyz.txt"]) for p in sections]
        self.camera_centers = self.read_camera_center()
        self.point_cloud = self.generate_point_cloud(
            self.random_level, color=self.generate_color, normal=self.generate_normal
        )

    def read_camera_center(self):
        camera_centers = []
        for i in range(len(self.camera_paths)):
            with open(self.camera_paths[i], "r") as f:
                line = f.readline()
            center = list(map(float, line.strip().split(" ")))
            camera_centers.append(np.asarray([center[0], center[1], center[2]]))
        return camera_centers

    def generate_point_cloud(self, random_level=0, color=False, normal=False):
        coords = []
        colors = []
        points = {}
        # normals = []

        # Getting Coordinates
        for i in range(len(self.depth_paths)):
            depth_img = cv2.imread(self.depth_paths[i], cv2.IMREAD_ANYDEPTH | cv2.IMREAD_ANYCOLOR)
            x_tick = 180.0 / depth_img.shape[0]
            y_tick = 360.0 / depth_img.shape[1]

            rgb_img = cv2.imread(self.rgb_paths[i])
            rgb_img = cv2.cvtColor(rgb_img, code=cv2.COLOR_BGR2RGB)
            # normal_img = cv2.imread(self.normal_paths[i])

            for x in range(0, depth_img.shape[0]):
                for y in range(0, depth_img.shape[1]):
                    # need 90 - -09
                    alpha = 90 - (x * x_tick)
                    beta = y * y_tick - 180

                    depth = depth_img[x, y] + np.random.random() * random_level

                    if depth > 500.0:
                        z_offset = depth * np.sin(np.deg2rad(alpha))
                        xy_offset = depth * np.cos(np.deg2rad(alpha))
                        x_offset = xy_offset * np.sin(np.deg2rad(beta))
                        y_offset = xy_offset * np.cos(np.deg2rad(beta))
                        point = np.asarray([x_offset, y_offset, z_offset])
                        coords.append(point + self.camera_centers[i])
                        colors.append(rgb_img[x, y])
                        # normals.append(normalize(normal_img[x, y].reshape(-1, 1)).ravel())

        coords = np.asarray(coords)
        colors = np.asarray(colors) / 255.0
        # normals = np.asarray(normals)

        coords[:, :2] = np.round(coords[:, :2] / 10) * 10.0
        coords[:, 2] = np.round(coords[:, 2] / 100) * 100.0
        unique_coords, unique_ind = np.unique(coords, return_index=True, axis=0)

        coords = coords[unique_ind]
        colors = colors[unique_ind]
        # normals = normals[unique_ind]

        points["coords"] = coords
        points["colors"] = colors
        # points['normals'] = normals

        print("Pointcloud size:", points["coords"].shape[0])
        return points

    def get_point_cloud(self):
        return self.point_cloud

    def generate_density(self, width=256, height=256):

        ps = self.point_cloud["coords"] * -1
        ps[:, 0] *= -1
        ps[:, 1] *= -1

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(ps)
        pcd.estimate_normals()

        # zs = np.round(ps[:,2] / 100) * 100
        # zs, zs_ind = np.unique(zs, return_index=True, axis=0)
        # ps_ind = ps[:, :2] ==
        # print("Generate density...")

        image_res = np.array((width, height))

        max_coords = np.max(ps, axis=0)
        min_coords = np.min(ps, axis=0)
        max_m_min = max_coords - min_coords

        max_coords = max_coords + 0.1 * max_m_min
        min_coords = min_coords - 0.1 * max_m_min

        normalization_dict = {}
        normalization_dict["min_coords"] = min_coords
        normalization_dict["max_coords"] = max_coords
        normalization_dict["image_res"] = image_res

        # coordinates = np.round(points[:, :2] / max_coordinates[None,:2] * image_res[None])
        coordinates = np.round(
            (ps[:, :2] - min_coords[None, :2]) / (max_coords[None, :2] - min_coords[None, :2]) * image_res[None]
        )
        coordinates = np.minimum(np.maximum(coordinates, np.zeros_like(image_res)), image_res - 1)

        density = np.zeros((height, width), dtype=np.float32)

        unique_coordinates, counts = np.unique(coordinates, return_counts=True, axis=0)
        # print(np.unique(counts))
        # counts = np.minimum(counts, 1e2)

        unique_coordinates = unique_coordinates.astype(np.int32)

        density[unique_coordinates[:, 1], unique_coordinates[:, 0]] = counts
        density = density / np.max(density)
        # print(np.unique(density))

        normals = np.array(pcd.normals)
        normals_map = np.zeros((density.shape[0], density.shape[1], 3))

        import time

        start_time = time.time()
        for i, unique_coord in enumerate(unique_coordinates):
            # print(normals[unique_ind])
            normals_indcs = np.argwhere(np.all(coordinates[::10] == unique_coord, axis=1))[:, 0]
            normals_map[unique_coordinates[i, 1], unique_coordinates[i, 0], :] = np.mean(
                normals[::10][normals_indcs, :], axis=0
            )

        print("Time for normals: ", time.time() - start_time)

        normals_map = (np.clip(normals_map, 0, 1) * 255).astype(np.uint8)

        # plt.figure()
        # plt.imshow(normals_map)
        # plt.show()

        return density, normals_map, normalization_dict

    def visualize(self, export_path=None):
        pcd = o3d.geometry.PointCloud()

        points = self.point_cloud["coords"]

        print(np.max(points, axis=0))
        indices = np.where(points[:, 2] < 2000)

        points = points[indices]
        points[:, 1] *= -1
        points[:, :] /= 1000
        pcd.points = o3d.utility.Vector3dVector(points)

        if self.generate_normal:
            normals = self.point_cloud["normals"]
            normals = normals[indices]
            pcd.normals = o3d.utility.Vector3dVector(normals)
        if self.generate_color:
            colors = self.point_cloud["colors"]
            colors = colors[indices]
            pcd.colors = o3d.utility.Vector3dVector(colors)

        # wireframe_geo_list = visualize_wireframe(annos, vis=False, ret=True)
        # o3d.visualization.draw_geometries([pcd] + wireframe_geo_list)
        # o3d.visualization.draw_geometries([pcd])

        pcd.estimate_normals()

        # radii = 0.01
        # mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(pcd, radii)

        # alpha = 0.1
        # tetra_mesh, pt_map = o3d.geometry.TetraMesh.create_from_point_cloud(pcd)
        # mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_alpha_shape(pcd, alpha, tetra_mesh, pt_map)

        o3d.visualization.draw_geometries([pcd])

        if export_path is not None:
            o3d.io.write_point_cloud(export_path, pcd)

        # o3d.visualization.draw_geometries([pcd])

    def export_ply(self, path):
        """
        ply
        format ascii 1.0
        comment Mars model by Paul Bourke
        element vertex 259200
        property float x
        property float y
        property float z
        property uchar r
        property uchar g
        property uchar b
        property float nx
        property float ny
        property float nz
        end_header
        """
        with open(path, "w") as f:
            f.write("ply\n")
            f.write("format ascii 1.0\n")
            f.write("element vertex %d\n" % self.point_cloud["coords"].shape[0])
            f.write("property float x\n")
            f.write("property float y\n")
            f.write("property float z\n")
            if self.generate_color:
                f.write("property uchar red\n")
                f.write("property uchar green\n")
                f.write("property uchar blue\n")
            if self.generate_normal:
                f.write("property float nx\n")
                f.write("property float ny\n")
                f.write("property float nz\n")
            f.write("end_header\n")
            for i in range(self.point_cloud["coords"].shape[0]):
                normal = []
                color = []
                coord = self.point_cloud["coords"][i].tolist()
                if self.generate_color:
                    color = list(map(int, (self.point_cloud["colors"][i] * 255).tolist()))
                if self.generate_normal:
                    normal = self.point_cloud["normals"][i].tolist()
                data = coord + color + normal
                f.write(" ".join(list(map(str, data))) + "\n")