update UI

app.py

@@ -2,26 +2,133 @@
 # thank you @NimaBoscarino
 
 import os
-import gradio as gr
-import googlemaps
-from skimage import io
+from textwrap import dedent
 from urllib import parse
+
+import googlemaps
+import gradio as gr
 import numpy as np
+from gradio.components import (
+    HTML,
+    Button,
+    Column,
+    Dropdown,
+    Image,
+    Markdown,
+    Radio,
+    Row,
+    Textbox,
+)
+from skimage import io
+
 from climategan_wrapper import ClimateGAN
-from textwrap import dedent
+
+HTMLS = [
+    dedent(
+        """
+        <p>
+            Climate change does not impact everyone equally.
+            This Space shows the effects of the climate emergency,
+            "one address at a time".
+        </p>
+        <p>
+            Visit the original experience at
+            <a href="https://thisclimatedoesnotexist.com/">
+                ThisClimateDoesNotExist.com
+            </a>
+        </p>
+        <br>
+        <p>
+            Enter an address or upload a Street View image, and ClimateGAN
+            will generate images showing how the location could be impacted
+            by flooding, wildfires, or smog if it happened there.
+        </p>
+        <br>
+        <p>
+            This is <strong>not</strong> an exercise in climate prediction,
+            rather an exercise of empathy, to put yourself in other's shoes,
+            as if Climate Change came crushing on your doorstep.
+        </p>
+        """
+    ),
+    dedent(
+        """
+        <br><br><br><br>
+        <p style='text-align: center'>
+            Visit
+            <a href='https://thisclimatedoesnotexist.com/'>
+                ThisClimateDoesNotExist.com
+            </a>
+            for more information
+            &nbsp;&nbsp;|&nbsp;&nbsp;
+            Original
+            <a href='https://github.com/cc-ai/climategan'>
+                ClimateGAN GitHub Repo
+            </a>
+        </p>
+        <br>
+        <p>
+            After you have selected an image and started the inference you
+            will see all the outputs of ClimateGAN, including intermediate
+            outputs such as the flood mask, the segmentation map and the
+            depth maps used to produce the 3 events.
+        </p>
+        <br>
+        <p>
+            This Space makes use of recent Stable Diffusion in-painting
+            pipelines to replace ClimateGAN's original Painter. If you
+            select 'Both' painters, you will see a comparison
+        </p>
+        <br>
+        <br>
+        <p>
+            Read the original
+            <a
+                href='https://openreview.net/forum?id=EZNOb_uNpJk'
+                target='_blank'>
+            ICLR 2021 ClimateGAN paper
+            </a>
+        </p>
+        """
+    ),
+]
+CSS = dedent(
+    """
+    a {
+        color: #0088ff;
+        text-decoration: underline;
+    }
+    strong {
+        color: #c34318;
+    }
+    """
+)
+
+
+def toggle(radio):
+    if "address" in radio.lower():
+        return [
+            gr.update(visible=True),
+            gr.update(visible=False),
+            gr.update(visible=True),
+        ]
+    else:
+        return [
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(visible=True),
+        ]
 
 
 def predict(cg: ClimateGAN, api_key):
     def _predict(*args):
-        image = place = painter = None
-        if len(args) == 2:
-            image = args[0]
-            painter = args[1]
+        image = place = painter = radio = None
+        if api_key:
+            radio, image, place, painter = args
         else:
-            assert len(args) == 3, "Unknown number of inputs {}".format(len(args))
-            image, place, painter = args
+            image, painter = args
 
-        if api_key and place:
+        if api_key and place and "address" in radio.lower():
             geocode_result = gmaps.geocode(place)
 
             address = geocode_result[0]["formatted_address"]
@@ -91,170 +198,72 @@ if __name__ == "__main__":
     )
     cg._setup_stable_diffusion()
 
-    with gr.Blocks(
-        css=dedent(
-            """
-            a {
-                color: #0088ff;
-                text-decoration: underline;
-            }
-            strong {
-                color: #c34318;
-            }
-            """
-        )
-    ) as blocks:
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("# ClimateGAN: Visualize Climate Change")
-                gr.HTML(
-                    dedent(
-                        """
-                    <p>
-                        Climate change does not impact everyone equally.
-                        This Space shows the effects of the climate emergency,
-                        "one address at a time".
-                    </p>
-
-                    <p>
-                        Visit the original experience at
-                        <a href="https://thisclimatedoesnotexist.com/">
-                            ThisClimateDoesNotExist.com
-                        </a>
-                    </p>
-
-                    <br>
-
-                    <p>
-                        Enter an address or upload a Street View image, and ClimateGAN
-                        will generate images showing how the location could be impacted
-                        by flooding, wildfires, or smog if it happened there.
-                    </p>
-
-                    <br>
+    radio = address = None
+    pred_ins = []
+    pred_outs = []
 
-                    <p>
-                        This is <strong>not</strong> an exercise in climate prediction,
-                        rather an exercise of empathy, to put yourself in other's shoes,
-                        as if Climate Change came crushing on your doorstep.
-                    </p>
-                        """
-                    )
-                )
-            with gr.Column():
-                gr.HTML(
-                    dedent(
-                        """
-
-                    <br>
-                    <br>
-                    <br>
-
-
-                    <p style='text-align: center'>
-                        Visit
-                        <a href='https://thisclimatedoesnotexist.com/'>
-                            ThisClimateDoesNotExist.com
-                        </a>
-                        for more information
-                        &nbsp;&nbsp;|&nbsp;&nbsp;
-                        Original
-                        <a href='https://github.com/cc-ai/climategan'>
-                            ClimateGAN GitHub Repo
-                        </a>
-                    </p>
-
-                    <br>
-
-                    <p>
-                        After you have selected an image and started the inference you
-                        will see all the outputs of ClimateGAN, including intermediate
-                        outputs such as the flood mask, the segmentation map and the
-                        depth maps used to produce the 3 events.
-                    </p>
-
-                    <br>
-
-                    <p>
-                        This Space makes use of recent Stable Diffusion in-painting
-                        pipelines to replace ClimateGAN's original Painter. If you
-                        select 'Both' painters, you will see a comparison
-                    </p>
-
-                    <br>
-
-                    <br>
-
-                    <p>
-                        Read the original
-                        <a
-                            href='https://openreview.net/forum?id=EZNOb_uNpJk'
-                            target='_blank'>
-                        ICLR 2021 ClimateGAN paper
-                        </a>
-                    </p>
-                    """
-                    )
-                )
-        with gr.Row():
-            gr.Markdown("## Inputs")
-        with gr.Row():
-            with gr.Column():
-                inputs = [gr.inputs.Image(label="Input Image")]
-            with gr.Column():
+    with gr.Blocks(css=CSS) as app:
+        with Row():
+            with Column():
+                Markdown("# ClimateGAN: Visualize Climate Change")
+                HTML(HTMLS[0])
+            with Column():
+                HTML(HTMLS[1])
+        with Row():
+            Markdown("## Inputs")
+        with Row():
+            with Column():
                 if api_key:
-                    inputs += [gr.inputs.Textbox(label="Address or place name")]
-                inputs += [
-                    gr.inputs.Dropdown(
+                    radio = Radio(["From Address", "From Image"], label="Input Type")
+                    pred_ins += [radio]
+                im_inp = Image(label="Input Image", visible=not api_key)
+                pred_ins += [im_inp]
+                if api_key:
+                    address = Textbox(label="Address or place name", visible=False)
+                    pred_ins += [address]
+            with Column():
+                pred_ins += [
+                    Dropdown(
                         choices=[
                             "ClimateGAN Painter",
                             "Stable Diffusion Painter",
                             "Both",
                         ],
                         label="Choose Flood Painter",
-                        default="Both",
+                        value="Both",
                     )
                 ]
-                btn = gr.Button("See for yourself!", label="Run", variant="primary")
-        with gr.Row():
-            gr.Markdown("## Outputs")
-        with gr.Row():
-            outputs = []
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Original image"),
-            )
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Masked input image"),
-            )
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Segmentation map"),
-            )
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Depth map"),
-            )
-        with gr.Row():
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="ClimateGAN-Flooded image"),
-            )
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Stable Diffusion-Flooded image"),
-            )
-            outputs.append(
-                gr.outputs.Image(
+                btn = Button(
+                    "See for yourself!",
+                    label="Run",
+                    variant="primary",
+                    visible=not api_key,
+                )
+        with Row():
+            Markdown("## Outputs")
+        with Row():
+            pred_outs += [Image(type="numpy", label="Original image")]
+            pred_outs += [Image(type="numpy", label="Masked input image")]
+            pred_outs += [Image(type="numpy", label="Segmentation map")]
+            pred_outs += [Image(type="numpy", label="Depth map")]
+        with Row():
+            pred_outs += [Image(type="numpy", label="ClimateGAN-Flooded image")]
+            pred_outs += [Image(type="numpy", label="Stable Diffusion-Flooded image")]
+            pred_outs += [
+                Image(
                     type="numpy",
                     label="Stable Diffusion-Flooded image (restricted to masked area)",
                 )
-            ),
-        with gr.Row():
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Comparison of previous images"),
-            )
-        with gr.Row():
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Wildfire"),
-            )
-            outputs.append(
-                gr.outputs.Image(type="numpy", label="Smog"),
-            )
-        btn.click(predict(cg, api_key), inputs=inputs, outputs=outputs)
-    blocks.launch()
+            ]
+        with Row():
+            pred_outs += [Image(type="numpy", label="Comparison of flood images")]
+        with Row():
+            pred_outs += [Image(type="numpy", label="Wildfire")]
+            pred_outs += [Image(type="numpy", label="Smog")]
+            Image(type="numpy", label="Empty on purpose", interactive=False)
+        btn.click(predict(cg, api_key), inputs=pred_ins, outputs=pred_outs)
+
+        if api_key:
+            radio.change(toggle, inputs=[radio], outputs=[address, im_inp, btn])
+
+    app.launch()

climategan_wrapper.py

@@ -241,7 +241,7 @@ class ClimateGAN:
         """
         if self.dev_mode:
             return {
-                "input": np.random.randint(0, 255, (640, 640, 3)),
+                "input": orig_image,
                 "mask": np.random.randint(0, 255, (640, 640)),
                 "masked_input": np.random.randint(0, 255, (640, 640, 3)),
                 "climategan_flood": np.random.randint(0, 255, (640, 640, 3)),