import gradio as gr
import rasterio
import matplotlib.pyplot as plt

def visualize_image(file_path):
    # Read the 10-band image using rasterio
    dataset = rasterio.open(file_path)

    # Get the number of bands in the image
    num_bands = dataset.count

    # Calculate the number of rows required for subplots
    num_rows = (num_bands + 1) // 2

    # Create subplots to visualize each band
    fig, axes = plt.subplots(num_rows, 2, figsize=(10, 6 * num_rows))

    # Loop through each band and plot it
    for band_index in range(1, num_bands + 1):
        # Read the band data
        band = dataset.read(band_index)

        # Determine the subplot position
        row_index = (band_index - 1) // 2
        col_index = (band_index - 1) % 2

        # Plot the band on the corresponding subplot
        axes[row_index, col_index].imshow(band, cmap='terrain')
        axes[row_index, col_index].set_title(f"Band {band_index}")

    # Remove any empty subplots if there are odd number of bands
    if num_bands % 2 != 0:
        fig.delaxes(axes[num_rows - 1, 1])

    # Adjust the layout and display the visualization
    plt.tight_layout()

    # Save the figure to a temporary file and close it to release resources
    temp_file = "/tmp/visualization.png"
    plt.savefig(temp_file)
    plt.close()

    # Return the temporary file path, which will be used by Gradio to display the visualization
    return temp_file

# Interface for the Gradio app
iface = gr.Interface(
    fn=visualize_image,
    inputs=gr.inputs.File(),
    outputs="image",
    live=True,
    capture_session=True,
    title="Earth Engine Image Segmentation",
    description="Upload a TIFF image and segmentation all its bands.",
)

# Launch the Gradio app
iface.launch()