Python GDAL – How to Mask Image Bands with Numpy and GDAL

gdalimagemaskingnumpypython

I have a raster with two bands, and I want to mask both bands to get data where band "yod" has a value of 2011 and export a new tif file with the same two bands masked.
The following script does not work correctly since the mask is not applied to the output raster.

import numpy as np
from osgeo import gdal, ogr

# Read a file
dataset = gdal.Open('test.tif')

# read raster bands into numpy array
yod = dataset.GetRasterBand(1).ReadAsArray()
mag = dataset.GetRasterBand(2).ReadAsArray()

# create a mask based on a specific band
mask = np.where(yod == 2011, True, False)

# Apply masks for all bands
yod_masked = np.ma.masked_array(yod, mask=mask)
mag_masked = np.ma.masked_array(mag, mask=mask)

# Create a new raster for masked bands
driver = gdal.GetDriverByName('GTiff')
output_path = 'output_masked_file.tif'
output_dataset = driver.Create(output_path, dataset.RasterXSize, dataset.RasterYSize, 2, gdal.GDT_Float32)
output_dataset.SetProjection(dataset.GetProjection())
output_dataset.SetGeoTransform(dataset.GetGeoTransform())

# Write the masked bands to the output raster file
output_dataset.GetRasterBand(1).WriteArray(yod_masked)
output_dataset.GetRasterBand(2).WriteArray(mag_masked)

# Save and close the output dataset
output_dataset.FlushCache()
output_dataset = None

Best Answer

A slightly simplified version of your solution by masking in fewer steps (assuming the band order is always the same):

import numpy as np
from osgeo import gdal

# Open the raster file
dataset = gdal.Open('test.tif')

# Read the raster bands into a numpy array
array_to_mask = dataset.ReadAsArray().astype(np.int16)

# Define the value threshold for masking
threshold = 2011  # Change this value to your desired threshold

# Do the masking in two steps
array_to_mask[0][array_to_mask[0] != threshold] = 0
array_to_mask[1][array_to_mask[0] != threshold] = 0

# Create a new raster file for the masked bands
driver = gdal.GetDriverByName('GTiff')
output_path = 'array_to_mask.tif'
output_dataset = driver.Create(output_path, dataset.RasterXSize, dataset.RasterYSize, 2, gdal.GDT_Int16)

# Set the projection and geotransform
output_dataset.SetProjection(dataset.GetProjection())
output_dataset.SetGeoTransform(dataset.GetGeoTransform())

# Write the masked bands to the output raster file
for i, masked_band in enumerate(array_to_mask):
    output_band = output_dataset.GetRasterBand(i + 1)
    output_band.WriteArray(masked_band)
    output_band.SetNoDataValue(0)

# Save and close the output dataset
output_dataset = None
Related Question