Upload 3 files

evoland/run.py

@@ -0,0 +1,31 @@
+import torch
+import opensr_test
+import matplotlib.pyplot as plt
+from utils import load_evoland, run_evoland
+
+# Load the model
+model = load_evoland()
+
+# Load the dataset
+dataset = opensr_test.load("naip")
+lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]
+
+# Predict a image
+results = run_evoland(
+    model=model, 
+    lr=lr_dataset[4],
+    hr=hr_dataset[4]
+)
+
+
+# Display the results
+fig, ax = plt.subplots(1, 3, figsize=(10, 5))
+ax[0].imshow(results["lr"].transpose(1, 2, 0)/3000)
+ax[0].set_title("LR")
+ax[0].axis("off")
+ax[1].imshow(results["sr"].transpose(1, 2, 0)/3000)
+ax[1].set_title("SR")
+ax[1].axis("off")
+ax[2].imshow(results["hr"].transpose(1, 2, 0) / 3000)
+ax[2].set_title("HR")
+plt.show()

evoland/utils.py

@@ -0,0 +1,73 @@
+import torch
+import pickle
+import numpy as np
+import opensr_test
+import onnxruntime as ort
+from typing import List, Union
+
+def load_evoland() -> np.ndarray:
+    # ONNX inference session options
+    so = ort.SessionOptions()
+    so.intra_op_num_threads = 10
+    so.inter_op_num_threads = 10
+    so.use_deterministic_compute = True
+
+    # Execute on cpu only
+    ep_list = ["CPUExecutionProvider"]
+    ep_list.insert(0, "CUDAExecutionProvider")
+
+    ort_session = ort.InferenceSession(
+        "weights/carn_3x3x64g4sw_bootstrap.onnx",
+        sess_options=so,
+        providers=ep_list
+    )
+    ort_session.set_providers(["CPUExecutionProvider"])
+    ro = ort.RunOptions()    
+
+    return [ort_session, ro]
+
+
+def run_evoland(
+    model: List,
+    lr: np.ndarray,
+    hr: np.ndarray 
+) -> dict:
+    
+    ort_session, ro = model
+
+    # Bands to use
+    bands = [1, 2, 3, 7, 4, 5, 6, 8, 10, 11]
+    lr = lr[bands]
+
+    if lr.shape[1] == 121:
+        # add padding
+        lr = torch.nn.functional.pad(
+            torch.from_numpy(lr[None]).float(),
+            pad=(3, 4, 3, 4),
+            mode='reflect'
+        ).squeeze().cpu().numpy()
+
+        # run the model
+        sr = ort_session.run(
+            None, 
+            {"input": lr[None]},
+            run_options=ro
+        )[0].squeeze()
+
+        # remove padding
+        sr = sr[:, 3*4:-4*4, 3*4:-4*4].astype(np.uint16)
+        lr = lr[:, 3:-4, 3:-4].astype(np.uint16)
+    else:
+        # run the model
+        sr = ort_session.run(
+            None, 
+            {"input": lr[None]},
+            run_options=ro
+        )[0].squeeze()
+    
+    # Run the model
+    return {
+        "lr": lr[[2, 1, 0]],
+        "sr": sr[[2, 1, 0]],
+        "hr": hr[0:3]
+    }

evoland/weights/carn_3x3x64g4sw_bootstrap.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:361233ee1cb1977a6e2c41e3bb40eb55cea8bdfe001e945d3b74b6eecaff6516
+size 10103338