test.py

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)