GeoTiff ファイルでディープラーニングするためのライブラリ Raster Vision

AI リサーチャーの松野です。

当社では衛星画像、航空画像、ドローンオルソ画像といった GeoTiff ファイルを対象としてディープラーニングを行うことがあります。 こういったファイルは、学習・推論の際に以下のような共通の前処理・後処理を適用することが多いです。

• 画像サイズが数万ピクセルx数万ピクセルと大きいため、モデルに入力する前に小さな画像ファイル（タイル）に分割する
• タイルごとに実施した推論結果を集約し、地理情報を付加する

当社ではこのような処理を社内で開発していますが、既存のライブラリで利用できるものがあれば便利なことがあるかもしれません。 そこで GeoTiff ファイルでディープラーニングするためのライブラリについて調べてみました。

結果としていくつかのライブラリを見つけることができました。本記事では、その中で最近でもアクティブに更新されていて、比較的使いやすそうな Raster Vision を紹介します。

github.com

docs.rastervision.io

Raster Vision は、GeoTiff ファイルや GeoJSON ファイルを用いてディープラーニングを実行することができる PyTorch ベースのオープンソースライブラリです。 当初は設定ファイルを記述して CLI でコマンドを実行することで学習・推論をおこなうローコードフレームワークだったのですが、2022年12月に Python ライブラリとしての機能がリリースされました。

以下でいくつかの機能を、セマンティックセグメンテーションの場合を題材として紹介します。詳細は公式のチュートリアルや API リファレンスをご覧ください。

GeoTiff ファイルや GeoJSON ファイルの利用

巨大な画像データから、自前でタイルに分割せずに Dataset を生成することができます。 label_raster_uri のかわりに label_vector_uri を指定することで GeoJSON ファイルを教師ラベルとして利用することもできます。

from rastervision.core.data import ClassConfig
from rastervision.pytorch_learner.dataset.semantic_segmentation_dataset import make_ss_geodataset, SemanticSegmentationSlidingWindowGeoDataset

config = ClassConfig(names=['0_background', ...],
                    null_class='0_background')
list_train_images_str = ['path/to/tif_1', 'path/to/tif_2', ...]
list_train_annotations_str = ['path/to/label_1', 'path/to_label_2', ...]
train_dataset = make_ss_geodataset(SemanticSegmentationSlidingWindowGeoDataset,
                                   image_uri=list_train_images_str,
                                   label_raster_uri=list_train_annotations_str,
                                   class_config=config,
                                   size=256,
                                   stride=256)

学習

データやソルバに関する設定を行ったうえで学習を実行します。

from rastervision.pytorch_learner import SemanticSegmentationGeoDataConfig, SolverConfig, SemanticSegmentationLearnerConfig, SemanticSegmentationLearner

data_cfg = SemanticSegmentationGeoDataConfig(class_names=config.names,
                                             class_colors=config.colors,
                                             num_workers=0)
solver_cfg = SolverConfig(batch_sz=32)
learner_cfg = SemanticSegmentationLearnerConfig(data=data_cfg, solver=solver_cfg)

learner = SemanticSegmentationLearner(cfg=learner_cfg,
                                      output_dir='/path/to/output/dir/',
                                      model=model,
                                      train_ds=train_dataset)
learner.train(epochs=10)

ただし Raster Vision 0.20.2 では、train_dataset から得られる教師ラベルの型は元データの型に依存するようで、PyTorch の CrossEntropyLoss() が受け取るべき long とは異なる場合があります。その場合、損失関数を上書きするなどして修正する必要があります。

推論

学習時と同様に make_ss_geodataset() により推論対象となるデータを生成した後、以下のように推論することで、セマンティックセグメンテーションの結果を GeoTiff 形式で保存することができます。

from rastervision.core.data import SemanticSegmentationLabels

dataset = make_ss_geodataset(SemanticSegmentationSlidingWindowGeoDataset,
                                 image_uri=str(test_image_path),
                                 label_raster_uri=str(test_annotation_path),
                                 class_config=config,
                                 pad_direction='both',
                                 size=512,
                                 stride=128)

predictions = learner.predict_dataset(dataset=dataset, numpy_out=True)
pred_labels = SemanticSegmentationLabels.from_predictions(
    dataset.windows,
    predictions,
    extent=dataset.scene.extent,
    num_classes=len(config),
    smooth=True,
    crop_sz=128
)

pred_labels.save(uri=str(test_result_dir),
                 crs_transformer=dataset.scene.raster_source.crs_transformer,
                 class_config=config,
                 profile_overrides={'compress': 'DEFLATE'})

推論時には stride や crop_sz を設定することで、タイルごとの推論結果に対して、それぞれの外縁部分を除外したり、残りの重複部分に関して平均をとったりした上で GeoTiff 形式のセグメンテーション結果へ集約することが可能となっています。

（画像は Medium 記事より引用）

medium.com

おわりに

今回紹介した以外にも、Raster Vision には以下のような機能もあります。

• タイル単位での画像分類
• 物体検出
• データセット化する領域 (Area of Interest) の指定

また、データセット生成や推論結果の GeoTiff ファイルへの集約処理部分だけを Raster Vision に処理してもらって、学習や推論にかかわる処理を PyTorch Lightning を用いておこなうことも可能です。

興味がある方は試してみてください。

202304_skymatix_engineer from skymatixHR

www.slideshare.net