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()
用python写yolov5 拍球计数
最新推荐文章于 2024-07-07 21:26:44 发布