improved interface

.gitignore

@@ -2,3 +2,5 @@ __pycache__/
 .ipynb_checkpoints/
 *.ipynb
 images/
+models/
+*.jpg

app.py

@@ -7,41 +7,22 @@ theme = gr.themes.Soft(primary_hue="cyan", secondary_hue="zinc", font=[gr.themes
 
 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.HTML('<img src="file/logos/COP-GEN-logo.png" style="max-width: 300px;">')
+        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)]')
-        
-         gr.Markdown('⚠️ NOTE: This is a protoype Beta model of COP-GEN. It is based on image thumbnails of Major TOM and does not yet support raw source data. The hillshade visualisation is used for elevation. The full model COP-GEN is coming soon.')
+        gr.Markdown('> ## ⚠️ NOTE: This is a prototype Beta model of COP-GEN. It is based on image thumbnails of Major TOM and does not yet support raw source data. The hillshade visualisation is used for elevation. The full model COP-GEN is coming soon.')
     
-    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.Column(elem_classes="Main app"):
 
         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.")
+            gr.Markdown("1. **Generate**: Click the `🏭 Generate` button to synthesize the output. The outputs will be shown below.")
+            gr.Markdown("2. **Define input**: If you want to condition your generation, you can upload your thumbnails manually or you can `🔄 Load` a random sample from Major TOM by clicking the button.")
+            gr.Markdown("3. **Select conditions**: Each input image can be used as a **conditioning** by selecting the `Active` checkbox. If no checkbox is selected, then you will still executy **unconditional generation**.")
+            gr.Markdown("4. **Additional Options**: You can control the number of generation steps (higher number might produce better quality, but will take more time), or set a fixed seed (for reproducible results).")
         
         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")
+            generate_button = gr.Button("🏭 Generate", variant="primary")
 
             gr.Markdown("## Outputs")
             with gr.Row():
@@ -50,25 +31,53 @@ with gr.Blocks(theme=theme) as demo:
                 s1rtc_output = gr.Image(label="S1 RTC (SAR)", interactive=False)
                 dem_output = gr.Image(label="DEM (Elevation)", interactive=False)
 
+            with gr.Row():
+                s2l1c_reuse_button = gr.Button("Reuse S2 L1C as input ⤵️", variant="primary")
+                s2l2a_reuse_button = gr.Button("Reuse S2 L2A as input ⤵️", variant="primary")
+                s1rtc_reuse_button = gr.Button("Reuse S1 RTC as input ⤵️", variant="primary")
+                dem_reuse_button = gr.Button("Reuse DEM as input ⤵️", variant="primary")
+
+            gr.Markdown("---")
+            with gr.Row():
+                gr.Markdown("## Input Conditions (Optional)")
+                load_button = gr.Button("🔄 Load a random sample from Major TOM 🗺", variant="secondary")
+            with gr.Row():
+                s2l1c_input = gr.Image(label="S2 L1C (Optical - Top of Atmosphere)", interactive=True)
+                s2l2a_input = gr.Image(label="S2 L2A (Optical - Bottom of Atmosphere)", interactive=True)
+                s1rtc_input = gr.Image(label="S1 RTC (SAR)", interactive=True)
+                dem_input = gr.Image(label="DEM (Elevation)", interactive=True)
+            gr.Markdown('### Ready? Go back up and press `🏭 Generate` again!')
+
         with gr.Accordion("Advanced Options", open=False) as advanced_options:
             num_inference_steps_slider = gr.Slider(minimum=10, maximum=1000, step=10, value=10, label="Inference Steps")
             with gr.Row():
                 seed_number = gr.Number(value=6378, label="Seed")
                 seed_checkbox = gr.Checkbox(value=True, label="Random")
+        
+        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.")
             
         load_button.click(
             fn=sample_shuffle,
-            outputs=[s2l1c_input, s2l1c_active, s2l2a_input,s2l2a_active, s1rtc_input, s1rtc_active, dem_input, dem_active]
+            outputs=[s2l1c_input,s2l2a_input,s1rtc_input,dem_input]
         )
             
         generate_button.click(
             fn=generate_output,
-            inputs=[s2l1c_input, s2l1c_active,
-                    s2l2a_input, s2l2a_active,
-                    s1rtc_input, s1rtc_active,
-                    dem_input, dem_active,
+            inputs=[s2l1c_input,
+                    s2l2a_input,
+                    s1rtc_input,
+                    dem_input,
                     num_inference_steps_slider, seed_number, seed_checkbox],
             outputs=[s2l1c_output, s2l2a_output, s1rtc_output, dem_output],
         )
 
-demo.queue().launch(share=True)
+        def pass_value(value):
+            return value
+
+        s2l1c_reuse_button.click(fn=pass_value, inputs=[s2l1c_output],outputs=[s2l1c_input])
+        s2l2a_reuse_button.click(fn=pass_value, inputs=[s2l2a_output],outputs=[s2l2a_input])
+        s1rtc_reuse_button.click(fn=pass_value, inputs=[s1rtc_output],outputs=[s1rtc_input])
+        dem_reuse_button.click(fn=pass_value, inputs=[dem_output],outputs=[dem_input])
+
+demo.queue().launch(share=True, allowed_paths=["logos/"])

src/backend.py

@@ -226,10 +226,15 @@ def custom_inference(images, generate_modalities, condition_modalities, num_infe
     
     return results
 
-def generate_output(s2l1c_input, s2l1c_active, s2l2a_input, s2l2a_active, s1rtc_input, s1rtc_active, dem_input, dem_active,num_inference_steps_slider, seed_number, ignore_seed):
+def generate_output(s2l1c_input, s2l2a_input, s1rtc_input, dem_input, num_inference_steps_slider, seed_number, ignore_seed):
 
     seed = seed_number if not ignore_seed else None
 
+    s2l2a_active = s2l2a_input is not None
+    s2l1c_active = s2l1c_input is not None
+    s1rtc_active = s1rtc_input is not None
+    dem_active = dem_input is not None
+
     images=[]
     condition_modalities=[]
     if s2l2a_active:

src/utils.py

@@ -145,7 +145,7 @@ def resize_and_crop(images, image_size=(1068, 1068), crop_size=(256, 256)):
 
     return [img.resize(image_size).crop((left, top, right, bottom)) for img in images]
 
-def sample_shuffle(interface=True):
+def sample_shuffle():
     """
     Randomly selects a 'grid_cell', retrieves corresponding images, and optionally prepares them for an interface.
 
@@ -160,15 +160,7 @@ def sample_shuffle(interface=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
+    return resize_and_crop(get_images(grid_cell))
 
 @spaces.GPU
 def generate_output():