from analysis_result.get_model_result import det_img from analysis_result.same_model_img import same_model_img_analysis_labels, model_labels_selet from model_load.model_load import Load_model from drawing_img.drawing_img import drawing_frame from analysis_data.data_rtsp import rtsp_para from analysis_data.data_dir_file import get_dir_file from analysis_data.config_load import get_configs from add_xml import add_xml from create_xml import create_xml from analysis_data.change_video import mp4_to_H264 import cv2 import os import time from datetime import datetime import json from loguru import logger import logging import logstash host = '192.168.10.96' xbank_logger = logging.getLogger('python-logstash-logger') xbank_logger.setLevel(logging.INFO) xbank_logger.addHandler(logstash.LogstashHandler(host, 5959, version=1)) def data_load(args): source = args[0] model_ymal = args[1] # 数据加载 rtsp_source = rtsp_para(source) dir_source = os.path.isdir(source) file_source = os.path.isfile(source) # # 模型加载 model_data = get_configs(model_ymal) model_inference = Load_model(model_file=model_data["model"], device=model_data["model_parameter"]['device'], cache_file=model_data["model_cache"]) if rtsp_source: rtsp_detect_process(source=source, model_data=model_data, model_inference=model_inference) if dir_source: dir_source_process(source, model_inference, model_data) if file_source: file_source_process(source, model_inference, model_data) def rtsp_detect_process(source, model_data, model_inference): cap = cv2.VideoCapture(source) logger.info(f"视频流{source}读取中...") # 视频流信息 fps = int(cap.get(cv2.CAP_PROP_FPS)) fps_num = fps*model_data['detect_time'] fps_num_small = fps*model_data['detect_time_small'] size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) i = 0 j = 0 n = 0 det_t_num = 0 nodet_t_num = 0 det_img = [] video_name_time = 0 det_fps_time = [] while True: try: ret, frame = cap.read() i += 1 j += 1 # 读取到当前视频帧时间 data_now = datetime.now() get_time = str(data_now.strftime("%H")) + \ str(data_now.strftime("%M")) + str(data_now.strftime("%S")) + \ str(data_now.strftime("%f")) imgframe_dict = {"path": source, 'frame': frame, 'get_fps': j, 'get_time': get_time} # 视频暂时保存路径 if video_name_time == 0: video_name_time = get_time video_path = video_name( video_name_base=video_name_time, save_path=model_data['save_videos'], save_file='temp') out_video = cv2.VideoWriter( video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, size) logger.info(f"视频{video_path}已经暂时保存...") # 模型推理 images_det_result = img_process( imgframe_dict, model_inference, model_data) images_update = save_process( imgframe_dict, images_det_result, model_data) # 结果判断，t if images_det_result: det_t_num += 1 if len(det_img) == 0: img_dict = images_update.copy() det_img.append(img_dict) if not images_det_result and len(det_img) > 0: nodet_t_num += 1 if (det_t_num + nodet_t_num) >= fps_num_small: para = determine_time( det_num=det_t_num, nodet_num=nodet_t_num, ratio_set=model_data['detect_ratio']) if para: first_fps_time = det_img[0] first_fps_time.update( {"dert_fps": (j-int(first_fps_time['get_fps'])+1)}) det_fps_time.append(first_fps_time) det_img.clear() det_t_num = 0 nodet_t_num = 0 # 视频保存 out_video.write(images_update['frame']) # 结果判断 ，T if j >= fps_num: try: out_video.release() except Exception: logger.exception(f"视频release失败") else: logger.info("视频release成功") # T时间截至，判断t时间结果。 if det_img: para = determine_time( det_num=det_t_num, nodet_num=nodet_t_num, ratio_set=model_data['detect_ratio']) first_fps_time = det_img[0] time_1 = (j-int(first_fps_time['get_fps'])+1) if para and time_1 >= (fps_num_small/2): first_fps_time.update( {"dert_fps": (j-int(first_fps_time['get_fps'])+1)}) det_fps_time.append(first_fps_time) if det_fps_time: det_fps_time = determine_duration(result_list=det_fps_time) # 转换后视频保存路径 save_video_name = os.path.basename(video_path) only_video_name = save_video_name.split('.')[0] save_video_path = os.path.join(model_data['save_videos'], (str(data_now.strftime( "%Y")) + str(data_now.strftime("%m")) + str(data_now.strftime("%d")))) save_video_path = os.path.join( save_video_path, only_video_name) # 路径 save_video = os.path.join(save_video_path, save_video_name) json_path = os.path.join( save_video_path, only_video_name + '.json') images_path = os.path.join(save_video_path, 'images') # 转换视频、保存视频 change_video = mp4_to_H264() change_video.convert_byfile(video_path, save_video) # 保存图片 update_det_fps = video_cut_images_save( det_list=det_fps_time, images_path=images_path) # print(update_det_fps) # 保存json文件 re_list, result_lables = json_get( time_list=update_det_fps, video_path=save_video, fps=fps) result_path = json_save(re_list, json_path) send_message(update_det_fps=update_det_fps, result_path=result_path, source=source, result_lables=result_lables) else: # print(video_path) os.remove(video_path) logger.info(f"未检测到目标信息的视频{video_path}删除成功") logger.info('开始信息重置') det_img.clear() det_fps_time.clear() det_t_num = 0 nodet_t_num = 0 video_name_time = 0 j = 0 # print('det_fps_time:', det_fps_time,'det_img:',det_img) # t2 = time.time() # tx = t2 - t1 # logger.info(f'检测一张图片的时间为：{tx}.') except Exception as e: logger.debug(f"读帧率失败{source}未读到...") logger.debug(e) cap.release() cap = cv2.VideoCapture(source) logger.info(f"摄像头{source}重新读取") # break def video_name(video_name_base, save_path, save_file): video_name_base = video_name_base savePath = os.path.join(save_path, save_file) if not os.path.exists(savePath): os.makedirs(savePath) video_path = os.path.join( savePath, video_name_base + '.mp4') return video_path def dir_source_process(source, model_inference, model_data): img_ext = [".jpg", ".JPG", ".bmp"] video_ext = [".mp4", ".avi", ".MP4"] img_list = get_dir_file(source, img_ext) video_list = get_dir_file(source, video_ext) if img_list: for img in img_list: t1 = time.time() images = cv2.imread(img) imgframe_dict = {"path": img, 'frame': images} images_update = img_process( imgframe_dict, model_inference, model_data) t2 = time.time() tx = t2 - t1 print('检测一张图片的时间为：', tx) if video_list: pass def file_source_process(source, model_inference, model_data): img_para = True if img_para: images = cv2.imread(source) imgframe_dict = {"path": source, 'frame': images} images_update = img_process( imgframe_dict, model_inference, model_data) def img_process(images, model_inference, model_data): # t1 = time.time() # 检测每帧图片，返回推理结果 results = det_img(model_inference=model_inference, images_frame=images['frame'], confidence=model_data["model_parameter"]['confidence'], label_name_list=model_data["model_parameter"]['label_names']) # print(results) # print(images['path']) # 根据需要挑选标注框信息 select_labels_list = model_labels_selet(example_list=model_data["model_parameter"]['compara_label_names'], result_dict_list=results) if model_data["model_parameter"]['compara_relevancy']: # 需要根据的逻辑判断标注框信息 determine_bbox = same_model_img_analysis_labels(example_list=model_data["model_parameter"]['compara_label_names'], result_dicts_list=select_labels_list, relevancy=model_data["model_parameter"]['compara_relevancy'], relevancy_para=model_data["model_parameter"]['relevancy_para']) else: determine_bbox = select_labels_list # print(determine_bbox) if model_data['model_parameter']['object_num_min']: if len(determine_bbox) >= model_data["model_parameter"]['object_num_min']: determine_bbox.clear() # logger.debug(f"正确获得检测后的信息{determine_bbox}...") # 返回检测后结果 return determine_bbox def save_process(images, determine_bbox, model_data): if determine_bbox: images.update({"results": determine_bbox}) if model_data['save_path_original']: imgname_original = images_save(images=images, save_path=model_data["save_path_original"]) img_save = drawing_frame( images_frame=images['frame'], result_list=determine_bbox) images.update({"frame": img_save}) if model_data["save_path"]: imgname = images_save( images=images, save_path=model_data["save_path"]) if model_data["save_annotations"]: if not os.path.exists(model_data["save_annotations"]): os.makedirs(model_data["save_annotations"]) save_annotations_xml( xml_save_file=model_data["save_annotations"], save_infors=determine_bbox, images=images['path']) else: pass else: # 没检测出来的图片是否保存 if model_data["test_path"]: imgname = images_save( images=images, save_path=model_data["test_path"]) # print('no:',images['path'],imgname) # 展示显示 # if images['path'] == 'rtsp://admin:@192.168.10.11': # cv2.namedWindow('11', cv2.WINDOW_NORMAL) # cv2.imshow('11',images['frame']) # cv2.waitKey(1) # cv2.destroyAllWindows() # t2 = time.time() return images def images_save(images, save_path): # 保存时候时间为图片名 images_name = images['get_time'] + '.jpg' if not os.path.exists(save_path): os.makedirs(save_path) full_name = os.path.join(save_path, images_name) cv2.imwrite(full_name, images['frame']) return full_name def save_annotations_xml(xml_save_file, save_infors, images): results = save_infors img = os.path.basename(images) img_frame = cv2.imread(images) xml_save_path = os.path.join(xml_save_file, img.split('.')[0] + '.xml') w, h, d = img_frame.shape img_shape = (w, h, d, img) if os.path.isfile(xml_save_path): add_xml(inforsDict=results, xmlFilePath=xml_save_path) else: create_xml(boxs=results, img_shape=img_shape, xml_path=xml_save_path) def determine_time(det_num, nodet_num, ratio_set): ratio = det_num / (det_num + nodet_num) if ratio >= ratio_set: return True else: return False def determine_duration(result_list): i = 0 while i < len(result_list) - 1: dict_i = result_list[i] dict_j = result_list[i + 1] if 'get_fps' in dict_i and 'dert_fps' in dict_i and 'get_fps' in dict_j: num_i = int(dict_i['get_fps']) dura_i = int(dict_i['dert_fps']) num_j = int(dict_j['get_fps']) if num_i + dura_i == num_j: dura_j = int(dict_j['dert_fps']) dura_update = dura_i + dura_j dict_i['dert_fps'] = dura_update result_list.pop(i + 1) else: i += 1 else: i += 1 return result_list # print('2:', result_list) def json_get(time_list, video_path, fps): result_dict = {'info': {'video_path': video_path, 'fps': fps}} re_dict = {} for i, det_dict in enumerate(time_list): list_hands = ["Keypad", "hands", "keyboard", "mouse", "phone"] list_sleep = ["sleep"] list_person = ["person"] if list(det_dict['results'][0].keys())[0] in list_hands: result_lables = 'playing_phone' if list(det_dict['results'][0].keys())[0] in list_sleep: result_lables = "sleep" if list(det_dict['results'][0].keys())[0] in list_person: result_lables = "person" fps_dict = {'time': det_dict['get_fps'], 'duration': det_dict['dert_fps'], 'images_path': det_dict['images_path']} re_dict.update({('id_' + str(i)): fps_dict}) result_dict.update({'result': re_dict}) return result_dict, result_lables def json_save(result_dict, json_path): result = json.dumps(result_dict) f = open(json_path, 'w') f.write(result + '\n') f.close return json_path def video_cut_images_save(det_list, images_path): for det_dict in det_list: images_path_full = images_save(images=det_dict, save_path=images_path) del det_dict['frame'] del det_dict['get_time'] det_dict.update({'images_path': images_path_full}) return det_list def send_message(update_det_fps, result_path, source, result_lables): for det_dict in update_det_fps: extra = { 'worker': 'xbank', 'time': det_dict['get_fps'], 'config_file': result_path, 'source': source, 'type': result_lables } xbank_logger.info('xBank_infer', extra=extra) logger.info(f'发送信息{extra}') # if __name__ == '__main__': # data_load(['rtsp://admin:@192.168.10.203', # 'E:/Bank_files/Bank_03/xbank_poc_test_use/config_phone.yaml'])