model.py

#  Here will be the inference to my model on hugging face!!
# Use a pipeline as a high-level helper

from timeit import default_timer as timer
from typing import Tuple, Dict
from transformers import pipeline

pipe = pipeline("image-classification", model="JYL480/vit-base-images")

image_path = "examples/melanocytic_Nevi.jpg"


def predict(image):
    start = timer()
    result = pipe(image)
    print(result)
    pred_time = round(timer() - start, 5)
    combined_dict = {item['label']: float(item['score']) for item in result}
    return combined_dict, pred_time

# combined_dict, pred_time = predict(image_path)
# print(combined_dict)