Spaces:
Sleeping
Sleeping
File size: 785 Bytes
ab899b5 efae727 0f0204b ab899b5 efae727 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
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
|