|
|
import numpy as np
|
|
|
import cv2
|
|
|
import os
|
|
|
import time
|
|
|
from tqdm import tqdm
|
|
|
from ultralytics import YOLO
|
|
|
from ultralytics.yolo.utils.plotting import Annotator
|
|
|
import queue
|
|
|
|
|
|
from yolov8_det import analysis_yolov8
|
|
|
import threading
|
|
|
from config import Q_SZ
|
|
|
|
|
|
|
|
|
class DealVideo():
|
|
|
|
|
|
def __init__(self,video_file,video_save_file):
|
|
|
|
|
|
self.video_file = video_file
|
|
|
self.video_save_file = video_save_file
|
|
|
|
|
|
|
|
|
|
|
|
# 图片检测后队列
|
|
|
self.videoQueue = queue.Queue(maxsize=Q_SZ)
|
|
|
self.frameQueue = queue.Queue(maxsize=0)
|
|
|
|
|
|
#线程
|
|
|
self.get_video_listThread = threading.Thread(target=self.get_video_list)
|
|
|
self.get_video_frameThread = threading.Thread(target=self.get_video_frame)
|
|
|
self.write_videoThread = threading.Thread(target=self.write_video)
|
|
|
|
|
|
# 图像文件夹
|
|
|
def get_video_list(self):
|
|
|
|
|
|
if os.path.isdir(self.video_file):
|
|
|
|
|
|
video_ext = [".mp4", ".avi",".MP4"]
|
|
|
for maindir, subdir, file_name_list in os.walk(self.video_file):
|
|
|
for filename in file_name_list:
|
|
|
apath = os.path.join(maindir, filename)
|
|
|
ext = os.path.splitext(apath)[1]
|
|
|
if ext in video_ext:
|
|
|
self.videoQueue.put(apath)
|
|
|
|
|
|
else:
|
|
|
self.videoQueue.put(self.video_file)
|
|
|
|
|
|
|
|
|
|
|
|
def get_video_frame(self):
|
|
|
|
|
|
while True:
|
|
|
if ~self.videoQueue.empty():
|
|
|
video_path = self.videoQueue.get()
|
|
|
|
|
|
video_basename = os.path.basename(video_path).split('.')[0]
|
|
|
|
|
|
# print('video_path:',video_path)
|
|
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
|
|
|
|
frame_list = []
|
|
|
count_fps = 0
|
|
|
|
|
|
while cap.isOpened():
|
|
|
success, frame = cap.read()
|
|
|
if not success:
|
|
|
print("Ignoring empty camera frame.")
|
|
|
break
|
|
|
count_fps += 1
|
|
|
|
|
|
frame_dict = {'fps':count_fps,'frame':frame}
|
|
|
frame_list.append(frame_dict)
|
|
|
|
|
|
video_dict = {'video_path':video_path,'frame_list':frame_list,'cap':cap}
|
|
|
|
|
|
self.frameQueue.put(video_dict)
|
|
|
|
|
|
def write_video(self):
|
|
|
|
|
|
while True:
|
|
|
if ~self.frameQueue.empty():
|
|
|
video_frame_dict = self.frameQueue.get()
|
|
|
|
|
|
video_basename = os.path.basename(video_frame_dict['video_path'])
|
|
|
|
|
|
video_name_save = os.path.join(self.video_save_file, video_basename)
|
|
|
|
|
|
# 原视频帧率和尺寸
|
|
|
cap = video_frame_dict['cap']
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
|
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
|
|
|
|
|
|
# 保存视频
|
|
|
videoWriter =cv2.VideoWriter(video_name_save,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size)
|
|
|
|
|
|
frame_list = video_frame_dict['frame_list']
|
|
|
|
|
|
for i in range(len(frame_list)):
|
|
|
|
|
|
if frame_list[i]["fps"] == i + 1:
|
|
|
|
|
|
videoWriter.write(frame_list[i]["frame"])
|
|
|
|
|
|
else:
|
|
|
break
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def run(self):
|
|
|
|
|
|
self.get_video_listThread.start()
|
|
|
self.get_video_frameThread.start()
|
|
|
self.write_videoThread.start()
|
|
|
|
|
|
|
|
|
class Process_tools():
|
|
|
|
|
|
# 图像文件夹
|
|
|
def get_video_list(path):
|
|
|
video_ext = [".mp4", ".avi",".MP4"]
|
|
|
video_names = []
|
|
|
for maindir, subdir, file_name_list in os.walk(path):
|
|
|
for filename in file_name_list:
|
|
|
apath = os.path.join(maindir, filename)
|
|
|
ext = os.path.splitext(apath)[1]
|
|
|
if ext in video_ext:
|
|
|
video_names.append(apath)
|
|
|
return video_names
|
|
|
|
|
|
|
|
|
# 截取裁剪需要的视频帧
|
|
|
def save_seg_video(video_name,frameToStart,frametoStop,videoWriter,bbox):
|
|
|
|
|
|
cap = cv2.VideoCapture(video_name)
|
|
|
count = 0
|
|
|
while True:
|
|
|
|
|
|
success, frame = cap.read()
|
|
|
|
|
|
if success:
|
|
|
|
|
|
count += 1
|
|
|
if count <= frametoStop and count > frameToStart: # 选取起始帧
|
|
|
print('correct= ', count)
|
|
|
|
|
|
#裁剪视频画面
|
|
|
frame_target = frame[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])] # (split_height, split_width)
|
|
|
|
|
|
videoWriter.write(frame_target)
|
|
|
|
|
|
if not success or count >= frametoStop:
|
|
|
break
|
|
|
|
|
|
print('end')
|
|
|
|
|
|
|
|
|
|
|
|
# 获得字典中所有values值（这个值是列表）
|
|
|
def get_dict_values(lst):
|
|
|
"""
|
|
|
获取列表中所有字典的 values 值（如果值是列表）
|
|
|
|
|
|
参数:
|
|
|
lst: 包含字典的列表
|
|
|
|
|
|
返回值:
|
|
|
values: 包含所有字典的 values 值的列表（如果值是列表）
|
|
|
"""
|
|
|
return [value for dictionary in lst for value in dictionary.values() if isinstance(value, list)]
|
|
|
|
|
|
|
|
|
|
|
|
# 解析检测后的结果，为检测后的结果排序
|
|
|
def analysis_sort_list(result_dict):
|
|
|
|
|
|
# print('result_dict:',result_dict)
|
|
|
|
|
|
# 获得检测列表
|
|
|
re_list = result_dict['start_bbox']
|
|
|
# print('re_list:',re_list)
|
|
|
|
|
|
# 获得列表中所有字典的values值
|
|
|
re_bbox_list = get_dict_values(re_list)
|
|
|
|
|
|
# 为检测出来的标注框排序
|
|
|
sorted_lst = sorted(re_bbox_list, key=lambda x: x[0])
|
|
|
|
|
|
return sorted_lst
|
|
|
|
|
|
|
|
|
#对比重叠率高的两个部分，并结合标注框，保存最大的标注框
|
|
|
def contrast_bbox(e_bbox,r_bbox):
|
|
|
|
|
|
e_bbox_min = e_bbox[:2]
|
|
|
r_bbox_min = r_bbox[:2]
|
|
|
|
|
|
bbox_min = [min(x, y) for x, y in zip(e_bbox_min, r_bbox_min)]
|
|
|
|
|
|
e_bbox_max = e_bbox[-2:]
|
|
|
r_bbox_max = r_bbox[-2:]
|
|
|
|
|
|
bbox_max = [max(x, y) for x, y in zip(e_bbox_max, r_bbox_max)]
|
|
|
|
|
|
bbox = bbox_min + bbox_max
|
|
|
|
|
|
return bbox
|
|
|
|
|
|
|
|
|
# 解析result_list列表
|
|
|
def analysis_re01_list(example_dict,result_dict):
|
|
|
|
|
|
# 第一次检测到目标的帧率和信息
|
|
|
example_dict_fps = list(example_dict.keys())[0]
|
|
|
example_sorted_lst = Process_tools.analysis_sort_list(example_dict)
|
|
|
|
|
|
# 当前帧检测结果中所有的检测结果数值
|
|
|
re_dict_fps = list(result_dict.keys())[0]
|
|
|
re_dict_sorted_lst = Process_tools.analysis_sort_list(result_dict)
|
|
|
|
|
|
# 保存前后帧率连续的范围、筛选出相同的部分
|
|
|
cut_list = []
|
|
|
example_temp = []
|
|
|
re_temp = []
|
|
|
|
|
|
for i,ex_bbox in enumerate(example_sorted_lst):
|
|
|
|
|
|
for j,re_bbox in enumerate(re_dict_sorted_lst):
|
|
|
|
|
|
iou = calculate_iou(box1=ex_bbox, box2=re_bbox)
|
|
|
|
|
|
# print(iou)
|
|
|
|
|
|
if iou > 0:
|
|
|
|
|
|
bbox = contrast_bbox(e_bbox=ex_bbox,r_bbox=re_bbox)
|
|
|
|
|
|
cut_list.append({i:bbox})
|
|
|
example_temp.append(ex_bbox)
|
|
|
re_temp.append(re_bbox)
|
|
|
|
|
|
break
|
|
|
|
|
|
else:
|
|
|
continue
|
|
|
|
|
|
example_sorted_lst = [item for item in example_sorted_lst if item not in example_temp]
|
|
|
re_dict_sorted_lst = [item for item in re_dict_sorted_lst if item not in re_temp]
|
|
|
|
|
|
return cut_list,example_sorted_lst,re_dict_sorted_lst
|
|
|
|
|
|
|
|
|
# 计算前后帧率重叠范围
|
|
|
def calculate_iou(box1, box2):
|
|
|
"""
|
|
|
计算两个边界框之间的IoU值
|
|
|
|
|
|
参数:
|
|
|
box1: 边界框1的坐标（x1, y1, x2, y2）
|
|
|
box2: 边界框2的坐标（x1, y1, x2, y2）
|
|
|
|
|
|
返回值:
|
|
|
iou: 两个边界框之间的IoU值
|
|
|
"""
|
|
|
x1 = max(box1[0], box2[0])
|
|
|
y1 = max(box1[1], box2[1])
|
|
|
x2 = min(box1[2], box2[2])
|
|
|
y2 = min(box1[3], box2[3])
|
|
|
|
|
|
# 计算交集区域面积
|
|
|
intersection_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1)
|
|
|
|
|
|
# 计算边界框1和边界框2的面积
|
|
|
box1_area = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)
|
|
|
box2_area = (box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1)
|
|
|
|
|
|
# 计算并集区域面积
|
|
|
union_area = box1_area + box2_area - intersection_area
|
|
|
|
|
|
# 计算IoU值
|
|
|
iou = intersection_area / union_area
|
|
|
|
|
|
return iou
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
|
|
|
# 每个视频的时长（单位秒）
|
|
|
dertTime = 5
|
|
|
|
|
|
video = "E:/Bank_files/Bank_02/dataset/video_person/after_1/"
|
|
|
video_save = 'videos_codes_2'
|
|
|
|
|
|
# get_seg_video(video_file=video,video_save_path=video_save,dertTime=dertTime)
|
|
|
deal = DealVideo(video_file=video,video_save_file=video_save)
|
|
|
deal.run()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|