"""
# File       : yolov5hub.py
# Time       ：22.5.16 16:47
# Author     ：FEANGYANG
# version    ：python 3.7
# Contact    ：1071082183@qq.com
# Description：
"""
import cv2
import torch
import glob


class Yolov5Detect:
    def __init__(self, weight_path, weight, class_f='\W+'):
        self.class_f = class_f
        self.model = torch.hub.load(weight_path, weight, source='local')


    def yolov5_detect(self, img):
        results = self.model(img)  # includes NMS
        all_results = results.pandas().xyxy[0]
        filter_cla = self._filter_class(all_results, self.class_f)
        # print(filter_cla)
        return filter_cla

    def _filter_class(self, all_results , class_f):
        filter_cla = all_results[all_results['name'].str.contains(class_f)]
        return filter_cla


if __name__ == '__main__':
    weight_path = './'
    weight = 'yolov5l'

    detect = Yolov5Detect(weight_path, weight, class_f='person')

    for img_path in glob.glob('./datasets/coco128/images/train20/*.jpg'):
        img = cv2.imread(img_path)[:,:,::-1]
        # img = cv2.imread(img_path)
        # img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        result = detect.yolov5_detect(img)
        print(result)
        print('-----------------------------------------------------')