import gradio as gr
import numpy as np
import rasterio
from rasterio.plot import reshape_as_image

def visualize_false_color(input_tiff):
    # Open the input TIFF image using rasterio
    with rasterio.open(input_tiff) as src:
        # Read all bands from the TIFF image
        bands = src.read()

    # Select bands 7, 3, and 2 for false color composit
    false_color_composite = np.stack([bands[6], bands[2], bands[1]], axis=-1)

    # Normalize the image to [0, 255] range
    false_color_normalized = (false_color_composite - false_color_composite.min()) / (false_color_composite.max() - false_color_composite.min()) * 255
    false_color_normalized = false_color_normalized.astype(np.uint8)

    return false_color_normalized

iface = gr.Interface(
    fn=visualize_false_color,
    inputs="file",
    outputs="image",
    description="Upload a satellite image to get a false color composite visualization",
)

iface.launch()