citynav/rasterize.py at main · water-cookie/citynav · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import os
from pathlib import Path

import numpy as np
import pdal
import rasterio
from PIL import Image


Z_NODATA = -9999


def generate_pipeline(ply_path: Path):
    return f"""
    [
        "{ply_path.as_posix()}",
        {{
            "type": "writers.gdal",
            "filename": "z.tif",
            "binmode": true,
            "dimension": "Z",
            "data_type": "double",
            "nodata": {Z_NODATA},
            "output_type": "max",
            "resolution": 0.1
        }},
        {{
            "type": "writers.gdal",
            "filename": "r.tif",
            "binmode": true,
            "dimension": "Red",
            "data_type": "uint8",
            "nodata": 0,
            "output_type": "mean",
            "resolution": 0.1
        }},
        {{
            "type": "writers.gdal",
            "filename": "g.tif",
            "binmode": true,
            "dimension": "Green",
            "data_type": "uint8",
            "nodata": 0,
            "output_type": "mean",
            "resolution": 0.1
        }},
        {{
            "type": "writers.gdal",
            "filename": "b.tif",
            "binmode": true,
            "dimension": "Blue",
            "data_type": "uint8",
            "nodata": 0,
            "output_type": "mean",
            "resolution": 0.1
        }}
    ]
    """


def rasterize_ply(ply_path, out_dir='rgbd'):

    ply_path = Path(ply_path)
    block_name = ply_path.stem
    out_dir = Path(out_dir)
    out_dir.mkdir(exist_ok=True, parents=True)

    # rasterize the point cloud
    pipeline_json = generate_pipeline(ply_path)
    pipeline = pdal.Pipeline(pipeline_json)
    pipeline.execute()

    # set z nodata to the minimum z value minus one
    with rasterio.open('z.tif') as z_dataset:
        z = z_dataset.read(1)
        min_z = (z != Z_NODATA).min()
        z[z == Z_NODATA] = min_z - 1
        with rasterio.open(out_dir/f'{block_name}.tif', 'w', **z_dataset.profile) as dst:
            dst.write(z, 1)

    # write rgb image
    rgb = np.stack([rasterio.open(f"{color}.tif").read(1) for color in ["r", "g", "b"]], axis=-1)
    Image.fromarray(rgb).save(out_dir/f"{block_name}.png")

    # clean up
    for temp_file in ["z.tif", "r.tif", "g.tif", "b.tif"]:
        os.remove(temp_file)


if __name__ == "__main__":

    import sys

    ply_dir = sys.argv[1]
    out_dir = sys.argv[2] if len(sys.argv) > 2 else 'data/rgbd'
    for ply_path in Path(ply_dir).rglob("*.ply"):
        if ply_path.stem not in ["cambridge_block_0", "cambridge_block_1", "cambridge_block_34"]:
            rasterize_ply(ply_path, out_dir)