import torch
import numpy
import logging

from towhee import register
from towhee.operator import NNOperator

logging.getLogger("yolov5").setLevel(logging.WARNING)

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

    def __call__(self, img: numpy.ndarray):
        # Get object detection results with YOLOv5 model
        results = self._model(img)

        boxes = [re[0:4] for re in results.xyxy[0]]
        boxes = [list(map(int, box)) for box in boxes]
        classes = list(results.pandas().xyxy[0].name)
        scores = list(results.pandas().xyxy[0].confidence)
        return boxes, classes, scores