UNIQ-DEV/Image-Classification-Benchmark

1import gradio as gr2import requests3import random4from src.classification_model import ClassificationModel5from src.util.extract import extract_image_urls6 7#only for dummy data8# response = requests.get("https://git.io/JJkYN")9# labels = response.text.split("\n")10 11print('start...')12clf = ClassificationModel()13model_names = clf.get_model_names()14output_labels = []15output_images = []16max_input_image = 1017 18def predict(models, img_url, img_files):19    print(f'model choosen: {models}')20    model_predictions = {}21 22    #set all labels visibility to false23    for label in output_labels:24        model_predictions[label] = gr.Label(label=f'# {name}', visible=False)25    #set all images visibility yo hidden26    for img in output_images:27        model_predictions[img] = gr.Image(visible=False)28    29    sources = extract_image_urls(img_url) + (img_files or [])30    for i, source in enumerate(sources):31        print(f'{i} type: {type(source)} --> {source}')32        if i >= max_input_image: break 33 34        for j, m in enumerate(models):35            results = clf.classify(m, source)36            print(f'{m} --> {results}')37 38            idx = j + (len(model_names)*i) #getting index of label39            label_value = {raw.class_name: raw.confidence for raw in results}        40            model_predictions[output_labels[idx]] =  gr.Label(label=f'# {m}, 3 seconds', value=label_value, visible=True) 41            model_predictions[output_images[i]] = gr.Image(visible=True, value=source, label=f'image {i}') # set image visibility to true42    43    return model_predictions44 45with gr.Blocks() as demo:46    gr.Markdown("# Image Classification Benchmark")47    gr.Markdown("You can input at maximum 10 images at once (urls or files)")48    49    with gr.Row():50        with gr.Column(scale=1):51            model = gr.Dropdown(choices=model_names, multiselect=True, label='Choose the model')52            img_urls = gr.Textbox(label='Image Urls (separated with comma)')    53            img_files = gr.File(label='Upload Files',file_count='multiple', file_types=['image'])54            apply = gr.Button("Classify", variant='primary')55        with gr.Column(scale=1):56            for i in range(max_input_image):57                output_images.append(gr.Image(interactive=False, visible= (i==0)))58                for name in clf.get_model_names():59                    output_labels.append(gr.Label(label=f'# {name}', visible= (i==0)))                  60 61    apply.click(fn=predict,62               inputs=[model, img_urls, img_files],63               outputs=output_images+output_labels)64 65 66# demo.launch()67demo.queue().launch()