fix fire inference

app.py

@@ -2,9 +2,9 @@
 # thank you @NimaBoscarino
 
 import os
+from datetime import datetime
 from textwrap import dedent
 from urllib import parse
-from requests import get
 
 import googlemaps
 import gradio as gr
@@ -20,8 +20,8 @@ from gradio.components import (
     Row,
     Textbox,
 )
+from requests import get
 from skimage import io
-from datetime import datetime
 
 from climategan_wrapper import ClimateGAN
 
@@ -76,9 +76,9 @@ TEXTS = [
               |  
             Read the original
             <a
-                href='https://openreview.net/forum?id=EZNOb_uNpJk'
+                href='https://arxiv.org/abs/2110.02871'
                 target='_blank'>
-            ICLR 2021 ClimateGAN paper
+            ICLR 2022 ClimateGAN paper
             </a>
         </p>
         """
@@ -217,7 +217,6 @@ def predict(cg: ClimateGAN, api_key):
 
 
 if __name__ == "__main__":
-
     ip = get("https://api.ipify.org").content.decode("utf8")
     print("My public IP address is: {}".format(ip))
 

climategan/fire.py

@@ -1,7 +1,8 @@
-import torch
-import torch.nn.functional as F
 import random
+
 import kornia
+import torch
+import torch.nn.functional as F
 from torchvision.transforms.functional import adjust_brightness, adjust_contrast
 
 from climategan.tutils import normalize, retrieve_sky_mask
@@ -105,9 +106,9 @@ def add_fire(x, seg_preds, fire_opts):
     kernel_size = (fire_opts.get("kernel_size", 301), fire_opts.get("kernel_size", 301))
     sigma = (fire_opts.get("kernel_sigma", 150.5), fire_opts.get("kernel_sigma", 150.5))
     border_type = "reflect"
-    kernel = torch.unsqueeze(
-        kornia.filters.kernels.get_gaussian_kernel2d(kernel_size, sigma), dim=0
-    ).to(x.device)
+    kernel = kornia.filters.kernels.get_gaussian_kernel2d(kernel_size, sigma)
+    if kernel.ndim == 2:
+        kernel = kernel.unsqueeze(0)
     sky_mask = filter2d(sky_mask, kernel, border_type)
 
     filter_ = torch.ones(wildfire_tens.shape, device=x.device)

requirements.txt

@@ -2,6 +2,7 @@ gradio==3.44.1
 torch
 torch-optimizer
 torchvision
+accelerate
 addict
 aiohttp
 aiosignal