GazeLocation/Visualization_Saliency

1import gradio as gr2from datasets import load_dataset3 4 5# +6def get_methods_and_arch(dataset):7    columns = dataset.column_names[5:]8    methods = []9    archs   = []10    for column in columns:11        methods.append(column.split('_')[0])12        archs.append('_'.join(column.split('_')[1:-2]))13    return list(set(methods)),list(set(archs))14 15def get_columns(arch,method):16    columns = dataset.column_names[5:]17    for col in columns:18        if f'{method}_{arch}' in col:19            return col20def button_fn(arch,method):21    column_heatmap = get_columns(arch,method)22    #print("Updated column: ",column_heatmap)23    return column_heatmap,index_default,dataset[index_default]["image"],dataset[index_default][column_heatmap]24 25def func_slider(index,column_textbox):26    #global column_heatmap27    example = dataset[index]28    return example['image'],example[column_textbox]29 30 31# -32 33dataset = load_dataset("GazeLocation/stimuli_heatmaps",split = 'train')34METHODS, ARCHS = get_methods_and_arch(dataset)35index_default = 036DEMO = False37 38if __name__ == '__main__':39    demo = gr.Blocks()40    with demo:41        gr.Markdown("# Heatmap Gaze Location")            42       43        with gr.Row():44            dropdown_arch   = gr.Dropdown(choices = ARCHS,45                                   value   =  'resnet50',46                                   label   = 'Model')47 48            dropdown_method = gr.Dropdown(choices = METHODS,49                                   value   =  'gradcam',50                                   label   = 'Method')51        with gr.Row():52            button = gr.Button(label = 'Update Heatmap Model - Method')53                    54        with gr.Row():55            hf_slider = gr.Slider(minimum=0, maximum=len(dataset)-1,step = 1)56        with gr.Row():57            column_textbox = gr.Textbox(label = 'column name',58                                        value = get_columns(ARCHS[0],METHODS[0]) )59        with gr.Row():60            with gr.Column():61                image_input = gr.Image(label="Input Image",value = dataset[index_default]["image"])62            with gr.Column():63                image_output = gr.Image(label="Output",value = dataset[index_default][get_columns('resnet50','gradcam')])64 65    66        button.click(fn = button_fn,67                     inputs = [dropdown_arch,dropdown_method],68                     outputs = [column_textbox,hf_slider,image_input,image_output])69        70        71        hf_slider.change(func_slider, 72                         inputs  = [hf_slider,column_textbox],73                         outputs = [image_input, image_output])74    if DEMO:75        demo.launch(share = True,debug = True)76    else:77        demo.launch()78 79 80 81 82