Upload mass_generate_examples.py

mass_generate_examples.py

@@ -43,9 +43,9 @@ if __name__ == "__main__":
     test_ridge = torch.from_numpy(imageio.imread('dataset_large/Ridge_11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
     test_basin = torch.from_numpy(imageio.imread('dataset_large/Basins_11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
     gt = torch.from_numpy(imageio.imread('dataset_large/11417648.tiff'))[None, None, :].to(dtype=torch.float16, device='cuda')
-    water_level = 300.0
-    num_steps = 10
-    num_images = 4
+    water_level = 0.0
+    num_steps = 50
+    num_images = 16
 
     test_basin = (test_basin >= water_level).to(torch.float16)
     test_ridge = test_ridge.expand(num_images, -1, -1, -1)
@@ -100,7 +100,7 @@ if __name__ == "__main__":
 
     # Plot condition image
     im1 = axes[1].imshow(basin_display, cmap='gray')
-    axes[1].set_title('Basin Condition', fontsize=12, pad=2)
+    axes[1].set_title(f'Basin Condition at level {water_level}', fontsize=12, pad=2)
     axes[1].set_axis_off()
 
     # Plot ground truth image