File size: 1,008 Bytes
caa7010 61e0235 caa7010 61e0235 caa7010 61e0235 caa7010 b01e57b 61e0235 b01e57b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import matplotlib.pyplot as plt
import opensr_test
import benchmark
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("naip", force=False)
lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"]
# Run the model
index = 5
results = run_diffuser(
model=model,
lr=lr_dataset[index],
hr=hr_dataset[index],
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(model, run_diffuser, "all", "ldm_baseline/")
# benchmark.run(["naip"])
# benchmark.plot(["naip"]) |