import spaces
import numpy as np
import yolov5
from yolov5.utils.plots import Annotator, colors
from huggingface_hub import get_token


def load_model(model_path, img_size=640):
    """Load model from HuggingFace Hub."""
    model = yolov5.load(model_path, hf_token=get_token())
    model.img_size = img_size  # add img_size attribute
    return model


@spaces.GPU
def inference(model, image):
    """Run inference on image and return annotated image."""
    results = model(image, size=model.img_size)
    annotator = Annotator(np.asarray(image))
    for *box, _, cls in reversed(results.pred[0]):
        # label = f'{model.names[int(cls)]} {conf:.2f}'
        # print(f'{cls} {conf:.2f} {box}')
        annotator.box_label(box, "", color=colors(cls, True))
    return annotator.im