| 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"]) |