Update demo.py

demo.py

@@ -80,7 +80,8 @@ def experiment_to_dict(experiment, skip=False):
     return info
     
 def generate_new_experiment() -> Experiment:
-    corruption = random.choice([f for f in list(Path(f"./images/{DATASET}").glob("*/*")) if f.is_dir()])
+    wanted_corruptions = ["spatter", "impulse_noise", "speckle_noise", "gaussian_noise", "pixelate", "jpeg_compression"]
+    corruption = random.choice([f for f in list(Path(f"./images/{DATASET}").glob("*/*")) if f.is_dir() and f.name in wanted_corruptions)
     image_id = random.choice(list(corruption.glob("*")))
     imgs_to_sample = (NUMBER_OF_IMAGES_PER_ROW * NUMBER_OF_ROWS) // 2