front end ready

app.py

@@ -0,0 +1,69 @@
+import gradio as gr
+from src.utils import *
+
+theme = gr.themes.Soft(primary_hue="amber", secondary_hue="orange", font=[gr.themes.GoogleFont("Source Sans 3", weights=(400, 600)),'arial'])
+
+with gr.Blocks(theme=theme) as demo:
+    with gr.Column(elem_classes="header"):
+        gr.Markdown("# 🗾 COP-GEN-Beta: Unified Generative Modelling of COPernicus Imagery Thumbnails")
+        gr.Markdown("### Miguel Espinosa, Valerio Marsocci, Yuru Jia, Elliot J. Crowley, Mikolaj Czerkawski")
+        gr.Markdown('[[Website](https://miquel-espinosa.github.io/cop-gen-beta/)] [[GitHub](https://github.com/miquel-espinosa/COP-GEN-Beta)] [[Model](https://huggingface.co/mespinosami/COP-GEN-Beta)] [[Dataset](https://huggingface.co/Major-TOM)]')
+
+    with gr.Column(elem_classes="abstract"):
+
+        with gr.Accordion("Abstract", open=False) as abstract:
+            gr.Markdown("In remote sensing, multi-modal data from various sensors capturing the same scene offers rich opportunities, but learning a unified representation across these modalities remains a significant challenge. Traditional methods have often been limited to single or dual-modality approaches. In this paper, we introduce COP-GEN-Beta, a generative diffusion model trained on optical, radar, and elevation data from the Major TOM dataset. What sets COP-GEN-Beta apart is its ability to map any subset of modalities to any other, enabling zero-shot modality translation after training. This is achieved through a sequence-based diffusion transformer, where each modality is controlled by its own timestep embedding. We extensively evaluate COP-GEN-Beta on thumbnail images from the Major TOM dataset, demonstrating its effectiveness in generating high-quality samples. Qualitative and quantitative evaluations validate the model's performance, highlighting its potential as a powerful pre-trained model for future remote sensing tasks.") # Replace with your abstract text
+
+        with gr.Accordion("Instructions", open=False) as abstract:
+            gr.Markdown("1. **Define input**: You can upload your thumbnails manually or you can get a random sample from Major TOM by clicking the button.")
+            gr.Markdown("2. **Select conditions**: Each input image can be used as a **conditioning** by selecting the `Active` checkbox. If no checkbox is selected, then you will observe **unconditional generation**.")
+            gr.Markdown("3. **Generate**: Click the `Generate` button to synthesize the output. The outputs will be shown below.")
+        
+        with gr.Column():
+            with gr.Row():
+                gr.Markdown("## Inputs (Optional)")
+                load_button = gr.Button("Load a random sample from Major TOM 🗺", variant="secondary")
+            with gr.Row():
+                with gr.Column():
+                    s2l1c_input = gr.Image(label="S2 L1C (Optical - Top of Atmosphere)", interactive=True)
+                    s2l1c_active = gr.Checkbox(value=False, label="Active", interactive=True)
+                with gr.Column():
+                    s2l2a_input = gr.Image(label="S2 L2A (Optical - Bottom of Atmosphere)", interactive=True)
+                    s2l2a_active = gr.Checkbox(value=False, label="Active", interactive=True)
+                with gr.Column():
+                    s1rtc_input = gr.Image(label="S1 RTC (SAR)", interactive=True)
+                    s1rtc_active = gr.Checkbox(value=False, label="Active", interactive=True)
+                with gr.Column():
+                    dem_input = gr.Image(label="DEM (Elevation)", interactive=True)
+                    dem_active = gr.Checkbox(value=False, label="Active", interactive=True)
+
+            generate_button = gr.Button("Generate", variant="primary")
+
+            gr.Markdown("## Outputs")
+            with gr.Row():
+                s2l1c_output = gr.Image(label="S2 L1C (Optical - Top of Atmosphere)", interactive=False)
+                s2l2a_output = gr.Image(label="S2 L2A (Optical - Bottom of Atmosphere)", interactive=False)
+                s1rtc_output = gr.Image(label="S1 RTC (SAR)", interactive=False)
+                dem_output = gr.Image(label="DEM (Elevation)", interactive=False)
+
+        with gr.Accordion("Advanced Options", open=False) as advanced_options:
+            num_inference_steps_slider = gr.Slider(minimum=10, maximum=1000, step=10, value=50, label="Inference Steps")
+            guidance_scale_slider = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, value=7.5, label="Guidance Scale")
+            with gr.Row():
+                seed_number = gr.Number(value=6378, label="Seed")
+                seed_checkbox = gr.Checkbox(value=True, label="Random")
+            
+        load_button.click(
+            fn=sample_shuffle,
+            outputs=[s2l1c_input, s2l1c_active, s2l2a_input,s2l2a_active, s1rtc_input, s1rtc_active, dem_input, dem_active]
+        )
+            
+        generate_button.click(
+            #fn=generate_output,
+            inputs=[s2l1c_input, s2l2a_input, s1rtc_input, dem_input, num_inference_steps_slider, guidance_scale_slider, seed_number, seed_checkbox],
+            outputs=[s2l1c_output, s2l2a_output, s1rtc_output, dem_output],
+        )
+
+demo.launch()
+
+demo.launch(share=True)

src/utils.py

@@ -0,0 +1,170 @@
+import os
+import pandas as pd
+
+# GLOBAL VARIABLES
+if os.path.isfile('data/s2l2a_metadata.parquet'):
+    l2a_meta_path = 'data/s2l2a_metadata.parquet'
+else:
+    DATASET_NAME = 'Major-TOM/Core-S2L2A'
+    l2a_meta_path = 'https://huggingface.co/datasets/{}/resolve/main/metadata.parquet'.format(DATASET_NAME)
+
+if os.path.isfile('data/s2l1c_metadata.parquet'):
+    l1c_meta_path = 'data/s2l1c_metadata.parquet'
+else:
+    DATASET_NAME = 'Major-TOM/Core-S2L1C'
+    l1c_meta_path = 'https://huggingface.co/datasets/{}/resolve/main/metadata.parquet'.format(DATASET_NAME)
+
+if os.path.isfile('/s1rtc_metadata.parquet'):
+    rtc_meta_path = 'data/s1rtc_metadata.parquet'
+else:
+    DATASET_NAME = 'Major-TOM/Core-S1RTC'
+    rtc_meta_path = 'https://huggingface.co/datasets/{}/resolve/main/metadata.parquet'.format(DATASET_NAME)
+
+if os.path.isfile('helpers/dem_metadata.parquet'):
+    dem_meta_path = 'data/dem_metadata.parquet'
+else:
+    DATASET_NAME = 'Major-TOM/Core-DEM'
+    dem_meta_path = 'https://huggingface.co/datasets/{}/resolve/main/metadata.parquet'.format(DATASET_NAME)
+
+print('Loading Major TOM meta...')
+l2a_df = pd.read_parquet(l2a_meta_path)
+l1c_df = pd.read_parquet(l1c_meta_path)
+rtc_df = pd.read_parquet(rtc_meta_path)
+dem_df = pd.read_parquet(dem_meta_path)
+
+# skip files with missing parts
+l2a_df = l2a_df[l2a_df.nodata == 0]
+l1c_df = l1c_df[l1c_df.nodata == 0]
+rtc_df = rtc_df[rtc_df.nodata == 0]
+dem_df = dem_df[dem_df.nodata == 0]
+
+# collect grid_cells, drop duplicates, and extract grid cell column only
+grid_cell_df = l2a_df[l2a_df.grid_cell.isin(l1c_df.grid_cell) &l2a_df.grid_cell.isin(rtc_df.grid_cell) & l2a_df.grid_cell.isin(dem_df.grid_cell)]
+gird_cell_df = grid_cell_df.drop_duplicates(subset=['grid_cell'])
+grid_cell_df = grid_cell_df.grid_cell
+print('[DONE]')
+
+import pyarrow.parquet as pq
+import fsspec
+from fsspec.parquet import open_parquet_file
+from io import BytesIO
+from PIL import Image
+import random
+
+def row2image(row, fullrow_read=True):
+    """
+    Extracts an image from a specific row in a Parquet file.
+
+    Args:
+        row: A row object containing information about the Parquet file and row index.
+            It is expected to have attributes 'parquet_row' (the row index within the Parquet file)
+            and 'parquet_url' (the URL or path to the Parquet file).
+        fullrow_read (bool, optional): Determines whether to read the entire Parquet file or just the 'thumbnail' column initially.
+            Defaults to True.
+            - If True, it opens the Parquet file using fsspec and reads the entire file.
+            - If False, it uses fsspec.parquet.open_parquet_file to only open the 'thumbnail' column.
+
+    Returns:
+        PIL.Image.Image: An Image object loaded from the 'thumbnail' data in the specified row.
+    """
+    parquet_row = row.parquet_row
+    parquet_url = row.parquet_url
+
+    if fullrow_read:
+        # Option 1: Read the entire Parquet file
+        f = fsspec.open(parquet_url)
+        temp_path = f.open()
+    else:
+        # Option 2: Read only the 'thumbnail' column initially
+        temp_path = open_parquet_file(parquet_url, columns=["thumbnail"])
+
+    with pq.ParquetFile(temp_path) as pf:
+        first_row_group = pf.read_row_group(parquet_row, columns=['thumbnail'])
+
+    stream = BytesIO(first_row_group['thumbnail'][0].as_py())
+    return Image.open(stream)
+
+# Example usage (assuming 'dem_df' is a Pandas DataFrame with the required structure):
+# row2image(dem_df.iloc[1000])
+
+def get_rows(grid_cell):
+    """
+    Retrieves the first row from multiple DataFrames based on a given 'grid_cell' value.
+
+    Args:
+        grid_cell: The value to filter the DataFrames by in the 'grid_cell' column.
+
+    Returns:
+        tuple: A tuple containing the first matching row from each of the following DataFrames:
+               l2a_df, l1c_df, rtc_df, and dem_df. It assumes these DataFrames are defined in the scope.
+               Each element of the tuple is a Pandas Series representing a row.
+    """
+    return l2a_df[l2a_df.grid_cell == grid_cell].iloc[0], \
+           l1c_df[l1c_df.grid_cell == grid_cell].iloc[0], \
+           rtc_df[rtc_df.grid_cell == grid_cell].iloc[0], \
+           dem_df[dem_df.grid_cell == grid_cell].iloc[0]
+
+def get_images(grid_cell):
+    """
+    Retrieves images corresponding to a specific 'grid_cell' by calling get_rows and row2image.
+
+    Args:
+        grid_cell: The grid cell identifier to fetch images for.
+
+    Returns:
+        list: A list of PIL.Image.Image objects, where each image is extracted from the rows
+              returned by the get_rows function for the given grid cell.
+    """
+    img_rows = get_rows(grid_cell)
+
+    imgs = []
+    for row in img_rows:
+        imgs.append(row2image(row))
+
+    return imgs
+
+def resize_and_crop(images, image_size=(1068, 1068), crop_size=(256, 256)):
+    """
+    Resizes a list of images to a specified size and then crops a random portion from each.
+
+    Args:
+        images (list): A list of PIL.Image.Image objects to be processed.
+        image_size (tuple, optional): The target size (width, height) to resize the images to.
+            Defaults to (1068, 1068).
+        crop_size (tuple, optional): The size (width, height) of the random crop to be taken
+            from the resized images. Defaults to (256, 256).
+
+    Returns:
+        list: A list of PIL.Image.Image objects, where each image has been resized and then cropped.
+    """
+    left = random.randint(0, image_size[0] - crop_size[0])
+    top = random.randint(0, image_size[1] - crop_size[1])
+    right = left + crop_size[0]
+    bottom = top + crop_size[1]
+
+    return [img.resize(image_size).crop((left, top, right, bottom)) for img in images]
+
+def sample_shuffle(interface=True):
+    """
+    Randomly selects a 'grid_cell', retrieves corresponding images, and optionally prepares them for an interface.
+
+    Args:
+        interface (bool, optional): If True, the function returns a list where each image is followed by True.
+            This might be intended for an interface that expects an image and a boolean flag.
+            If False, it returns just the list of processed images. Defaults to True.
+
+    Returns:
+        list: If interface is False, returns a list of resized and cropped PIL.Image.Image objects.
+              If interface is True, returns a list where each image is followed by the boolean value True.
+    """
+    grid_cell = grid_cell_df.sample().iloc[0]
+
+    images = resize_and_crop(get_images(grid_cell))
+
+    if not interface:
+        return images
+    else:
+        out = []
+        for el in images:
+            out += [el, True]
+        return out