clip_raster_from_shp.py

import os

from osgeo import gdal


def rasterize(input_shp, out_raster):
    # run gdal_rasterizer
    layer_name = os.path.basename(input_shp).split('.')[0]
    gdal_rasterize_cmd = f'gdal_rasterize -burn 1 -ts 1000 1000 -l {layer_name} {input_shp} {out_raster}'
    print(gdal_rasterize_cmd)
    os.system(gdal_rasterize_cmd)


def calculate(pop_raster, rasterized_shp, out_clipped_raster):
    # run gdal_calc.py
    gdal_calc_cmd = f'gdal_calc.py -A {pop_raster} -B {rasterized_shp} --outfile={out_clipped_raster} --calc="A*B"'
    print(gdal_calc_cmd)
    os.system(gdal_calc_cmd)


def crop_raster(dest, src, out_path):
    """
    This will crop a src raster to a dest raster using same resolution.
    """
    src_ds = gdal.Open(src)
    src_proj = src_ds.GetProjection()
    data_type = src_ds.GetRasterBand(1).DataType
    n_bands = src_ds.RasterCount

    dest_ds = gdal.Open(dest)
    dest_proj = dest_ds.GetProjection()
    dest_geotrans = dest_ds.GetGeoTransform()
    w = dest_ds.RasterXSize
    h = dest_ds.RasterYSize

    out_filename = src.replace(".tif", "_crop.tiff")
    out_filename = out_path + out_filename.split('/')[-1]
    out_ds = gdal.GetDriverByName('gtiff').Create(out_filename,
                                                  w, h, n_bands, data_type)
    out_ds.SetGeoTransform(dest_geotrans)
    out_ds.SetProjection(dest_proj)

    gdal.ReprojectImage(src_ds, out_ds, src_proj,
                        dest_proj, gdal.GRA_NearestNeighbour)

    return out_filename


# out_dir: path where output is written. Must be cleared at every run
out_dir = '/Users/pcorti/temp/out/'
# input shapefile
input_shp = '/Users/pcorti/data/shapefile/Ultra_Rural_V8/Ultra_Rural_V8.shp'
# output path for the rasterized shapefile
rasterized = f'{out_dir}rasterized.tif'
# input raster to use
input_raster = '/Users/pcorti/data/raster/zaf_ppp_2020.tif'

rasterize(input_shp, rasterized)
cropped_rasterized = crop_raster(rasterized, input_raster, out_dir)
calculate(cropped_rasterized, rasterized, f'{out_dir}result.tiff')