unfinished work

app.py

@@ -74,18 +74,18 @@ def generate_imgs(dataset: EvalDataset, idx: int,
 
     return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
-def update_random_idx_and_generate_imgs(dataset: EvalDataset, 
-                                        model: EvalModel, 
-                                        baseline: BaselineModel, 
+def update_random_idx_and_generate_imgs(dataset: EvalDataset,
+                                        model: EvalModel,
+                                        baseline: BaselineModel,
                                         physics: PhysicsWithGenerator,
                                         use_gen: bool,
                                         metrics: List[Metric]):
     idx = random.randint(0, len(dataset)-1)
-    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs(dataset, 
-                                                                                                            idx, 
-                                                                                                            model, 
-                                                                                                            baseline, 
-                                                                                                            physics, 
+    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs(dataset,
+                                                                                                            idx,
+                                                                                                            model,
+                                                                                                            baseline,
+                                                                                                            physics,
                                                                                                             use_gen, 
                                                                                                             metrics)
     return idx, x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline
@@ -145,6 +145,17 @@ def get_model(model_name, ckpt_pth):
     else:
         return get_eval_model_on_DEVICE_STR(model_name, ckpt_pth)
 
+AVAILABLE_PHYSICS = PhysicsWithGenerator.all_physics
+
+def get_dataset(dataset_name):
+    global AVAILABLE_PHYSICS
+    if dataset_name = 'MRI':
+        AVAILABLE_PHYSICS = ['MRI']
+    elif dataset_name = 'CT':
+        AVAILABLE_PHYSICS = ['CT']
+    else:
+        AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard', 'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
+    return get_dataset_on_DEVICE_STR(dataset_name)
 
 ### Gradio Blocks interface