| import benchmark | |
| import opensr_test | |
| import matplotlib.pyplot as plt | |
| from opensrmodel.utils import create_opensr_model, run_opensr_model | |
| # Load the model | |
| model = create_opensr_model(device="cuda") | |
| # Load the dataset | |
| dataset = opensr_test.load("naip") | |
| lr_dataset, hr_dataset = dataset["L2A"], dataset["HRharm"] | |
| # Run the model | |
| results = run_opensr_model( | |
| model=model, | |
| lr=lr_dataset[7], | |
| hr=hr_dataset[7], | |
| device="cuda" | |
| ) | |
| # 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_opensr_model, "all", "opensrmodel/", device="cuda") | |
| # benchmark.run(["naip"]) | |
| # benchmark.plot(["naip"]) |