Upload 4 files

opensr_model/run.py

@@ -0,0 +1,30 @@
+import opensr_test
+import matplotlib.pyplot as plt
+from utils import create_opensr_model, run_opensr_model
+
+# Load the model
+model = create_opensr_model(device="cpu")
+
+# Load the dataset
+dataset = opensr_test.load("naip")
+lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]
+
+# Run the model
+results = run_opensr_model(
+    model=model,
+    lr=lr_dataset[7],
+    hr=hr_dataset[7],
+    device="cpu"
+)
+
+# 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()

opensr_model/utils.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+import opensr_model
+from typing import Union
+
+def create_opensr_model(
+    device: Union[str, torch.device] = "cpu"
+) -> opensr_model:
+    """ Create the super image model
+    Returns:
+        HanModel: The super image model
+    """
+    model = opensr_model.SRLatentDiffusion(device=device)
+    model.load_pretrained("./weights/opensr_10m_v4_v5.ckpt")
+    model.eval()
+    return model
+
+
+def run_opensr_model(
+    model: opensr_model,
+    lr: np.ndarray,
+    hr: np.ndarray,
+    device: Union[str, torch.device] = "cpu"
+) -> dict:
+    # Convert the input to torch tensors    
+    lr_img = torch.from_numpy(lr[[3, 2, 1, 7]] / 10000).to(device).float()
+    hr_img = hr[0:3]
+
+    if lr_img.shape[1] == 121:
+        # add padding
+        lr_img = torch.nn.functional.pad(
+            lr_img[None],
+            pad=(3, 4, 3, 4),
+            mode='reflect'
+        ).squeeze()
+                            
+        # Run the model
+        with torch.no_grad():
+            sr_img = model(lr_img[None]).squeeze()
+
+        # take out padding
+        lr_img = lr_img[:, 3:-4, 3:-4]
+        sr_img = sr_img[:, 3*4:-4*4, 3*4:-4*4]
+    else:
+        # Run the model
+        with torch.no_grad():
+            sr_img = model(lr_img[None]).squeeze()
+
+    # Convert the output to numpy
+    lr_img = (lr_img.cpu().numpy()[0:3] * 10000).astype(np.uint16)
+    sr_img = (sr_img.cpu().numpy()[0:3] * 10000).astype(np.uint16)
+    hr_img = hr_img
+    
+    # Return the results
+    return {
+        "lr": lr_img,
+        "sr": sr_img,
+        "hr": hr_img
+    }

opensr_model/weights/opensr_10m_v4_v5.ckpt

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