The Spatial Space

from osgeo import gdal
import os
import numpy as np
from scipy.ndimage.filters import median_filter, gaussian_filter
from PIL import Image


class SatelliteImageProcessor:
    def __init__(self, band_folder):
        self.band_folder = band_folder
        self.band_dict = {}

    def load_band_files(self):
        band_files = os.listdir(self.band_folder)
        for i in band_files:
            if i.endswith(".jp2"):
                self.band_dict[f"Band_{i[-7:-4]}"] = os.path.join(self.band_folder, i)

    def calculate_band_combination(self, band_key1, band_key2, band_key3):
        band1 = gdal.Open(self.band_dict[band_key1]).ReadAsArray().astype(np.float32)
        band2 = gdal.Open(self.band_dict[band_key2]).ReadAsArray().astype(np.float32)
        band3 = gdal.Open(self.band_dict[band_key3]).ReadAsArray().astype(np.float32)

        # Normalize the bands to 0-1 range
        band1 = band1 / np.max(band1)
        band2 = band2 / np.max(band2)
        band3 = band3 / np.max(band3)

        # Apply gamma correction for brightness adjustment
        gamma = 1.8
        band1 = np.power(band1, 1 / gamma)
        band2 = np.power(band2, 1 / gamma)
        band3 = np.power(band3, 1 / gamma)

        # Apply contrast stretch to enhance the image
        p2, p98 = np.percentile(band1, (2, 98))
        band1 = np.interp(band1, (p2, p98), (0, 1))
        p2, p98 = np.percentile(band2, (2, 98))
        band2 = np.interp(band2, (p2, p98), (0, 1))
        p2, p98 = np.percentile(band3, (2, 98))
        band3 = np.interp(band3, (p2, p98), (0, 1))

        # Apply median filter to reduce speckles
        band1 = median_filter(band1, size=3)
        band2 = median_filter(band2, size=3)
        band3 = median_filter(band3, size=3)

        # Apply Gaussian filter to further reduce speckles
        sigma = 0.7
        band1 = gaussian_filter(band1, sigma=sigma)
        band2 = gaussian_filter(band2, sigma=sigma)
        band3 = gaussian_filter(band3, sigma=sigma)

        # Scale the bands to 0-255 range
        band1 = (band1 * 255).astype(np.uint8)
        band2 = (band2 * 255).astype(np.uint8)
        band3 = (band3 * 255).astype(np.uint8)

        # Resize bands to have the same dimensions
        band1 = self.resize_array(band1, band2.shape)
        band3 = self.resize_array(band3, band2.shape)

        # Stack the bands to create the output image
        output_image = np.dstack((band1, band2, band3))

        return output_image

    def calculate_ndvi(self):
        red = gdal.Open(self.band_dict["Band_B04"]).ReadAsArray().astype(np.float32)
        nir = gdal.Open(self.band_dict["Band_B08"]).ReadAsArray().astype(np.float32)

        ndvi = (nir - red) / (nir + red)

        # Normalize NDVI to 0-1 range
        ndvi = (ndvi - np.min(ndvi)) / (np.max(ndvi) - np.min(ndvi))

        # Scale NDVI to 0-255 range
        ndvi = (ndvi * 255).astype(np.uint8)

        return ndvi

    def resize_array(self, array, shape):
        image = Image.fromarray(array)
        image = image.resize((shape[1], shape[0]), resample=Image.BILINEAR)
        return np.array(image)

    def save_output_image(self, output_path, image_data):
        height, width, bands = image_data.shape

        driver = gdal.GetDriverByName("GTiff")
        output_dataset = driver.Create(output_path, width, height, bands, gdal.GDT_Byte)

        for i in range(bands):
            output_band = output_dataset.GetRasterBand(i + 1)
            output_band.WriteArray(image_data[:, :, i])
            output_band.SetNoDataValue(0)

        reference_dataset = gdal.Open(self.band_dict["Band_B04"])
        if reference_dataset is not None:
            output_dataset.SetGeoTransform(reference_dataset.GetGeoTransform())
            output_dataset.SetProjection(reference_dataset.GetProjection())

        output_dataset = None

        print(f"Output image saved: {output_path}")


# Assign S2 folder to a variable
band_folder = r"W:/GIS/SatelliteImagery/Bands"

# Create an instance of SatelliteImageProcessor
processor = SatelliteImageProcessor(band_folder)

# Load band files
processor.load_band_files()

# Process and save additional band combinations
band_combinations = {
    "Natural Colour": ["Band_B04", "Band_B03", "Band_B02"],
    "Colour Infrared": ["Band_B08", "Band_B04", "Band_B03"],
    "Short-Wave Infrared": ["Band_B12", "Band_B08", "Band_B04"],
    "Agriculture": ["Band_B11", "Band_B08", "Band_B02"],
    "Geology": ["Band_B12", "Band_B11", "Band_B04"],
    "Bathymetric": ["Band_B04", "Band_B03", "Band_B01"],
    "Vegetation Index": ["Band_B08", "Band_B04", "Band_B03"],
    "Moisture Index": ["Band_B11", "Band_B08", "Band_B12"],
    "NDVI": ["Band_B04", "Band_B08"]
}

output_folder = r"W:/GIS/SatelliteImagery/Output"

for combination_name, band_keys in band_combinations.items():
    if combination_name == "NDVI":
        output_image = processor.calculate_ndvi()
    else:
        output_image = processor.calculate_band_combination(*band_keys)

    output_path = os.path.join(output_folder, f"{combination_name.lower().replace(' ', '_')}.tif")
    processor.save_output_image(output_path, output_image)

print("Image processing completed.")