Upload 2 files

superimage/run.py

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

superimage/utils.py

@@ -0,0 +1,50 @@
+import torch
+import numpy as np
+from super_image import HanModel
+from typing import Union
+
+def create_superimage_model(
+    device: Union[str, torch.device] = "cuda"
+) -> HanModel:
+    """ Create the super image model
+
+    Returns:
+        HanModel: The super image model
+    """
+    return HanModel.from_pretrained('eugenesiow/han', scale=4).to(device)
+
+
+def run_superimage(
+    model: HanModel,
+    lr: np.ndarray,
+    hr: np.ndarray,
+    device: Union[str, torch.device] = "cuda"    
+):
+    """ Run the super image model
+
+    Args:
+        model (HanModel): The super image model
+        lr (np.ndarray): The low resolution image
+        hr (np.ndarray): The high resolution image
+        device (Union[str, torch.device], optional): The device to run the model on. Defaults to "cuda".
+
+    Returns:
+        dict: The results
+    """
+    # Convert the images to tensors
+    lr_tensor = (torch.from_numpy(lr[[3, 2, 1]]).to(device) / 2000).float()
+    
+    # Run the model
+    with torch.no_grad():
+        sr_tensor = model(lr_tensor[None])
+
+    # Convert the tensors to numpy arrays
+    lr = (lr_tensor.cpu().numpy() * 2000).astype(np.uint16)
+    sr = (sr_tensor.cpu().numpy() * 2000).astype(np.uint16)
+
+    # Return the results
+    return {
+        "lr": lr.squeeze(),
+        "hr": hr[0:3].squeeze(),
+        "sr": sr.squeeze()
+    }