Create new raster (.tif) from standard deviation stretched bands, works with dstack but not to write a new file, Python - python-3.x

Question

I am sorry if the title is unclear, I am new to python and my vocabulary is limited.
What I am trying to do is apply a standard deviation stretch to each band in a .tif raster and then create a new raster (.tif) by stacking those bands using GDAL (Python).
I able to create new false color rasters with differing band combinations and save them, and I am able to create my desired IMAGE in python using dstack (first block of code), but I am unable to save that image as a georectified .tif file.
So to create the stretched image using dstack my code looks like:
import os
import numpy as np
import matplotlib.pyplot as plt
import math
from osgeo import gdal
# code from my prof
def std_stretch_data(data, n=2):
"""Applies an n-standard deviation stretch to data."""
# Get the mean and n standard deviations.
mean, d = data.mean(), data.std() * n
# Calculate new min and max as integers. Make sure the min isn't
# smaller than the real min value, and the max isn't larger than
# the real max value.
new_min = math.floor(max(mean - d, data.min()))
new_max = math.ceil(min(mean + d, data.max()))
# Convert any values smaller than new_min to new_min, and any
# values larger than new_max to new_max.
data = np.clip(data, new_min, new_max)
# Scale the data.
data = (data - data.min()) / (new_max - new_min)
return data
# open the raster
img = gdal.Open(r'/Users/Rebekah/ThesisData/TestImages/OG/OG_1234.tif')
#open the bands
red = img.GetRasterBand(1).ReadAsArray()
green = img.GetRasterBand(2).ReadAsArray()
blue = img.GetRasterBand(3).ReadAsArray()
# create alpha band where a 0 indicates a transparent pixel and 1 is a opaque pixel
# (this is from class and i dont FULLY understand it)
alpha = np.where(red + green + blue ==0, 0, 1).astype(np.byte)
red_stretched = std_stretch_data(red, 1)
green_stretched = std_stretch_data(green, 1)
blue_stretched = std_stretch_data(blue, 1)
data_stretched = np.dstack((red_stretched, green_stretched, blue_stretched, alpha))
plt.imshow(data_stretched)
plt.show()
And that gives me a beautiful image of exactly what I want in a separate window. But no where in that code is an option to assign projections, or save it as a multiband tif.
So I took that and applied it the best I could to the code I use to create false color images and it fails (code below). If I create a 4 band tif with the alpha band the output is an empty tif, and if I create a 3 band tif and omit the alpha band the output is an entirely black tif.
import os
import numpy as np
import matplotlib.pyplot as plt
import math
from osgeo import gdal
#code from my professor
def std_stretch_data(data, n=2):
"""Applies an n-standard deviation stretch to data."""
# Get the mean and n standard deviations.
mean, d = data.mean(), data.std() * n
# Calculate new min and max as integers. Make sure the min isn't
# smaller than the real min value, and the max isn't larger than
# the real max value.
new_min = math.floor(max(mean - d, data.min()))
new_max = math.ceil(min(mean + d, data.max()))
# Convert any values smaller than new_min to new_min, and any
# values larger than new_max to new_max.
data = np.clip(data, new_min, new_max)
# Scale the data.
data = (data - data.min()) / (new_max - new_min)
return data
#open image
img = gdal.Open(r'/Users/Rebekah/ThesisData/TestImages/OG/OG_1234.tif')
# get geotill driver
gtiff_driver = gdal.GetDriverByName('GTiff')
# read in bands
red = img.GetRasterBand(1).ReadAsArray()
green = img.GetRasterBand(2).ReadAsArray()
blue = img.GetRasterBand(3).ReadAsArray()
# create alpha band where a 0 indicates a transparent pixel and 1 is a opaque pixel
# (this is from class and i dont FULLY understand it)
alpha = np.where(red + green + blue ==0, 0, 1).astype(np.byte)
# apply the 1 standard deviation stretch
red_stretched = std_stretch_data(red, 1)
green_stretched = std_stretch_data(green, 1)
blue_stretched = std_stretch_data(blue, 1)
# create empty tif file
NewImg = gtiff_driver.Create('/Users/riemann/ThesisData/TestImages/FCI_tests/1234_devst1.tif', img.RasterXSize, img.RasterYSize, 4, gdal.GDT_Byte)
if NewImg is None:
raise IOerror('could not create new raster')
# set the projection and geo transform of the new raster to be the same as the original
NewImg.SetProjection(img.GetProjection())
NewImg.SetGeoTransform(img.GetGeoTransform())
# write new bands to the new raster
band1 = NewImg.GetRasterBand(1)
band1.WriteArray(red_stretched)
band2 = NewImg.GetRasterBand(2)
band2.WriteArray(green_stretched)
band3= NewImg.GetRasterBand(3)
band3.WriteArray(blue_stretched)
alpha_band = NewImg.GetRasterBand(4)
alpha_band.WriteArray(alpha)
del band1, band2, band3, img, alpha_band
I am not entirely sure how to go from here and create a new file displaying the stretch on the different bands.
The image is just a 4 band raster (NAIP) downloaded from earthexplorer, I can upload the specific image I am using for my test if needed but there is nothing inherently special about this file compared to other NAIP images.

Answer 1

You should close the new Dataset (NewImg) as well by either adding it to the del list you already have, or setting it to None.
That properly closes the file and makes sure all data is written to disk.
There is however another issue, you are scaling your data between 0 and 1, but storing it as a Byte. So either change the output datatype from gdal.GDT_Byte to something like gdal.GDT_Float32. Or multiply your scaled data to fit the output datatype, in the case of Byte multiple with 255 (don't forget the alpha), you should properly round it for accuracy, GDAL will otherwise truncate to the nearest integer.
You can use np.iinfo() to check what the range of a datatype is, in case you are unsure what multiplication to use for other datatypes.
Depending on your use case, it might be easiest to use gdal.Translate for the scaling. If you would modify your scaling function a little to return the scaling parameteters instead of the data, you could use something like:
ds = gdal.Translate(output_file, input_file, outputType=gdal.GDT_Byte, scaleParams=[
[old_min_r, old_max_r, new_min_r, new_max_r], # red
[old_min_g, old_max_g, new_min_g, new_max_g], # green
[old_min_b, old_max_b, new_min_b, new_max_b], # blue
[old_min_a, old_max_a, new_min_a, new_max_a], # alpha
])
ds = None
You could also add the exponents keyword for non-linear stretching.
Using gdal.Translate would save you from all the standard file creation boilerplate, you still would need to think about the datatype, since that might change compared to the input file.