Rustで点群データからGeoTiffを作成する

これは [MIERUNE AdventCalendar 2023] 22日目の記事です。
昨日は @bordoray さんによるQField cloudを使ってみたでした。

MIERUNEでは毎週火曜日から木曜日まで朝会のあと、ちょっとしたLTを持ち回りで行うが、メンバーの一人がPythonで点群からGeoTiffを作成するコードを作成し、その実行速度を紹介していた。

コードは以下のようなもので、ちょうどRustを勉強しているところだし、試しにRust版だとどうなるのか興味本位で作成してみたわけである。

点群データは静岡の点群データを25ファイル用意した。
(08nd8030.txt 〜 08nd8074.txt)

各CSVファイルのイメージ

1,-39999.75,-114900.25,635.30,0
2,-39999.25,-114900.25,634.90,1
3,-39998.75,-114900.25,634.60,1
4,-39998.25,-114900.25,634.20,1
5,-39997.75,-114900.25,633.70,0
.
.
.

Pythonのコード

from pathlib import Path
from shutil import rmtree
from time import time

import pandas as pd
import numpy as np
import rasterio
from rasterio.transform import from_origin

INPUT_DIR = Path("./data/")
assert INPUT_DIR.exists()

input_paths = sorted(INPUT_DIR.glob("*txt"))
len(input_paths)


OUTPUT_DIR_BASE = Path("./output/")
OUTPUT_DIR_BASE.mkdir(exist_ok=True)

OUTPUT_DIR_LOOP = OUTPUT_DIR_BASE / "loop"
OUTPUT_DIR_VECTORIZED = OUTPUT_DIR_BASE / "vectorized"

RESOLUTION = 0.5

for p in [OUTPUT_DIR_LOOP, OUTPUT_DIR_VECTORIZED]:
    if p.exists():
        rmtree(p)
    p.mkdir(exist_ok=True)


def xyz2raster_loop(df_xyz):
    """ループ版"""

    xmin, xmax = df_xyz["x"].min(), df_xyz["x"].max()
    ymin, ymax = df_xyz["y"].min(), df_xyz["y"].max()

    n_rows = int((ymax - ymin) / RESOLUTION) + 1
    n_cols = int((xmax - xmin) / RESOLUTION) + 1
    raster_array = np.zeros((n_rows, n_cols))

    # ↓ これが遅い
    for _, row in df_xyz.iterrows():
        x, y, z = row["x"], row["y"], row["z"]
        row_idx = int((ymax - y) / RESOLUTION)
        col_idx = int((x - xmin) / RESOLUTION)
        raster_array[row_idx, col_idx] = z

    return raster_array, xmin, ymax


def xyz2raster_vectorized(df_xyz):
    """ベクタライズ版"""

    xmin, xmax = df_xyz["x"].min(), df_xyz["x"].max()
    ymin, ymax = df_xyz["y"].min(), df_xyz["y"].max()

    n_rows = int((ymax - ymin) / RESOLUTION) + 1
    n_cols = int((xmax - xmin) / RESOLUTION) + 1
    raster_array = np.zeros((n_rows, n_cols))

    # 1行ごとのループではなく、ベクトル化（vecrotization）により配列をまとめて処理
    x_arr = df_xyz["x"].values
    y_arr = df_xyz["y"].values
    z_arr = df_xyz["z"].values

    row_indices = ((ymax - y_arr) / RESOLUTION).astype(int)
    col_indices = ((x_arr - xmin) / RESOLUTION).astype(int)

    raster_array[row_indices, col_indices] = z_arr

    return raster_array, xmin, ymax


start_time = time()

for in_path in input_paths:
    df_in = pd.read_csv(in_path, header=None, usecols=[
                        1, 2, 3], names=["x", "y", "z"])

    raster_array, xmin, ymax = xyz2raster_vectorized(df_in)
    #raster_array, xmin, ymax = xyz2raster_loop(df_in)

    out_path = OUTPUT_DIR_VECTORIZED / f"{in_path.stem}.tif"
    with rasterio.open(
        out_path,
        mode="w",
        driver="GTiff",
        height=raster_array.shape[0],
        width=raster_array.shape[1],
        count=1,
        dtype=str(raster_array.dtype),
        crs="+init=EPSG:6676",
        transform=from_origin(xmin, ymax, RESOLUTION, RESOLUTION),
    ) as dst:
        dst.write(raster_array, 1)

end_time = time()
elapsed_time = end_time - start_time
print(f"ベクタライズ版: {elapsed_time:.2f} 秒")

最も時間がかかる、いわいるボトルネック部分は点群データのｘ,ｙ,ｚ値を計算しながら配列に落とし込む部分だったことが判明。

• Python単純Forループ

  計測速度 183.88 秒

• PythonのForループは遅いのでPandasのベクトル化（vecrotization）により配列をまとめて処理することにより高速化

  計測速度 2.25 秒

単純ForループではなくPandasの機能を使えばかなりの高速化が実現できた。

Rustではどのように実装するのだろう？

• まず最初にGeoTiffを生成する方法を見つけなければならない。
  • これはgdalをバインディングしたクレートを利用した
    https://crates.io/crates/gdal
• 次にCSVファイルを読み込むためのクレートを用意
  • https://crates.io/crates/csv
• さらにCSVファイルをデシリアライズするため、定番とされるserdeクレートを利用した
  • https://crates.io/crates/serde

  実際に利用してみると、使い方が適切でないのか、自前でデシリアライズしたほうが早かったので結局自前実装に変更した

// serdeクレートを利用するコード(一部抜粋)
#[derive(Debug)]
struct Demdata {
    xmin: f64,
    xmax: f64,
    ymin: f64,
    ymax: f64,
    data: Vec<Record>,
}

impl Default for Demdata {
    fn default() -> Demdata {
        Demdata {
            xmin: f64::MAX,
            xmax: f64::MIN,
            ymin: f64::MAX,
            ymax: f64::MIN,
            data: Vec::new(),
        }
    }
}

fn read_csv(file_path: String) -> Result<Demdata, Box<dyn Error>> {

    let mut demdata = Demdata::default();
    let csv_text = fs::read_to_string(file_path)?;
    let mut rdr = csv::Reader::from_reader(csv_text.as_bytes());
    for result in rdr.records() {
        let record = result?.deserialize(None)?;
        demdata.add(record);
    }
    Ok(demdata)
}

• もっとも効果があったのはイテレーションを自動でマルチスレッド化するrayon (マルチコアCPU最高！！)
  • https://crates.io/crates/rayon

Rustでプログラムを組んでみて実行してみた

• Python単純Forループ

  計測速度 183.88 秒

• PythonのForループは遅いのでPandasのベクトル化（vecrotization）により配列をまとめて処理することにより高速化

  計測速度 2.25 秒

• Rust版

  計測速度 0.876 秒

やはりネイティブコードになるとかなりの高速化が期待できた。

ある意味当たり前の結果ではある

作成したRustコードを掲載

• まだまだ甘そうだが、さらなる高速化のためGPUパワー(Cuda)を借りて試してみるのも良いかもしれない。(実現できたら記事にしたい)

use gdal::raster::Buffer;
use gdal::spatial_ref::SpatialRef;
use gdal::{Driver, DriverManager};
use rayon::prelude::*;
use std::fs::DirEntry;
use std::time::Instant;
use std::{error::Error, fs};

static RESOLUTION: f64 = 0.5;

#[derive(Debug)]
struct Record {
    l1: u32,
    x: f64,
    y: f64,
    z: f64,
    l5: i32,
}

impl Record {
    fn new(data: &str) -> Self {
        let data_vec: Vec<&str> = data.trim_end().split(',').collect();
        Self {
            l1: 0,
            x: data_vec[1].parse().unwrap(),
            y: data_vec[2].parse().unwrap(),
            z: data_vec[3].parse().unwrap(),
            l5: 0,
        }
    }
}
#[derive(Debug)]
struct Demdata {
    xmin: f64,
    xmax: f64,
    ymin: f64,
    ymax: f64,
    zmin: f64,
    zmax: f64,
    data: Vec<Record>,
}

impl Default for Demdata {
    fn default() -> Demdata {
        Demdata {
            xmin: f64::MAX,
            xmax: f64::MIN,
            ymin: f64::MAX,
            ymax: f64::MIN,
            zmin: f64::MAX,
            zmax: f64::MIN,
            data: Vec::new(),
        }
    }
}

fn main() {

    let start = Instant::now(); // 計測開始

    // GDALのGeoTIFF形式のドライバを取得
    let driver = DriverManager::get_driver_by_name("GTiff").unwrap();
    let folder_name = "./data";
    if let Ok(entries) = fs::read_dir(folder_name) {
        entries.for_each(|entry| {
            if let Ok(entry) = entry {
                exec(&entry, &driver);
            }
        });
    }
    println!("処理時間 end: {:?}", start.elapsed());
}

fn exec(entry: &DirEntry, driver: &Driver) {
    #[cfg(debug_assertions)]
    println!("{:?}", entry.file_name());

    let file_name = entry.file_name().into_string().unwrap();
    let in_file_path = format!(
        "{}/{}",
        entry.path().parent().unwrap().to_str().unwrap(),
        file_name
    );

    let dem_data = read_csv(&in_file_path).unwrap();
    make_tiff(&driver, &file_name, &dem_data).unwrap();
}

fn make_tiff(
    driver: &Driver,
    out_file_name: &String,
    dem_data: &Demdata,
) -> Result<(), Box<dyn Error>> {

    #[cfg(debug_assertions)]
    let start = Instant::now(); // 計測開始

    // サイズを計算
    let x_size = (((dem_data.xmax - dem_data.xmin) / RESOLUTION).ceil() + 1.0) as isize;
    let y_size = (((dem_data.ymax - dem_data.ymin) / RESOLUTION).ceil() + 1.0) as isize;

    // 新しいデータセットを作成
    let mut dataset = driver.create_with_band_type::<f64, _>(
        format!("{}.tiff", out_file_name),
        x_size,
        y_size,
        1,
    )?;

    let sr = SpatialRef::from_epsg(6676)?;
    dataset.set_projection(&sr.to_wkt().unwrap())?;

    // GeoTransformを設定
    // 引数 transformation - 変換の係数：
    // 左上隅ピクセルのx座標 (x-offset)
    // ピクセルの幅 (x-resolution)
    // 行の回転（通常はゼロ）
    // 左上隅ピクセルのy座標
    // 列の回転（通常はゼロ）
    // ピクセルの高さ（y解像度，通常は負）
    dataset.set_geo_transform(&[dem_data.xmin, RESOLUTION, 0.0, dem_data.ymin, 0.0, -0.5])?;

    let buffer: Vec<f64> = dem_data.data.par_iter().map(|x| x.z).collect();
    let band_buffer = Buffer::new((x_size as usize, y_size as usize), buffer);
    dataset
        .rasterband(1)?
        .write((0, 0), (x_size as usize, y_size as usize), &band_buffer)?;

    #[cfg(debug_assertions)]
    println!("make_tiff 処理時間 : {:?}", start.elapsed());

    Ok(())
}

fn read_csv(file_path: &String) -> Result<Demdata, Box<dyn Error>> {
    fn make_demdata_from_record(records: Vec<Record>) -> Demdata {
        let mut demdata = Demdata::default();
        let mut demdata = Demdata { xmin: f64::MAX, .. };
        let xmin = records
            .par_iter()
            .map(|x| x.x)
            .min_by(|x, y| x.partial_cmp(y).unwrap());
        let xmax = records
            .par_iter()
            .map(|x| x.x)
            .max_by(|x, y| x.partial_cmp(y).unwrap());
        let ymin = records
            .par_iter()
            .map(|x| x.y)
            .min_by(|x, y| x.partial_cmp(y).unwrap());
        let ymax = records
            .par_iter()
            .map(|x| x.y)
            .max_by(|x, y| x.partial_cmp(y).unwrap());

        demdata.xmax = xmax.unwrap();
        demdata.xmin = xmin.unwrap();
        demdata.ymax = ymax.unwrap();
        demdata.ymin = ymin.unwrap();
        demdata.data = records;
        demdata
    }

    #[cfg(debug_assertions)]
    let start = Instant::now(); // 計測開始

    let contents = fs::read_to_string(&file_path).unwrap();
    let csv_line: Vec<&str> = contents.trim_end().split('\n').collect();

    #[cfg(debug_assertions)]
    println!("read_csv 処理時間: {:?}", start.elapsed());

    #[cfg(debug_assertions)]
    let start2 = Instant::now(); // 計測開始

    let output: Vec<Record> = csv_line.par_iter().map(|&data| Record::new(data)).collect();

    #[cfg(debug_assertions)]
    println!("make Record 処理時間: {:?}", start2.elapsed());

    #[cfg(debug_assertions)]
    let start2 = Instant::now(); // 計測開始

    let demdata = make_demdata_from_record(output);

    #[cfg(debug_assertions)]
    println!("make_demdata 処理時間: {:?}", start2.elapsed());

    Ok(demdata)
}

[package]
name = "csv2geotiff"
version = "0.1.0"
edition = "2021"

[dependencies]
csv = "1.2.2"
gdal = "0.16.0"
rayon = "1.7.0"
#serde = { version = "1.0", features = ["derive"] }
#serde = "1.0.188"
# serde_json is just for the example, not required in general
#serde_json = "1.0"

明日は@northprintさんによる、SvelteKitでURLクエリパラメーターの操作をするです！お楽しみに!