gtsrb / test.py
Bazyl
use pipeline
b124b3b
from transformers import ViTImageProcessor, ViTForImageClassification, AutoModelForImageClassification, AutoTokenizer, pipeline
from PIL import Image
import requests
import os
import gradio as gr
from timeit import default_timer as timer
from typing import Tuple, Dict
import torch
start_time = timer()
# config = AutoConfig.from_pretrained('bazyl/gtsrb-model')
image = Image.open('examples/00009.png')
classifier = pipeline("image-classification", model="bazyl/gtsrb-model", tokenizer="bazyl/gtsrb-model")
result = classifier(image)
response = {result[i]["label"]: result[i]["score"] for i in range(len(result))}
# Calculate the prediction time
pred_time = round(timer() - start_time, 5)
print(classifier(image), pred_time)