Upload 5 files

satlas/run.py

@@ -0,0 +1,22 @@
+import torch
+import opensr_test
+import matplotlib.pyplot as plt
+from utils import load_satlas_sr, run_satlas
+
+# Load the model
+model = load_satlas_sr(device="cuda")
+
+# Load the dataset
+dataset = opensr_test.load("naip")
+lr_dataset, hr_dataset = dataset["L1C"], dataset["HRharm"]
+
+# Predict a image
+index = 20
+lr = torch.from_numpy(lr_dataset[index][[3, 2, 1]]/3558).float().to("cuda").clamp(0, 1)
+sr = run_satlas(model=model, lr=lr, cropsize=32, overlap=0)
+
+# Run the model
+fig, ax = plt.subplots(1, 2, figsize=(10, 5))
+ax[0].imshow(lr.cpu().numpy().transpose(1, 2, 0))
+ax[1].imshow(sr.cpu().numpy().transpose(1, 2, 0))
+plt.show()

satlas/utils.py

@@ -0,0 +1,76 @@
+import torch
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from typing import Union
+import itertools
+import numpy as np
+
+
+def load_satlas_sr(device: Union[str, torch.device] = "cuda") -> RRDBNet:
+    # Load the weights
+    weights_file = "weights/esrgan_1S2.pth"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Create the model
+    model = RRDBNet(
+        num_in_ch=3,
+        num_out_ch=3,
+        num_feat=64,
+        num_block=23,
+        num_grow_ch=32,
+        scale=4
+    ).to(device)
+
+    # Setup the weights
+    state_dict = torch.load(weights_file)
+    model.load_state_dict(state_dict['params_ema'])
+    model.eval()
+
+    # no gradients
+    for param in model.parameters():
+        param.requires_grad = False
+
+    return model
+
+
+def run_satlas(model, lr, cropsize: int = 32, overlap: int = 0):
+    # Select the raster with the lowest resolution
+    tshp = lr.shape
+    
+    # if the image is too small, return (0, 0)
+    if (tshp[1] < cropsize) and (tshp[2] < cropsize):
+        return [(0, 0)]
+
+    # Define relative coordinates.
+    xmn, xmx, ymn, ymx = (0, tshp[1], 0, tshp[2])
+
+    if overlap > cropsize:
+        raise ValueError("The overlap must be smaller than the cropsize")
+
+    xrange = np.arange(xmn, xmx, (cropsize - overlap))
+    yrange = np.arange(ymn, ymx, (cropsize - overlap))
+
+    # If there is negative values in the range, change them by zero.
+    xrange[xrange < 0] = 0
+    yrange[yrange < 0] = 0
+
+    # Remove the last element if it is outside the tensor
+    xrange = xrange[xrange - (tshp[1] - cropsize) <= 0]
+    yrange = yrange[yrange - (tshp[2] - cropsize) <= 0]
+
+    # If the last element is not (tshp[1] - cropsize) add it!
+    if xrange[-1] != (tshp[1] - cropsize):
+        xrange = np.append(xrange, tshp[1] - cropsize)
+    if yrange[-1] != (tshp[2] - cropsize):
+        yrange = np.append(yrange, tshp[2] - cropsize)
+
+    # Create all the relative coordinates
+    mrs = list(itertools.product(xrange, yrange))
+
+    # Predict the image
+    sr = torch.zeros(3, tshp[1]*4, tshp[2]*4)
+    for x, y in mrs:
+        crop = lr[:, x:x+cropsize, y:y+cropsize]        
+        sr_crop = model(crop[None])[0]
+        sr[:, x*4:(x+cropsize)*4, y*4:(y+cropsize)*4] = sr_crop
+    
+    return sr

satlas/weights/esrgan_1S2.pth

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