Implement multiprocessing to write GeoTiff with GDAL on deep learning prediction

user_15629049 0
Viewed 289

I'm running a deep learning segmentation model on drone footage. The drone orthomosaic is cut in 256x256px tiles and I'm getting the prediction for each. For this I'm importing the original RGB geotiff tile, extracting the geo information and create a new geotiff with those same info but with the predictions output instead of the RGB. It is running fine with a for loop. However it takes a very long time. I'd like to use multiprocessing to speep up the process. Any help on how to modify the following code to do that ?

Thank you very much

from osgeo import gdal, osr

# Lets export predicted data
pred_dataset = Dataset_pred(x_pred_dir,)
os.listdir(x_pred_dir)

ids = list(range(0,len(os.listdir(x_pred_dir))))

for i in ids:
    
    image = pred_dataset[i]
    image = np.expand_dims(image, axis=0)
    pr_mask = model.predict(image)
    pr_mask = pr_mask.squeeze()
    image_name = os.listdir(x_pred_dir)[i]

    driver = gdal.GetDriverByName('GTiff')
    rows, cols, no_bands = pr_mask.shape
    x_geotif_dir = os.path.join(DATA_pred_DIR, 'images_raw/',os.listdir(x_pred_dir)[i])
    raster_ds = gdal.Open(x_geotif_dir)
    geo_transform=raster_ds.GetGeoTransform()
    projection=raster_ds.GetProjection()

    image_name = os.listdir(x_pred_dir)[i]
    images_fps = os.path.join(y_pred_dir, image_name)

    DataSet = driver.Create(images_fps, cols, rows, no_bands, gdal.GDT_Byte)
    DataSet.SetGeoTransform(geo_transform)
    DataSet.SetProjection(projection)
    data = np.moveaxis(pr_mask, -1, 0)

    for j, image in enumerate(data, 1):
        DataSet.GetRasterBand(j).WriteArray(image)
    DataSet = None
raster deep-learning gdal multiprocessing python
1 Answers

Something like this:

import os
import multiprocessing as mp
from osgeo import gdal, osr


def process_image(dataset_image, image_filename):
    image = np.expand_dims(dataset_image, axis=0)
    pr_mask = model.predict(image)
    pr_mask = pr_mask.squeeze()

    driver = gdal.GetDriverByName('GTiff')
    rows, cols, no_bands = pr_mask.shape
    x_geotif_dir = os.path.join(DATA_pred_DIR, 'images_raw/', image_filename)
    raster_ds = gdal.Open(x_geotif_dir)
    geo_transform=raster_ds.GetGeoTransform()
    projection=raster_ds.GetProjection()

    images_fps = os.path.join(y_pred_dir, image_filename)

    DataSet = driver.Create(images_fps, cols, rows, no_bands, gdal.GDT_Byte)
    DataSet.SetGeoTransform(geo_transform)
    DataSet.SetProjection(projection)
    data = np.moveaxis(pr_mask, -1, 0)

    for j, image in enumerate(data, 1):
        DataSet.GetRasterBand(j).WriteArray(image)
    DataSet = None


def main():
    pred_dataset = Dataset_pred(x_pred_dir,)
    dir_images = os.listdir(x_pred_dir)
    with mp.pool() as pool:
        pool.starmap(process_image, zip(pred_dataset, dir_images)


if __name__ == '__main__':
    main()
user_19253880
Related