add inference time, remove global variable for user-specific variable, consistency in idx_slider

app.py

@@ -1,6 +1,7 @@
 import json
 import os
 import random
+import time
 from functools import partial
 from pathlib import Path
 from typing import List
@@ -20,7 +21,6 @@ torch.set_grad_enabled(False)  # stops tracking values for gradients
 
 
 ### Gradio Utils
-
 def generate_imgs_from_user(image,
                             model: EvalModel, baseline: BaselineModel,
                             physics: PhysicsWithGenerator, use_gen: bool,
@@ -60,37 +60,43 @@ def generate_imgs(x: torch.Tensor,
                   physics: PhysicsWithGenerator, use_gen: bool,
                   metrics: List[Metric]):
 
-    with torch.no_grad():
-        ### Compute y
-        y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
-
-        ### Compute x_hat
-        out = model(y=y, physics=physics.physics)
-        out_baseline = baseline(y=y, physics=physics.physics)
-
-        ### Process tensors before metric computation
-        if "Blur" in physics.name:
-            w_1, w_2 = (x.shape[2] - y.shape[2]) // 2, (x.shape[2] + y.shape[2]) // 2
-            h_1, h_2 = (x.shape[3] - y.shape[3]) // 2, (x.shape[3] + y.shape[3]) // 2
-
-            x = x[..., w_1:w_2, h_1:h_2]
-            out = out[..., w_1:w_2, h_1:h_2]
-            if out_baseline.shape != out.shape:
-                out_baseline = out_baseline[..., w_1:w_2, h_1:h_2]
-
-        ### Metrics
-        metrics_y = ""
-        metrics_out = ""
-        metrics_out_baseline = ""
-        for metric in metrics:
-            if y.shape == x.shape:
-                metrics_y += f"{metric.name} = {metric(y, x).item():.4f}" + "\n"
-            metrics_out += f"{metric.name} = {metric(out, x).item():.4f}" + "\n"
-            metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
+    ### Compute y
+    y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
+
+    ### Compute x_hat from RAM & DPIR
+    ram_time = time.time()
+    out = model(y=y, physics=physics.physics)
+    ram_time = time.time() - ram_time
+
+    dpir_time = time.time()
+    out_baseline = baseline(y=y, physics=physics.physics)
+    dpir_time = time.time() - dpir_time
+
+    ### Process tensors before metric computation
+    if "Blur" in physics.name:
+        w_1, w_2 = (x.shape[2] - y.shape[2]) // 2, (x.shape[2] + y.shape[2]) // 2
+        h_1, h_2 = (x.shape[3] - y.shape[3]) // 2, (x.shape[3] + y.shape[3]) // 2
+
+        x = x[..., w_1:w_2, h_1:h_2]
+        out = out[..., w_1:w_2, h_1:h_2]
+        if out_baseline.shape != out.shape:
+            out_baseline = out_baseline[..., w_1:w_2, h_1:h_2]
+
+    ### Metrics
+    metrics_y = ""
+    metrics_out = f"Inference time = {ram_time:.3f}s" + "\n"
+    metrics_out_baseline = f"Inference time = {dpir_time:.3f}s" + "\n"
+    for metric in metrics:
+        if y.shape == x.shape:
+            metrics_y += f"{metric.name} = {metric(y, x).item():.4f}" + "\n"
+        metrics_out += f"{metric.name} = {metric(out, x).item():.4f}" + "\n"
+        metrics_out_baseline += f"{metric.name} = {metric(out_baseline, x).item():.4f}" + "\n"
 
     ### Process y when y shape is different from x shape
-    if physics.name == "MRI" in physics.name:
+    if physics.name == "MRI":
         y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
+    elif physics.name == "CT":
+        y_plot = physics.physics.A_adjoint(y)
     else:
         y_plot = y.clone()
 
@@ -114,28 +120,26 @@ get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
 get_physics_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
 
-AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
-                     'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
 def get_dataset(dataset_name):
-    global AVAILABLE_PHYSICS
     if dataset_name == 'MRI':
-        AVAILABLE_PHYSICS = ['MRI']
-        baseline_name = 'DPIR_MRI'
+        available_physics = ['MRI']
         physics_name = 'MRI'
+        baseline_name = 'DPIR_MRI'
     elif dataset_name == 'CT':
-        AVAILABLE_PHYSICS = ['CT']
-        baseline_name = 'DPIR_CT'
+        available_physics = ['CT']
         physics_name = 'CT'
+        baseline_name = 'DPIR_CT'
     else:
-        AVAILABLE_PHYSICS = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
+        available_physics = ['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
                              'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard']
-        baseline_name = 'DPIR'
         physics_name = 'MotionBlur_easy'
+        baseline_name = 'DPIR'
 
     dataset = get_dataset_on_DEVICE_STR(dataset_name)
+    idx = 0
     physics = get_physics_on_DEVICE_STR(physics_name)
     baseline = get_baseline_model_on_DEVICE_STR(baseline_name)
-    return dataset, physics, baseline
+    return dataset, idx, physics, baseline, available_physics
 
 
 ### Gradio Blocks interface
@@ -144,28 +148,30 @@ title = "Inverse problem playground"  # displayed on gradio tab and in the gradi
 with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
     gr.Markdown("## " + title)
 
-    # DEFAULT VALUES
+    ### DEFAULT VALUES
     # Issue: giving directly a `torch.nn.module` to `gr.State(...)` since it has __call__ method
     # Solution: using lambda expression
     model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
+    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))
+
     dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
     physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
-    idx_placeholder = gr.State(0)
+    available_physics_placeholder = gr.State(['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
+                                              'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard'])
 
-    metric_names = ["PSNR"]
-    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(metric_names))
 
-    @gr.render(inputs=[dataset_placeholder, physics_placeholder])
-    def dynamic_layout(dataset, physics):
-        ### LAYOUT
+    ### LAYOUT
+    # Display images
+    with gr.Row():
+        gt_img = gr.Image(label="Ground-truth IMAGE", interactive=True)
+        observed_img = gr.Image(label="Observed IMAGE", interactive=False)
+        model_a_out = gr.Image(label="RAM OUTPUT", interactive=False)
+        model_b_out = gr.Image(label="DPIR OUTPUT", interactive=False)
 
-        # Display images
-        with gr.Row():
-            gt_img = gr.Image(label=f"Ground-truth IMAGE", interactive=True)
-            observed_img = gr.Image(label=f"Observed IMAGE", interactive=False)
-            model_a_out = gr.Image(label="RAM OUTPUT", interactive=False)
-            model_b_out = gr.Image(label="DPIR OUTPUT", interactive=False)
+    @gr.render(inputs=[dataset_placeholder, physics_placeholder, available_physics_placeholder])
+    def dynamic_layout(dataset, physics, available_physics):
+        ### LAYOUT
 
         # Manage datasets and display metric values
         with gr.Row():
@@ -174,7 +180,7 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
                 choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
                                           label="Datasets",
                                           value=dataset.name)
-                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index")
+                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", key=0)
                 with gr.Row():
                     load_button = gr.Button("Run on index image from dataset")
                     load_random_button = gr.Button("Run on random image from dataset")
@@ -191,14 +197,14 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
         # Manage physics
         with gr.Row():
             with gr.Column(scale=1):
-                choose_physics = gr.Radio(choices=AVAILABLE_PHYSICS,
+                choose_physics = gr.Radio(choices=available_physics,
                                           label="Physics",
                                           value=physics.name)
                 use_generator_button = gr.Checkbox(label="Generate physics parameters during inference")
             with gr.Column(scale=1):
                 with gr.Row():
                     key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()),
-                                                label="Updatable Parameter Key")
+                                               label="Updatable Parameter Key")
                     value_text = gr.Textbox(label="Update Value")
                 update_button = gr.Button("Manually update parameter value")
             with gr.Column(scale=2):
@@ -211,7 +217,7 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
 
         choose_dataset.change(fn=get_dataset,
                               inputs=choose_dataset,
-                              outputs=[dataset_placeholder, physics_placeholder, model_b_placeholder])
+                              outputs=[dataset_placeholder, idx_slider, physics_placeholder, model_b_placeholder, available_physics_placeholder])
         choose_physics.change(fn=get_physics_on_DEVICE_STR,
                               inputs=choose_physics,
                               outputs=[physics_placeholder])