benchmark script

benchmark.py

@@ -0,0 +1,56 @@
+import rasterio
+import pathlib
+
+from typing import Callable
+from rasterio.transform import from_origin
+
+
+def create_geotiff(
+    fn: Callable,
+    dataset_snippet: str,
+    output_path: str
+) -> pathlib.Path:
+    """Create all the GeoTIFFs for a specific dataset snippet 
+
+    Args:
+        fn (Callable): A function that return a dictionary with the following keys:
+            - "lr": Low resolution image
+            - "sr": Super resolution image
+            - "hr": High resolution image
+        dataset_snippet (str): The dataset snippet to use to run the fn function.
+        output_path (str): The output path to save the GeoTIFFs.
+
+    Returns:
+        pathlib.Path: The output path where the GeoTIFFs are saved.
+    """
+    pass
+
+
+def run(
+    model_path: str
+) -> pathlib.Path:
+    """Run the all metrics for a specific model.
+
+    Args:
+        model_path (str): The path to the model folder.
+    
+    Returns:
+        pathlib.Path: The output path where the metrics are 
+        saved as a pickle file.
+    """
+    pass
+
+
+def plot(
+    model_path: str
+) -> pathlib.Path:
+    """Generate the plots and tables for a specific model.
+
+    Args:
+        model_path (str): The path to the model folder.
+    
+    Returns:
+        pathlib.Path: The output path where the plots and tables are 
+        saved.
+    """
+    pass

ldm_baseline/metadata.json

@@ -0,0 +1,10 @@
+{
+    "name": "ldm-baseline",
+    "authors": ["CompVis team"],
+    "affiliations": ["None"],
+    "description": "A baseline of LDM models trained on the Open Images dataset.",
+    "code": "open-source",
+    "scale": "x4",
+    "url": "https://huggingface.co/CompVis/ldm-super-resolution-4x-openimages",
+    "license": "apache-2.0"
+  }

ldm_baseline/run.py

@@ -0,0 +1,35 @@
+import matplotlib.pyplot as plt
+import opensr_test
+
+from ldm_baseline.utils import create_stable_diffusion_model, run_diffuser
+
+# set the device
+device = "cuda:0"
+
+# Load the model
+model = create_stable_diffusion_model(device=device)
+
+# Load the dataset
+dataset = opensr_test.load("spain_crops")
+lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]
+
+# Run the model
+results = run_diffuser(model=model, lr=lr_dataset[5], hr=hr_dataset[5], device=device)
+
+# 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()
+
+# Run the experiment
+#
+# benchmark.create_geotiff(run_diffuser, "all", "ldm_baseline/")
+# benchmark.run("all")
+# benchmark.plot("all")

ldm_baseline/utils.py

@@ -0,0 +1,73 @@
+import pickle
+from typing import Union
+
+import numpy as np
+import opensr_test
+import torch
+from diffusers import LDMSuperResolutionPipeline
+
+
+def create_stable_diffusion_model(
+    device: Union[str, torch.device] = "cuda"
+) -> LDMSuperResolutionPipeline:
+    """Create the stable diffusion model
+
+    Returns:
+        LDMSuperResolutionPipeline: The model to use for
+            super resolution.
+    """
+    model_id = "CompVis/ldm-super-resolution-4x-openimages"
+    pipeline = LDMSuperResolutionPipeline.from_pretrained(model_id)
+    pipeline = pipeline.to(device)
+    return pipeline
+
+
+def run_diffuser(
+    model: LDMSuperResolutionPipeline,
+    lr: torch.Tensor,
+    hr: torch.Tensor,
+    device: Union[str, torch.device] = "cuda",
+) -> dict:
+    """Run the model on the low resolution image
+
+    Args:
+        model (LDMSuperResolutionPipeline): The model to use
+        lr (torch.Tensor): The low resolution image
+        hr (torch.Tensor): The high resolution image
+        device (Union[str, torch.device], optional): The device
+            to use. Defaults to "cuda".
+
+    Returns:
+        dict: The results of the model
+    """
+
+    # move the images to the device
+    lr = (torch.from_numpy(lr[[3, 2, 1]]) / 2000).to(device).clamp(0, 1)
+
+    if lr.shape[1] == 121:
+        # add padding
+        lr = torch.nn.functional.pad(
+            lr[None], pad=(3, 4, 3, 4), mode="reflect"
+        ).squeeze()
+
+        # run the model
+        with torch.no_grad():
+            sr = model(lr[None], num_inference_steps=100, eta=1)
+            sr = torch.from_numpy(np.array(sr.images[0]) / 255).permute(2, 0, 1).float()
+
+        # remove padding
+        sr = sr[:, 3 * 4 : -4 * 4, 3 * 4 : -4 * 4]
+        lr = lr[:, 3:-4, 3:-4]
+    else:
+        # run the model
+        with torch.no_grad():
+            sr = model(lr[None], num_inference_steps=100, eta=1)
+            sr = torch.from_numpy(np.array(sr.images[0]) / 255).permute(2, 0, 1).float()
+
+    lr = (lr.cpu().numpy() * 2000).astype(np.uint16)
+    hr = ((hr[0:3] / 2000).clip(0, 1) * 2000).astype(np.uint16)
+    sr = (sr.cpu().numpy() * 2000).astype(np.uint16)
+
+    results = {"lr": lr, "hr": hr, "sr": sr}
+
+    return results