Fix image output dims

app.py

@@ -17,7 +17,7 @@ model = ClimateGAN(model_path="config/model/masker")
 def predict(place):
     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}"
+    static_map_url = f"https://maps.googleapis.com/maps/api/streetview?size=640x640&location={loc['lat']},{loc['lng']}&source=outdoor&key={API_KEY}"
 
     img_np = io.imread(static_map_url)
     flood, wildfire, smog = model.inference(img_np)
@@ -29,14 +29,19 @@ gr.Interface(
     inputs=[
         gr.inputs.Textbox(label="Address or place name")
     ],
-    outputs=["image", "image", "image", "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",
     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"
+        "Vancouver Art Gallery",
+        "Chicago Bean",
+        "Duomo Siracusa"
     ],
     css=".footer{display:none !important}",
-).launch()
+).launch(cache_examples=True)

inferences.py

@@ -82,9 +82,8 @@ class ClimateGAN():
 
     # Does all three inferences at the moment.
     def inference(self, orig_image):
-        image, new_size = self._preprocess_image(orig_image)
+        image = self._preprocess_image(orig_image)
 
-        image = np.stack(image)
         # Retreive numpy events as a dict {event: array[BxHxWxC]}
         outputs = self.trainer.infer_all(
             image,
@@ -92,7 +91,11 @@ class ClimateGAN():
             bin_value=0.5,
         )
 
-        return outputs['flood'], outputs['wildfire'], outputs['smog']
+        return (
+            outputs['flood'].squeeze(),
+            outputs['wildfire'].squeeze(),
+            outputs['smog'].squeeze()
+        )
 
     def _preprocess_image(self, img):
         # rgba to rgb
@@ -100,8 +103,7 @@ class ClimateGAN():
 
         # 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
+        return data