import torch
import numpy
from towhee.operator import NNOperator

@register(outputschema=['boxes', 'classes', 'scores'])
class Yolov5(NNOperator):
    def __init__(self, model_name):
        super().__init__()
        self._model = torch.hub.load("ultralytics/yolov5", model_name, pretrained=True)

    def __call__(self, img: numpy.ndarray):
        # Get object detection results with YOLOv5 model
        results = self._model(img)
        
        boxes = results.xyxy[0]
        classes = list(results.pandas().xyxy[0].name)
        scores = list(results.pandas().xyxy[0].confidence)
        return boxes, classes, scores