File size: 840 Bytes
720ff4c
 
 
 
 
 
 
 
 
 
aae03c9
720ff4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import torch
from transformers import pipeline
import torch

import gradio as gr

with gr.Blocks() as demo:
    def submit(image):
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        pipe = pipeline("image-classification", model="./checkpoint-600")
        output = pipe(images=[image])

        result = {}

        for index, item in enumerate(output[0]):
            result[item["label"]] = item["score"]

        return result

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Input image", type="filepath")
            
        with gr.Column():
            label = gr.Label()
            submit_button = gr.Button(value="Submit", variant="primary")
            submit_button.click(submit, inputs=[image_input], outputs=label)

demo.launch()