Efficient way to multiply specific raster with specific variable python

Question

I have 4 separate raster bands of the same image and 4 different coefficients (one for each band) stored as variables. I would like to iterate over the raster folder, and multiply each band with its corresponded coefficient. I can recognize each band by its file name. Let's say I have in folder c:\bands these bands:

sn2_b1.tif
sn2_b2.tif
sn2_b3.tif
sn2_b3.tif

And let's say these are my coefficients:

coeff_b1 = 10
coeff_b2 = 20
Coeff_b3 = 30
Coeff_b4 = 40

My goal is iterate over the bands, multiply each one with its corresponded coefficient and save the result as a new raster. My solution so far is quite cumbersome:

import gdal
import os
src = r"c:\bands"
dst = r"c:\cal_bands"
for i in os.listdir(src):
 if i.endswith(".tif"):
 orig_tif = os.path.join(src,i)
 band_ds = gdal.Open(os.path.join(src,i))
 band_DN = band_ds.GetRasterBand(1).ReadAsArray() 
 geotransform = band_ds.GetGeoTransform()
 geoproj = band_ds.GetProjection()
 xsize = band_ds.RasterXSize
 ysize = band_ds.RasterYSize
 band_DN = band_DN.astype('f')
 format1 = "GTiff"
 driver = gdal.GetDriverByName(format1)
 dst_datatype = gdal.GDT_Float32
 if i.split("_")[1] == 'b1.tif':
 cal_band = band_DN * coeff_b1
 elif i.split("_")[1] == 'b2.tif':
 cal_band = band_DN * coeff_b2
 elif i.split("_")[1] == 'b3.tif':
 cal_band = band_DN * coeff_b3 
 elif i.split("_")[1] == 'b4.tif':
 cal_band = band_DN * coeff_b4
 new_name = "cal_" + i
 output_path = os.path.join(dst,new_name)
 dst_ds = driver.Create(output_path ,xsize,ysize,1,dst_datatype)
 dst_ds.GetRasterBand(1).WriteArray(cal_band)
 dst_ds.SetGeoTransform(geotransform)
 dst_ds.SetProjection(geoproj)
 dst_ds = None

This script works great but is there a way to make it less cumbersome? In this example I have only 4 bands, which is easy to write all the ifs, but I might have 20 bands or more.

Answer 1

It looks as though you are doing different operations inside a single loop, which you can break off into functions.

If I understand correctly, you:

1. Have a band-coefficient lookup table (LUT)
2. Unknown number of tifs in a source folder
3. Image band number is determined by the characters after the final underscore in the file name.

And you can match them together by extracting the band name from the string. The natural way to write LUTs in Python is with built-in dictionaries.

Code

import gdal
import os
src = r'c:\bands'
dst = r'c:\cal_bands'
COEFF_LUT = {'b1': 10, 'b2': 20, 'b3': 30, 'b4': 40}
FORMAT = 'GTiff'
def band_num(f):
 """Given raster filename with band number info, returns band number"""
 filename = f.split('.')[0]
 band = filename.split('_')[-1]
 return band
def multiply_raster_by_coeff(file, band):
 """Given filename and band number of image, multiplies it by saves a copy
 in specified destination folder"""
 coeff = COEFF_LUT[band]
 dst_datatype = gdal.GDT_Float32
 driver = gdal.GetDriverByName(FORMAT)
 # Read file, get metadata
 band_ds = gdal.Open(os.path.join(src, file))
 band_DN = band_ds.GetRasterBand(1).ReadAsArray()
 geotransform = band_ds.GetGeoTransform()
 geoproj = band_ds.GetProjection()
 xsize = band_ds.RasterXSize
 ysize = band_ds.RasterYSize
 band_DN = band_DN.astype('f')
 # Multiply by band coefficient
 cal_band = band_DN * coeff
 # Write to disk
 new_name = 'cal_' + file
 output_path = os.path.join(dst, new_name)
 dst_ds = driver.Create(output_path, xsize, ysize, 1, dst_datatype)
 dst_ds.GetRasterBand(1).WriteArray(cal_band)
 dst_ds.SetGeoTransform(geotransform)
 dst_ds.SetProjection(geoproj)
if __name__ == '__main__':
 files = [f for f in os.listdir(path) if f.endswith('.tif')]
 image_bands = {f: band_num(f) for f in files}
 for file, band in image_bands:
 multiply_raster_by_coeff(file, band)

Conclusion

You can now handle an arbitrary number of bands just by updating the COEFF_LUT dictionary.

You could further break the multiply_raster_by_coeff function. Of course, performance-wise the two scripts are equivalent, but future you (and collaborators) will appreciate this if it needs changes.

Answer 2

I think what you are looking for is How to get a variable name as a string in Python?. Using your script, you should be able to do it like this:

import gdal
import os
# define your variables
coeff_b1 = 10
coeff_b2 = 20
coeff_b3 = 30
coeff_b4 = 40
# create a dictionary with their names and values
coeff_dict = dict((name, eval(name)) for name in ["coeff_b1", "coeff_b2", "coeff_b3", "coeff_b4"])
src = r"c:\bands"
dst = r"c:\cal_bands"
for i in os.listdir(src):
 if i.endswith(".tif"):
 band_ds = gdal.Open(os.path.join(src,i))
 band_DN = band_ds.GetRasterBand(1).ReadAsArray() 
 geotransform = band_ds.GetGeoTransform()
 geoproj = band_ds.GetProjection()
 xsize = band_ds.RasterXSize
 ysize = band_ds.RasterYSize
 band_DN = band_DN.astype("f")
 format1 = "GTiff"
 driver = gdal.GetDriverByName(format1)
 dst_datatype = gdal.GDT_Float32
 # look for your matching variable
 band_num = i.split("_")[1].split(".tif")[0]
 coeff_name = "coeff_" + band_num
 coeff = coeff_dict[coeff_name]
 # do the math
 cal_band = band_DN * coeff
 # do the rest
 new_name = "cal_" + i
 output_path = os.path.join(dst,new_name)
 dst_ds = driver.Create(output_path ,xsize,ysize,1,dst_datatype)
 dst_ds.GetRasterBand(1).WriteArray(cal_band)
 dst_ds.SetGeoTransform(geotransform)
 dst_ds.SetProjection(geoproj)
 dst_ds = None

• user88484: Thanks, my actual script is a bit more complex but I used this dictionary approach you suggested and it made my script much shorter and elegant