Refactor, polish (WIP)

app.py

@@ -1,41 +1,27 @@
 import os
 import gradio as gr
 import googlemaps
-import requests
-from PIL import Image
+from skimage import io
 
 os.system('gdown https://drive.google.com/u/0/uc?id=18OCUIy7JQ2Ow_-cC5xn_hhDn-Bp45N1K')
 os.system('unzip release-github-v1.zip')
 os.system('mkdir config')
 os.system('mv model config')
 
+from inferences import ClimateGAN
+
 API_KEY = os.environ.get("API_KEY")
 gmaps = googlemaps.Client(key=API_KEY)
+model = ClimateGAN(model_path="config/model/masker")
 
 def predict(place):
-    # I don't think I need to do any error handling, I'll just let Gradio manage it. (Test it with a bad place though)
     geocode_result = gmaps.geocode(place)
     loc = geocode_result[0]['geometry']['location']
     static_map_url = f"https://maps.googleapis.com/maps/api/streetview?size=400x400&location={loc['lat']},{loc['lng']}&fov=80&heading=70&pitch=0&key={API_KEY}"
 
-    response = requests.get(static_map_url)
-    img = response.content
-
-    file = open("./inputs/loc_image.png", "wb")
-    file.write(img)
-    file.close()
-
-    # Next improvement is to refactor apply_events so that I can load the model in memory when the interface launces
-    # Would save ~10 seconds according to the logs.
-    # I would also like to not have to write the image to the filesystem, and instead just hold it in memory. As it is right now it would mess up with more than one person using it.
-    os.system('python apply_events.py -b 1 -i ./inputs -r config/model/masker --output_path ./outputs --overwrite')
-
-    # Also if I refactor then I wouldn't have to read from the filesystem...
-    # I would like to show all three images
-    os.system('ls -R ./outputs')
-
-    out_img = Image.open("./outputs/loc_image_flood_640.png")
-    return out_img # I actually want to return all 3 images.
+    img_np = io.imread(static_map_url)
+    flood, wildfire, smog = model.inference(img_np)
+    return img_np, flood, wildfire, smog
 
 
 gr.Interface(
@@ -43,10 +29,14 @@ gr.Interface(
     inputs=[
         gr.inputs.Textbox(label="Address or place name")
     ],
-    outputs="image",
+    outputs=["image", "image", "image", "image"],
     title="ClimateGAN",
-    # description="These CryptoPunks do not exist. Generate random punks with an initial seed!",
-    # article="<p style='text-align: center'><a href='https://arxiv.org/pdf/1511.06434.pdf'>Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks</a> | <a href='https://github.com/teddykoker/cryptopunks-gan'>Github Repo</a></p>",   ---- should be ClimateGAN citation
-    # examples=[[123, 15], [42, 29], [456, 8], [1337, 35]],
-    # css=".footer{display:none !important}",
+    description="Enter an address or place name, and ClimateGAN will generate images showing how the location could be impacted by flooding, wildfires, or smog.",
+    article="<p style='text-align: center'>This project is a clone of <a href='https://thisclimatedoesnotexist.com/'>ThisClimateDoesNotExist</a> | <a href='https://github.com/cc-ai/climategan'>ClimateGAN GitHub Repo</a></p>",
+    examples=[
+        "Kafka's Great Northern Way, Vancouver",
+        "Simon Fraser University",
+        "Duomo, Milano"
+    ],
+    css=".footer{display:none !important}",
 ).launch()

inferences.py

@@ -0,0 +1,107 @@
+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, new_size = self._preprocess_image(orig_image)
+
+        image = np.stack(image)
+        # Retreive numpy events as a dict {event: array[BxHxWxC]}
+        outputs = self.trainer.infer_all(
+            image,
+            numpy=True,
+            bin_value=0.5,
+        )
+
+        return outputs['flood'], outputs['wildfire'], outputs['smog']
+
+    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)
+        new_size = (self.target_size, self.target_size)
+
+        # resize() produces [0, 1] images, rescale to [-1, 1]
+        data = to_m1_p1(data)
+        return data, new_size