用python写yolov5 拍球计数

import numpy as np
import torch
import cv2
from PIL import Image, ImageDraw, ImageFont

def cv2_add_chinese_text(img, text, position, text_color=(0, 255, 0), text_size=30):
    img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img)
    font_style = ImageFont.truetype("./fonts/songti.ttc", text_size, encoding="utf-8")
    # 绘制文本
    text_color = (text_color[2], text_color[1], text_color[0])
    draw.text(position, text, text_color, font=font_style)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

class FireDetection:
    def __init__(self):
        self.detector = torch.hub.load('G:\\yolov\\yolov5-5.0', 'custom', 'G:\\yolov\\yolov5-5.0\\weights\\yolov5s.pt', source='local')
        self.detector.conf = 0.1

    def detect(self):
        cap = cv2.VideoCapture('./basketball.mp4')

        ballY = 0
        ballY_Old = 0
        down = False
        up = False
        count = 0


        while True:
            ret, frame = cap.read()
            if not ret or frame is None:
                break

            # 转为RGB
            img_cvt = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            results = self.detector(img_cvt)
            pd = results.pandas().xyxy[0]

            # 绘制检测框
            for obj in pd.to_numpy():
                box_l, box_t = int(obj[0]), int(obj[1])
                box_r, box_b = int(obj[2]), int(obj[3])
                confidence  = obj[4]
                name = obj[5]
                obj_name = obj[6]

                if obj_name == 'sports ball':
                    box_color = (0, 255, 0)
                    box_txt = '检测球'
                    ballY = box_b
                else:
                    continue

                # if confidence < 0.6:
                #     continue

                if (ballY > ballY_Old):
                    down = True
                elif (ballY < ballY_Old) and down:
                    up = True

                ballY_Old = ballY

                if True == down and True == up:
                    count = count + 1
                    down = False
                    up = False

                frame = cv2_add_chinese_text(frame, "个数：" + str(count) +  " 置信度：" + str(confidence), (5, 50), box_color, 35)
                frame = cv2.rectangle(frame, (box_l, box_t), (box_r, box_b), box_color, 2)
                frame = cv2_add_chinese_text(frame, box_txt, (box_l, box_t-40), box_color, 25)

            cv2.imshow('basketball detection', frame)

            if cv2.waitKey(5) & 0xFF == ord('q'):
                break

if __name__ == '__main__':
    fire_detection = FireDetection()
    fire_detection.detect()