yolov8_runner.py

Go to the documentation of this file.

from ultralytics import YOLO
 
class YoloV8Runner(object):
    """
    YOLO 모델을 동작시키기 위한 클래스
    args:
        checkpoint (str): 모델 파라미터 저장 경로, engine 파일(TensorRT 모델) 또는 pt 파일(torch 모델)
        kwargs (str): 모델 동작을 위한 기타 파라미터
    """
    def __init__(self, checkpoint, **kwargs):
        self.model = YOLO(checkpoint)
        self.kwargs = kwargs
 
    def run_detector(self, image):
        """
        객체 검출기를 동작시키는 기능
        args:
            image (np.ndarray): RGB? BGR? 이미지
        """
        
        results = self.model.predict(source=image, **self.kwargs)
 
        data = {}
        data['det_bboxes'] = results[0].boxes.data.cpu()[:, :5]
        data['det_labels'] = results[0].boxes.cls.cpu()
        
        return data