Spaces:
Sleeping
Sleeping
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() |