187 lines
7.7 KiB
Python
187 lines
7.7 KiB
Python
# 导入相关依赖
|
||
import os
|
||
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
|
||
import cv2
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
import torch
|
||
import argparse
|
||
|
||
def exec_inference_video(file_path,line):
|
||
if torch.cuda.is_available():
|
||
model = YOLO("weights/yolov8s.pt") # gpu
|
||
print('gpu推理')
|
||
else:
|
||
model = YOLO('weights/yolov8s.onnx')
|
||
print('cpu推理')
|
||
# 基本逻辑:行人从视频顶部往下部行走并进触碰线,视为 下行 low;从视频画面下部往上部行走并进触碰线,视为 上行 up。
|
||
# 其实我们只在视频中画了一条线，实际上在程序判断的时候，线的上下都有一段区域(四边形,以我们之前画的线为一边，高度20,宽度是线的宽度)
|
||
# 如果行人先出现在线的上边区域再出现线的下边区域,证明是从上往下,反之,从下往上
|
||
## 加载要使用的视频
|
||
kamera = cv2.VideoCapture(file_path)
|
||
font = cv2.FONT_HERSHEY_DUPLEX
|
||
# 获取视频的高度与宽度
|
||
width, height = kamera.get(3), kamera.get(4)
|
||
fps = kamera.get(cv2.CAP_PROP_FPS)
|
||
# 注意opencv的坐标系是从左上方开始的,向下为y轴,向右为x轴
|
||
end_x = int(width)
|
||
# 注意线的位置应该根据视频而定,不是一成不变的
|
||
|
||
l_x1,l_y1,l_x2,l_y2=line[0][0],line[0][1],line[1][0],line[1][1]
|
||
#防止线的位置太靠边缘而使线上下区域超出画面
|
||
if l_x1-5<0:
|
||
l_x1=6
|
||
if l_x2-5<0:
|
||
l_x2=6
|
||
if l_x1+5>width:
|
||
l_x1=width-6
|
||
if l_x2+5>width:
|
||
l_x2=width-6
|
||
if l_y1-20<0:
|
||
l_y1=21
|
||
if l_y1+20>height:
|
||
l_y1=height-21
|
||
if l_y2-20<0:
|
||
l_y2=21
|
||
if l_y2+20>height:
|
||
l_y2=height-21
|
||
l_x1,l_y1,l_x2,l_y2=int(l_x1),int(l_y1),int(l_x2),int(l_y2)
|
||
#线的上面区域
|
||
region_1=[(l_x1,l_y1-50),(l_x2,l_y2-50),(l_x2,l_y2),(l_x1,l_y1)]
|
||
#线的下面的区域
|
||
region_2=[(l_x1,l_y1),(l_x2,l_y2),(l_x2,l_y2+50),(l_x1,l_y1+50)]
|
||
region1 = np.array(region_1)
|
||
region1 = region1.reshape((-1, 1, 2))
|
||
|
||
region2 = np.array(region_2)
|
||
region2 = region2.reshape((-1, 1, 2))
|
||
|
||
# 下行的行人id保存在这个变量中
|
||
low_id = set()
|
||
|
||
# 上行行人的id保存在这个变量中
|
||
up_id = set()
|
||
#进入上面(region1)区域的行人id保存在这个变量中
|
||
total_region1=set()
|
||
# 进入下面(region2)区域的行人id保存到这里
|
||
total_region2= set()
|
||
|
||
# output_path = file_path.split(".")[0]+ "_output.mp4"
|
||
# print(output_path)
|
||
# fourcc = cv2.VideoWriter_fourcc(*'MJPG')
|
||
# out = cv2.VideoWriter(output_path, fourcc, 30, (int(width), int(height)))
|
||
|
||
up_cnt = 0
|
||
down_cnt = 0
|
||
|
||
while True:
|
||
# 逐渐帧读取视频内容
|
||
ret, frame = kamera.read()
|
||
if not ret:
|
||
break
|
||
# opnecv读取的图片色彩空间为BGR,我们把它转化成RGB
|
||
rgb_img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
cv2.line(frame, (l_x1,l_y1),(l_x2,l_y2), (255, 0, 255), 3)
|
||
# 使用yolo8的追踪模式(采用byte-tracker多目标跟踪算法)
|
||
# 遍历yolo的处理结果，结果全部保存在返回的results(list类型)变量中,里面保存着结果信息,可以打印一下看一看(打印results[0])
|
||
results = model.track(rgb_img, conf=0.5, classes=[0],tracker="bytetrack.yaml", persist=True, verbose=False)
|
||
for i in range(len(results[0].boxes)):
|
||
# 框的四个坐标位置
|
||
x1, y1, x2, y2 = results[0].boxes.xyxy[i]
|
||
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||
# 将跟踪成功的物体 用框标记出来(跟踪失败的不显示)
|
||
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
|
||
|
||
# 识别对象的id,可能为None(目标跟踪失败)
|
||
if results[0].boxes.id == None:
|
||
continue
|
||
else:
|
||
ids = results[0].boxes.id[i]
|
||
if ids == None:
|
||
continue
|
||
|
||
# float转化int
|
||
ids = int(ids)
|
||
cv2.putText(frame, str(ids), (x1+10, y1), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 2)
|
||
|
||
# 计算物体中心位置的坐标
|
||
cx = int(x1 / 2 + x2 / 2)
|
||
# cy = int(y1 / 2 + y2 / 2)
|
||
cy = y2
|
||
|
||
# 判断从视频顶部进入画面的物体的中心位置与区域1的关系(在内部还是外部)
|
||
# 判断视频中行人的中心位置与区域1的关系(在内部还是外部)
|
||
inside_region1 = cv2.pointPolygonTest(region1, (cx, cy), False)
|
||
|
||
# inside_region1为1表示在区域内部
|
||
if inside_region1 > 0:
|
||
#判断当前行人id是否进入过region2区域
|
||
if ids in total_region2:
|
||
#如果当前行人id进入过region2区域,现在又在region1区域内
|
||
#证明其是从视频下方到视频上方(region2区域->region1区域),视为上行,保存到
|
||
up_id.add(ids)
|
||
# 行人越线,线的颜色改变
|
||
cv2.line(frame, (l_x1,l_y1),(l_x2,l_y2), (255, 255, 255), 3)
|
||
#进入region1区域的行人id保存在这个变量中
|
||
total_region1.add(ids)
|
||
# 保存行人图片
|
||
# img_name = file_path.split(".")[0] + "down_" + str(down_cnt) +".jpg"
|
||
# cv2.imwrite(img_name, frame[y1:y2,x1:x2])
|
||
# down_cnt = down_cnt + 1
|
||
|
||
inside_region2 = cv2.pointPolygonTest(region2, (cx, cy), False)
|
||
|
||
if inside_region2 > 0:
|
||
if ids in total_region1:
|
||
cv2.line(frame, (l_x1,l_y1),(l_x2,l_y2), (255, 255, 255), 3)
|
||
low_id.add(ids)
|
||
total_region2.add(ids)
|
||
# 保存行人图片
|
||
# img_name = file_path.split(".")[0] + "up_" + str(up_cnt) +".jpg"
|
||
# cv2.imwrite(img_name, frame[y1:y2,x1:x2])
|
||
# up_cnt = up_cnt + 1
|
||
# 统计进出物体,并设置为字符串,下一步准备在画面上显示
|
||
first_low = 'LOW:' + str(len(low_id))
|
||
first_up = 'UP:' + str(len(up_id))
|
||
|
||
# 设置画面左右上角的背景色
|
||
frame[0:40, 0:120] = (102, 0, 153)
|
||
frame[0:40, end_x - 120:end_x] = (102, 0, 153)
|
||
# 将进出物体统计的个数输出到图片,展示
|
||
cv2.putText(frame, first_low, (0, 30), font, 1, (255, 255, 255), 1)
|
||
cv2.putText(frame, first_up, (end_x - 120, 30), font, 1, (255, 255, 255), 1)
|
||
|
||
# 显示两个区域
|
||
zeros = np.zeros((frame.shape), dtype=np.uint8)
|
||
mask1 = cv2.fillPoly(zeros, [region1], color=(255, 0, 0))
|
||
mask2 = cv2.fillPoly(zeros, [region2], color=(0, 255, 0))
|
||
frame1 = cv2.addWeighted(frame, 1, mask1, 0.3, 0)
|
||
frame2 = cv2.addWeighted(frame1, 1, mask2, 0.05, 0)
|
||
# 逐帧显示处理后的画面
|
||
# out.write(frame2)
|
||
cv2.namedWindow('frame',cv2.WINDOW_NORMAL)
|
||
cv2.imshow("frame", frame2)
|
||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||
break
|
||
if cv2.getWindowProperty('frame', cv2.WND_PROP_VISIBLE) < 1:
|
||
break
|
||
kamera.release()
|
||
# out.release()
|
||
cv2.destroyAllWindows()
|
||
print('最终统计结果如下:')
|
||
print('-----------------------------------------------------------------')
|
||
print(f'LOW--->{len(low_id)}')
|
||
print(f'UP--->{len(up_id)}')
|
||
print('-----------------------------------------------------------------')
|
||
|
||
def parse_option():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--file_path', type=str)
|
||
parser.add_argument('--line', type=str)
|
||
opt = parser.parse_args()
|
||
return opt
|
||
if __name__=='__main__':
|
||
opt=parse_option()
|
||
file_path=opt.file_path
|
||
line=eval(opt.line)
|
||
exec_inference_video(file_path,line) |