[Feat] Apply methods on user image

app.py

@@ -18,14 +18,31 @@ DEVICE_STR = 'cuda'
 
 
 ### Gradio Utils
-def generate_imgs(dataset: EvalDataset, idx: int,
-                  model: EvalModel, baseline: BaselineModel,
-                  physics: PhysicsWithGenerator, use_gen: bool,
-                  metrics: List[Metric]):
+
+def generate_imgs_from_dataset(dataset: EvalDataset, idx: int,
+                               model: EvalModel, baseline: BaselineModel,
+                               physics: PhysicsWithGenerator, use_gen: bool,
+                               metrics: List[Metric]):
     ### Load 1 image
     x = dataset[idx]    # shape : (3, 256, 256)
     x = x.unsqueeze(0)  # shape : (1, 3, 256, 256)
 
+    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+
+def generate_imgs_from_user(image,
+                            model: EvalModel, baseline: BaselineModel,
+                            physics: PhysicsWithGenerator, use_gen: bool,
+                            metrics: List[Metric]):
+    # PIL image -> torch.Tensor
+    x = transforms.ToTensor()(image).unsqueeze(0)
+
+    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+    
+def generate_imgs(x: torch.Tensor,
+                  model: EvalModel, baseline: BaselineModel,
+                  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
@@ -71,23 +88,18 @@ def generate_imgs(dataset: EvalDataset, idx: int,
     out = to_pil(process_img(out)[0].to('cpu'))
     out_baseline = to_pil(process_img(out_baseline)[0].to('cpu'))
 
-
     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,
+def generate_random_imgs_from_dataset(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,
-                                                                                                            use_gen, 
-                                                                                                            metrics)
+    x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs_from_dataset(
+        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
 
 
@@ -151,7 +163,7 @@ with gr.Blocks(title=title, css=custom_css) as interface:
             with gr.Column():
                 with gr.Row():
                     with gr.Column():
-                        clean = gr.Image(label=f"{dataset_name} IMAGE", interactive=False)
+                        clean = gr.Image(label=f"{dataset_name} IMAGE", interactive=True)
                         physics_params = gr.Textbox(label="Physics parameters", elem_classes=["fixed-textbox"], value=physics.display_saved_params())
                     with gr.Column():
                         y_image = gr.Image(label=f"{physics_name} IMAGE", interactive=False)
@@ -189,9 +201,10 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                                                   value=metric_names,
                                                   label="Choose metrics you are interested")
             use_generator_button = gr.Checkbox(label="Generate valid physics parameters", scale=1)
+            run_button = gr.Button("Run current image")
             with gr.Column(scale=1):
-                load_button = gr.Button("Load images...")
-                load_random_button = gr.Button("Load randomly...")
+                load_button = gr.Button("Load images from dataset...")
+                load_random_button = gr.Button("Load randomly from dataset...")
 
         ### Event listeners
         choose_dataset.change(fn=get_dataset,
@@ -204,7 +217,15 @@ with gr.Blocks(title=title, css=custom_css) as interface:
         choose_metrics.change(fn=get_list_metrics_on_DEVICE_STR,
                               inputs=choose_metrics,
                               outputs=metrics_placeholder)
-        load_button.click(fn=generate_imgs,
+        run_button.click(fn=generate_imgs_from_user,
+                         inputs=[clean,
+                                 model_a_placeholder,
+                                 model_b_placeholder,
+                                 physics_placeholder,
+                                 use_generator_button,
+                                 metrics_placeholder],
+                         outputs=[clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
+        load_button.click(fn=generate_imgs_from_dataset,
                           inputs=[dataset_placeholder,
                                   idx_slider,
                                   model_a_placeholder,
@@ -213,7 +234,7 @@ with gr.Blocks(title=title, css=custom_css) as interface:
                                   use_generator_button,
                                   metrics_placeholder],
                           outputs=[clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])
-        load_random_button.click(fn=update_random_idx_and_generate_imgs,
+        load_random_button.click(fn=generate_random_imgs_from_dataset,
                                  inputs=[dataset_placeholder,
                                          model_a_placeholder,
                                          model_b_placeholder,