app.py

import json
import os
import random
from functools import partial
from pathlib import Path
from typing import List

import deepinv as dinv
import gradio as gr
import torch
from PIL import Image
from torchvision import transforms

from evals import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDataset, Metric


DEVICE_STR = 'cuda'


### Gradio Utils
def generate_imgs(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)

    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"

    ### Process y when y shape is different from x shape
    if physics.name == "MRI" or "SR" in physics.name:
        y_plot = physics.physics.prox_l2(physics.physics.A_adjoint(y), y, 1e4)
    else:
        y_plot = y.clone()

    ### Processing images for plotting :
    #     - clip value outside of [0,1]
    #     - shape (1, C, H, W) -> (C, H, W)
    #     - torch.Tensor object -> Pil object
    process_img = partial(dinv.utils.plotting.preprocess_img, rescale_mode="clip")
    to_pil = transforms.ToPILImage()
    x = to_pil(process_img(x)[0].to('cpu'))
    y = to_pil(process_img(y_plot)[0].to('cpu'))
    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, 
                                        model: EvalModel, 
                                        baseline: BaselineModel, 
                                        physics: PhysicsWithGenerator,
                                        use_gen: bool,
                                        metrics: List[Metric]):
    idx = random.randint(0, len(dataset))
    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

def save_imgs(dataset: EvalDataset, idx: int, physics: PhysicsWithGenerator,
              model_a: EvalModel | BaselineModel, model_b: EvalModel | BaselineModel,
              x: Image.Image, y: Image.Image,
              out_a: Image.Image, out_b: Image.Image,
              y_metrics_str: str,
              out_a_metric_str : str, out_b_metric_str: str) -> None:

    ### PROCESSES STR
    physics_params_str = ""
    for param_name, param_value in physics.saved_params["updatable_params"].items():
        physics_params_str += f"{param_name}_{param_value}-"
    physics_params_str = physics_params_str[:-1] if physics_params_str.endswith("-") else physics_params_str
    y_metrics_str = y_metrics_str.replace(" = ", "_").replace("\n", "-")
    y_metrics_str = y_metrics_str[:-1] if y_metrics_str.endswith("-") else y_metrics_str
    out_a_metric_str = out_a_metric_str.replace(" = ", "_").replace("\n", "-")
    out_a_metric_str = out_a_metric_str[:-1] if out_a_metric_str.endswith("-") else out_a_metric_str
    out_b_metric_str = out_b_metric_str.replace(" = ", "_").replace("\n", "-")
    out_b_metric_str = out_b_metric_str[:-1] if out_b_metric_str.endswith("-") else out_b_metric_str

    save_path = SAVE_IMG_DIR / f"{dataset.name}+{idx}+{physics.name}+{physics_params_str}+{y_metrics_str}+{model_a.name}+{out_a_metric_str}+{model_b.name}+{out_b_metric_str}.png"
    titles = [f"{dataset.name}[{idx}]",
              f"y = {physics.name}(x)",
              f"{model_a.name}",
              f"{model_b.name}"]

    # Pil object -> torch.Tensor
    to_tensor = transforms.ToTensor()
    x = to_tensor(x)
    y = to_tensor(y)
    out_a = to_tensor(out_a)
    out_b = to_tensor(out_b)

    dinv.utils.plot([x, y, out_a, out_b], titles=titles, show=False, save_fn=save_path)

get_list_metrics_on_DEVICE_STR = partial(Metric.get_list_metrics, device_str=DEVICE_STR)
get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
get_eval_model_on_DEVICE_STR = partial(EvalModel, device_str=DEVICE_STR)
get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
get_physics_generator_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)

def get_model(model_name, ckpt_pth):
    if model_name in BaselineModel.all_baselines:
        return get_baseline_model_on_DEVICE_STR(model_name)
    else:
        return get_eval_model_on_DEVICE_STR(model_name, ckpt_pth)


### Gradio Blocks interface

# Define custom CSS
custom_css = """
.fixed-textbox textarea {
    height: 90px !important;  /* Adjust height to fit exactly 4 lines */
    overflow: scroll;         /* Add a scroll bar if necessary */
    resize: none;             /* User can resize vertically the textbox */
}
"""

title = "Inverse problem playground"  # displayed on gradio tab and in the gradio page
with gr.Blocks(title=title, css=custom_css) as interface:
    gr.Markdown("## " + title)

    # Loading things
    model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))  # lambda expression to instanciate a callable in a gr.State
    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DRUNET"))  # lambda expression to instanciate a callable in a gr.State
    dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
    physics_placeholder = gr.State(lambda: get_physics_generator_on_DEVICE_STR("MotionBlur_easy"))  # lambda expression to instanciate a callable in a gr.State
    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))

    @gr.render(inputs=[model_a_placeholder, model_b_placeholder, dataset_placeholder, physics_placeholder, metrics_placeholder])
    def dynamic_layout(model_a, model_b, dataset, physics, metrics):
        ### LAYOUT
        model_a_name = model_a.base_name
        model_a_full_name = model_a.name
        model_b_name = model_b.base_name
        model_b_full_name = model_b.name
        dataset_name = dataset.name
        physics_name = physics.name
        metric_names = [metric.name for metric in metrics]

        # Components: Inputs/Outputs + Load EvalDataset/PhysicsWithGenerator/EvalModel/BaselineModel
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    with gr.Column():
                        clean = gr.Image(label=f"{dataset_name} IMAGE", interactive=False)
                        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)
                        y_metrics = gr.Textbox(label="Metrics(y, x)", elem_classes=["fixed-textbox"],)
                with gr.Row():
                    choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
                                              label="List of EvalDataset",
                                              value=dataset_name,
                                              scale=2)
                    idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", scale=1)
        
                choose_physics = gr.Radio(choices=PhysicsWithGenerator.all_physics,
                                          label="List of PhysicsWithGenerator",
                                          value=physics_name)
                with gr.Row():
                    key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()), 
                                               label="Updatable Parameter Key",
                                               scale=2)
                    value_text = gr.Textbox(label="Update Value", scale=2)
                    with gr.Column(scale=1):
                        update_button = gr.Button("Update Param")
                        use_generator_button = gr.Checkbox(label="Use param generator")
                    
            with gr.Column():
                with gr.Row():      
                    with gr.Column():
                        model_a_out = gr.Image(label=f"{model_a_full_name} OUTPUT", interactive=False)
                        out_a_metric = gr.Textbox(label="Metrics(model_a(y), x)", elem_classes=["fixed-textbox"])
                        load_model_a = gr.Button("Load model A...", scale=1)
                    with gr.Column():
                        model_b_out = gr.Image(label=f"{model_b_full_name} OUTPUT", interactive=False)
                        out_b_metric = gr.Textbox(label="Metrics(model_b(y), x)", elem_classes=["fixed-textbox"])
                        load_model_b = gr.Button("Load model B...", scale=1)
                with gr.Row():
                    choose_model_a = gr.Dropdown(choices=EvalModel.all_models + BaselineModel.all_baselines,
                                                 label="List of Model A",
                                                 value=model_a_name,
                                                 scale=2)
                    path_a_str = gr.Textbox(value=model_a.ckpt_pth, label="Checkpoint path", scale=3)
                with gr.Row():
                    choose_model_b = gr.Dropdown(choices=EvalModel.all_models + BaselineModel.all_baselines,
                                                 label="List of Model B",
                                                 value=model_b_name,
                                                 scale=2)
                    path_b_str = gr.Textbox(value=model_b.ckpt_pth, label="Checkpoint path", scale=3)
        
        # Components: Load Metric + Load/Save Buttons
        with gr.Row():
            with gr.Column():
                choose_metrics = gr.CheckboxGroup(choices=Metric.all_metrics,
                                                  value=metric_names,
                                                  label="Choose metrics you are interested")
            with gr.Column():
                load_button = gr.Button("Load images...")
                load_random_button = gr.Button("Load randomly...")
        save_button = gr.Button("Save images...")

        ### Event listeners
        choose_dataset.change(fn=get_dataset_on_DEVICE_STR,
                              inputs=choose_dataset,
                              outputs=dataset_placeholder)
        choose_physics.change(fn=get_physics_generator_on_DEVICE_STR,
                              inputs=choose_physics,
                              outputs=physics_placeholder)
        update_button.click(fn=physics.update_and_display_params, inputs=[key_selector, value_text], outputs=physics_params)
        load_model_a.click(fn=get_model,
                           inputs=[choose_model_a, path_a_str],
                           outputs=model_a_placeholder)
        load_model_b.click(fn=get_model,
                           inputs=[choose_model_b, path_b_str],
                           outputs=model_b_placeholder)
        choose_metrics.change(fn=get_list_metrics_on_DEVICE_STR,
                              inputs=choose_metrics,
                              outputs=metrics_placeholder)
        load_button.click(fn=generate_imgs, 
                          inputs=[dataset_placeholder, 
                                  idx_slider, 
                                  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_random_button.click(fn=update_random_idx_and_generate_imgs, 
                                 inputs=[dataset_placeholder, 
                                         model_a_placeholder,
                                         model_b_placeholder,
                                         physics_placeholder,
                                         use_generator_button,
                                         metrics_placeholder], 
                                 outputs=[idx_slider, clean, y_image, model_a_out, model_b_out, physics_params, y_metrics, out_a_metric, out_b_metric])

interface.launch()