Upload 3 files

satlas/run.py

@@ -11,12 +11,23 @@ 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)
+results = run_satlas(
+    model=model,
+    lr=lr_dataset[4],
+    hr=hr_dataset[4],
+    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))
+# Display the results
+fig, ax = plt.subplots(1, 3, figsize=(10, 5))
+ax[0].imshow(results["lr"].transpose(1, 2, 0)/10000)
+ax[0].set_title("LR")
+ax[0].axis("off")
+ax[1].imshow(results["sr"].transpose(1, 2, 0)/10000)
+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()
+

satlas/utils.py

@@ -32,7 +32,17 @@ def load_satlas_sr(device: Union[str, torch.device] = "cuda") -> RRDBNet:
     return model
 
 
-def run_satlas(model, lr, cropsize: int = 32, overlap: int = 0):
+def run_satlas(
+    model: RRDBNet,
+    lr: torch.Tensor,
+    hr: torch.Tensor,
+    cropsize: int = 32,
+    overlap: int = 0,
+    device: Union[str, torch.device] = "cuda"
+) -> torch.Tensor:
+    # Load the LR image
+    lr = torch.from_numpy(lr[[3, 2, 1]]/3558).float().to(device).clamp(0, 1)
+
     # Select the raster with the lowest resolution
     tshp = lr.shape
     
@@ -73,4 +83,11 @@ def run_satlas(model, lr, cropsize: int = 32, overlap: int = 0):
         sr_crop = model(crop[None])[0]
         sr[:, x*4:(x+cropsize)*4, y*4:(y+cropsize)*4] = sr_crop
     
-    return sr
+    # Save the result
+    results = {
+        "lr": (lr.cpu().numpy() * 10000).astype(np.uint16),
+        "sr": (sr.cpu().numpy() * 10000).astype(np.uint16),
+        "hr": hr[0:3]
+    }
+
+    return results