Update opensr_model/run.py

opensr_model/run.py

@@ -1,9 +1,10 @@
+import benchmark
 import opensr_test
 import matplotlib.pyplot as plt
-from utils import create_opensr_model, run_opensr_model
+from opensrmodel.utils import create_opensr_model, run_opensr_model
 
 # Load the model
-model = create_opensr_model(device="cpu")
+model = create_opensr_model(device="cuda")
 
 # Load the dataset
 dataset = opensr_test.load("naip")
@@ -14,7 +15,7 @@ results = run_opensr_model(
     model=model,
     lr=lr_dataset[7],
     hr=hr_dataset[7],
-    device="cpu"
+    device="cuda"
 )
 
 # Display the results
@@ -28,3 +29,9 @@ 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"])