You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
import cv2
|
|
from yolov8_det import analysis_yolov8
|
|
from ultralytics import YOLO
|
|
|
|
class VideoAnalysis():
|
|
@staticmethod
|
|
def detect_head(image, person_coord, confidence, yolo_head_model, yolo_hand_model):
|
|
"""
|
|
截取后图像保存视频
|
|
"""
|
|
x_min, y_min, x_max, y_max = person_coord['person']
|
|
# 提取出人的区域
|
|
person_frame = image[int(y_min)-100:int(y_max)+100, int(x_min)-100:int(x_max)+100, :]
|
|
|
|
# 使用 YOLO 模型分析人体中头部和手部区域
|
|
head_result = analysis_yolov8(person_frame, YOLO(yolo_head_model), confidence)
|
|
hand_result = analysis_yolov8(person_frame, YOLO(yolo_hand_model), confidence)
|
|
|
|
# 如果检测到头部，返回头部坐标
|
|
if not head_result:
|
|
return head_result[1]
|
|
|
|
# 如果检测到手部，返回手部坐标
|
|
elif not head_result:
|
|
return hand_result[3]
|
|
|
|
|
|
@staticmethod
|
|
def frame_analysis(test_video_path, yolo_model, confidence, yolo_head_model, yolo_hand_model):
|
|
# 读取测试视频
|
|
video_capture = cv2.VideoCapture(test_video_path)
|
|
coords = []
|
|
while True:
|
|
result, frame_img = video_capture.read()
|
|
if result:
|
|
# 使用 YOLO 模型分析人的区域
|
|
person_result = analysis_yolov8(frame_img, YOLO(yolo_model), confidence)
|
|
# 检测头部，返回头部和手部坐标
|
|
coord = VideoAnalysis.detect_head(frame_img, person_result[0], confidence, yolo_head_model, yolo_hand_model)
|
|
|
|
if coord is not None:
|
|
coords.append(coord)
|
|
else:
|
|
break
|
|
video_capture.release()
|
|
print(coords)
|
|
|
|
# 返回坐标信息列表
|
|
return coords
|
|
|
|
def main():
|
|
test_video_path = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\video\test01_0.avi"
|
|
yolo_model = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\all_labels.pt"
|
|
yolo_head_model = r"E:\git-Xznsh\XZNSH-Code-AI\xznsh_flow2\head.pt"
|
|
yolo_hand_model = r"E:\git-Xznsh\XZNSH-Code-AI\7.27\best.pt"
|
|
confidence = {"person": 0.5, "head": 0.5, "phone": 0.5, "hand": 0.5}
|
|
VideoAnalysis.frame_analysis(test_video_path, yolo_model, confidence, yolo_head_model, yolo_hand_model)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|