How to mask bands of an image based on values of one of the bands using numpy and gdal

Question

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

Answer 1

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

Answer 2

I managed to solve the problem with the following code:

import numpy as np
from osgeo import gdal

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

# Read the raster bands into numpy arrays
yod = dataset.GetRasterBand(1).ReadAsArray().astype(np.int16)
mag = dataset.GetRasterBand(2).ReadAsArray().astype(np.int16)

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

# Create a mask based on the specified band
mask = np.where(yod != threshold, True, False)

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

# Create a new masked array with both bands
masked_array = np.ma.array([yod_masked, mag_masked])

# Convert the masked array to a regular numpy array
masked_array_data = masked_array.filled(fill_value=0)

# Create a new raster file for the 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_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(masked_array_data):
    output_band = output_dataset.GetRasterBand(i + 1)
    output_band.WriteArray(masked_band)
    output_band.SetNoDataValue(0)

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