add gradio app

README.md

@@ -10,6 +10,9 @@ title: ClimateGAN
 emoji: 🌎
 colorFrom: blue
 colorTo: green
+sdk: gradio
+sdk_version: 4.6
+app_file: app.py
 # datasets:
 # -
 ---

app.py

@@ -0,0 +1,70 @@
+# based on https://huggingface.co/spaces/NimaBoscarino/climategan/blob/main/app.py # noqa: E501
+# thank you @NimaBoscarino
+
+import os
+import gradio as gr
+import googlemaps
+from skimage import io
+from urllib import parse
+from inferences import ClimateGAN
+
+
+def predict(api_key):
+    def _predict(*args):
+        print("args: ", args)
+        image = place = None
+        if len(args) == 1:
+            image = args[0]
+        else:
+            assert len(args) == 2, "Unknown number of inputs {}".format(len(args))
+            image, place = args
+
+        if api_key and place:
+            geocode_result = gmaps.geocode(place)
+
+            address = geocode_result[0]["formatted_address"]
+            static_map_url = f"https://maps.googleapis.com/maps/api/streetview?size=640x640&location={parse.quote(address)}&source=outdoor&key={api_key}"
+            img_np = io.imread(static_map_url)
+        else:
+            img_np = image
+        flood, wildfire, smog = model.inference(img_np)
+        return img_np, flood, wildfire, smog
+
+    return _predict
+
+
+if __name__ == "__main__":
+
+    api_key = os.environ.get("GMAPS_API_KEY")
+    gmaps = None
+    if api_key is not None:
+        gmaps = googlemaps.Client(key=api_key)
+
+    model = ClimateGAN(model_path="config/model/masker")
+
+    inputs = inputs = [gr.inputs.Image(label="Input Image")]
+    if api_key:
+        inputs += [gr.inputs.Textbox(label="Address or place name")]
+
+    gr.Interface(
+        predict(api_key),
+        inputs=[
+            gr.inputs.Textbox(label="Address or place name"),
+            gr.inputs.Image(label="Input Image"),
+        ],
+        outputs=[
+            gr.outputs.Image(type="numpy", label="Original image"),
+            gr.outputs.Image(type="numpy", label="Flooding"),
+            gr.outputs.Image(type="numpy", label="Wildfire"),
+            gr.outputs.Image(type="numpy", label="Smog"),
+        ],
+        title="ClimateGAN: Visualize Climate Change",
+        description='Climate change does not impact everyone equally. This Space shows the effects of the climate emergency, "one address at a time". Visit the original experience at <a href="https://thisclimatedoesnotexist.com/">ThisClimateDoesNotExist.com</a>.<br>Enter an address or place name, and ClimateGAN will generate images showing how the location could be impacted by flooding, wildfires, or smog.',  # noqa: E501
+        article="<p style='text-align: center'>This project is an unofficial clone of <a href='https://thisclimatedoesnotexist.com/'>ThisClimateDoesNotExist</a> | <a href='https://github.com/cc-ai/climategan'>ClimateGAN GitHub Repo</a></p>",  # noqa: E501
+        # examples=[
+        #     "Vancouver Art Gallery",
+        #     "Chicago Bean",
+        #     "Duomo Siracusa",
+        # ],
+        css=".footer{display:none !important}",
+    ).launch()

climategan/generator.py

@@ -101,6 +101,10 @@ class OmniGenerator(nn.Module):
             if self.verbose > 0:
                 print("  - Add Empty Painter")
 
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
     def __str__(self):
         return strings.generator(self)
 
@@ -381,7 +385,7 @@ class OmniGenerator(nn.Module):
                 val_painter_opts = Dict(yaml.safe_load(f))
 
             # load checkpoint
-            state_dict = torch.load(ckpt_path)
+            state_dict = torch.load(ckpt_path, map_location=self.device)
 
             # create dummy painter from loaded opts
             painter = create_painter(val_painter_opts)

climategan/masker.py

@@ -186,18 +186,18 @@ class MaskSpadeDecoder(nn.Module):
         for i in range(self.num_layers):
             self.spade_blocks.append(
                 SPADEResnetBlock(
-                    int(self.z_nc / (2 ** i)),
+                    int(self.z_nc / (2**i)),
                     int(self.z_nc / (2 ** (i + 1))),
                     cond_nc,
                     spade_use_spectral_norm,
                     spade_param_free_norm,
                     spade_kernel_size,
                     spade_activation,
-                ).cuda()
+                )
             )
         self.spade_blocks = nn.Sequential(*self.spade_blocks)
 
-        self.final_nc = int(self.z_nc / (2 ** self.num_layers))
+        self.final_nc = int(self.z_nc / (2**self.num_layers))
         self.mask_conv = Conv2dBlock(
             self.final_nc,
             1,

inferences.py

@@ -0,0 +1,108 @@
+# based on https://huggingface.co/spaces/NimaBoscarino/climategan/blob/main/inferences.py # noqa: E501
+# thank you @NimaBoscarino
+
+import torch
+from skimage.color import rgba2rgb
+from skimage.transform import resize
+import numpy as np
+
+from climategan.trainer import Trainer
+
+
+def uint8(array):
+    """
+    convert an array to np.uint8 (does not rescale or anything else than changing dtype)
+    Args:
+        array (np.array): array to modify
+    Returns:
+        np.array(np.uint8): converted array
+    """
+    return array.astype(np.uint8)
+
+
+def resize_and_crop(img, to=640):
+    """
+    Resizes an image so that it keeps the aspect ratio and the smallest dimensions
+    is `to`, then crops this resized image in its center so that the output is `to x to`
+    without aspect ratio distortion
+    Args:
+        img (np.array): np.uint8 255 image
+    Returns:
+        np.array: [0, 1] np.float32 image
+    """
+    # resize keeping aspect ratio: smallest dim is 640
+    h, w = img.shape[:2]
+    if h < w:
+        size = (to, int(to * w / h))
+    else:
+        size = (int(to * h / w), to)
+
+    r_img = resize(img, size, preserve_range=True, anti_aliasing=True)
+    r_img = uint8(r_img)
+
+    # crop in the center
+    H, W = r_img.shape[:2]
+
+    top = (H - to) // 2
+    left = (W - to) // 2
+
+    rc_img = r_img[top : top + to, left : left + to, :]
+
+    return rc_img / 255.0
+
+
+def to_m1_p1(img):
+    """
+    rescales a [0, 1] image to [-1, +1]
+    Args:
+        img (np.array): float32 numpy array of an image in [0, 1]
+        i (int): Index of the image being rescaled
+    Raises:
+        ValueError: If the image is not in [0, 1]
+    Returns:
+        np.array(np.float32): array in [-1, +1]
+    """
+    if img.min() >= 0 and img.max() <= 1:
+        return (img.astype(np.float32) - 0.5) * 2
+    raise ValueError(f"Data range mismatch for image: ({img.min()}, {img.max()})")
+
+
+# No need to do any timing in this, since it's just for the HF Space
+class ClimateGAN:
+    def __init__(self, model_path) -> None:
+        torch.set_grad_enabled(False)
+        self.target_size = 640
+        self.trainer = Trainer.resume_from_path(
+            model_path,
+            setup=True,
+            inference=True,
+            new_exp=None,
+        )
+
+    # Does all three inferences at the moment.
+    def inference(self, orig_image):
+        image = self._preprocess_image(orig_image)
+
+        # Retrieve numpy events as a dict {event: array[BxHxWxC]}
+        outputs = self.trainer.infer_all(
+            image,
+            numpy=True,
+            bin_value=0.5,
+        )
+
+        return (
+            outputs["flood"].squeeze(),
+            outputs["wildfire"].squeeze(),
+            outputs["smog"].squeeze(),
+        )
+
+    def _preprocess_image(self, img):
+        # rgba to rgb
+        data = img if img.shape[-1] == 3 else uint8(rgba2rgb(img) * 255)
+
+        # to args.target_size
+        data = resize_and_crop(data, self.target_size)
+
+        # resize() produces [0, 1] images, rescale to [-1, 1]
+        data = to_m1_p1(data)
+        return data