render_custom_point_cloud.py

"""
Render cropped las file with same camera paramerters as real one.
"""
import cv2
import torch
import pytorch3d

import laspy as lp
import numpy as np
import pandas as pd

from pathlib import Path
from fire import Fire
from datetime import datetime

from pytorch3d.structures import Pointclouds
from pytorch3d.renderer import (look_at_view_transform,
                                FoVPerspectiveCameras,
                                PerspectiveCameras,
                                PointsRasterizationSettings,
                                PointsRenderer,
                                PointsRasterizer,
                                AlphaCompositor)


def main(
    device_type=None
):
    assert device_type in [None, "cpu", "cuda"]
    if device_type is not None:
        device = torch.device(device_type)
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    las_fn = Path("croppped-from=16L-at=2487972.539-6959599.716.las")
    las = lp.read(las_fn)

    points = torch.tensor([las.x, las.y, las.z], dtype=torch.float32).T
    rgb = np.stack([las.red, las.green, las.blue]).T

    rgb = rgb / 256.0
    rgb = torch.tensor(rgb, dtype=torch.float32)
    pc = Pointclouds(points=[points], features=[rgb])
    pc = pc.to(device)

    frame_x = 2487972.539
    frame_y = 6959599.716
    frame_z = 468.7058266

    R = torch.tensor([[
        [-1.0,  0.0,  0.0],
        [ 0.0,  0.0,  1.0],
        [ 0.0,  1.0,  0.0],
    ]], dtype=torch.float32)

    T = torch.tensor([[frame_x, -frame_z, -frame_y]], dtype=torch.float32)


    print(f"r matrix: {R}")
    print(f"f matrix dtype: {R.dtype}")
    print(f"r matrix shape: {R.shape}")
    print(f"t matrix: {T}")
    print(f"t matrix dtype: {T.dtype}")
    print(f"t matrix shape: {T.shape}")

    image_size = (512, 512)

    focal_length = torch.tensor((1.0 + 0.0131234, 1.0 + 0.0131234)).unsqueeze(0)
    cameras = PerspectiveCameras(device=device,
                                 focal_length=focal_length,
                                 image_size=[image_size],
                                 R=R, T=T)

    raster_settings = PointsRasterizationSettings(
        image_size=image_size,
        radius=0.005,
        points_per_pixel=1,
    )

    rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)

    renderer = PointsRenderer(
        rasterizer=rasterizer,
        compositor=AlphaCompositor()
    )
    images = renderer(pc)

    img = images[0, ..., :3].cpu().numpy()

    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    run_name = "render-custom-pc"
    time_str = datetime.now().strftime("%d-%m-%H-%M-%S")
    save_name = f"name={run_name}-device={device}-time={time_str}.png"
    cv2.imwrite(save_name, img)


if __name__ == "__main__":
    Fire(main)