diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/main.py b/main.py index 2646193..8da2462 100644 --- a/main.py +++ b/main.py @@ -45,4 +45,4 @@ uvicorn_logger.setLevel(logging.INFO) - uvicorn.run(app, host="0.0.0.0", port=9299, log_config=None) + uvicorn.run(app, host="0.0.0.0", port=9000, log_config=None) diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/main.py b/main.py index 2646193..8da2462 100644 --- a/main.py +++ b/main.py @@ -45,4 +45,4 @@ uvicorn_logger.setLevel(logging.INFO) - uvicorn.run(app, host="0.0.0.0", port=9299, log_config=None) + uvicorn.run(app, host="0.0.0.0", port=9000, log_config=None) diff --git a/scene_handler/block_scene_handler.py b/scene_handler/block_scene_handler.py index c1696af..db0b841 100644 --- a/scene_handler/block_scene_handler.py +++ b/scene_handler/block_scene_handler.py @@ -9,7 +9,7 @@ from scipy.spatial import ConvexHull from algo.stream_loader import OpenCVStreamLoad -from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area +from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area, is_overlapping from common.device_status_manager import DeviceStatusManager from common.display_frame_manager import DisplayFrameManager from common.global_logger import logger @@ -42,6 +42,7 @@ 1 未检测到报警 2 人未穿戴报警 3 其他 +4 人员检测到报警 ''' ALARM_DICT = [ { @@ -90,7 +91,7 @@ }, { 'alarmCategory': 0, - 'alarmType': '2', + 'alarmType': '18', 'handelType': 2, 'category_order': -1, 'class_idx': [18], @@ -98,6 +99,31 @@ 'alarmContent': '未佩戴安全帽', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', 'label': '未佩戴安全帽', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '19', # todo + 'handelType': 4, + 'category_order': -1, + 'class_idx': [4], + 'alarm_name': 'cigarette', + 'alarmContent': '吸烟', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '吸烟', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '2', + 'handelType': 4, # todo + 'category_order': -1, + 'class_idx': [5], + 'alarm_name': 'phone', + 'alarmContent': '打电话', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '打电话', + 'model_type': 'safe', }, # todo 明火 { @@ -138,17 +164,32 @@ 'alarmType': '3', 'handelType': 2, 'category_order': 4, - 'class_idx': [24], + 'class_idx': [3], 'alarm_name': 'break_in_alarm', 'alarmContent': '非法闯入', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x37\x00\xCB', 'label': '非法闯入', + 'model_type': 'safe', }, ] COLOR_RED = (0, 0, 255) COLOR_BLUE = (255, 0, 0) +HEALTH_DEVICE_TYPE = '2' # 安全帽设备类型 +HARMFUL_DEVICE_TYPE = '4' # 四合一设备类型 + +def get_group_device_list(device_code): + health_device_codes = [] + harmful_device_codes = [] + url = f'http://111.198.10.15:22006/v3/device/listGroupDevs?devcode={device_code}' + response = get_request(url) + if response and response.get('code') == 200 and response.get('data'): + data = response.get('data') + for item in data: + health_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HEALTH_DEVICE_TYPE] + harmful_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HARMFUL_DEVICE_TYPE] + return health_device_codes, harmful_device_codes class BlockSceneHandler(BaseSceneHandler): @@ -170,8 +211,9 @@ self.display_frame_manager = DisplayFrameManager() # todo 要改成通过后台接口读取设备编号 - self.health_device_codes = ['HWIH061000056395'] - self.harmful_device_codes = ['862635063168165A'] + # self.health_device_codes = ['HWIH061000056395'] + # self.harmful_device_codes = ['862635063168165A'] + self.health_device_codes, self.harmful_device_codes = get_group_device_list(device.code) self.thread_pool.submit_task(self.gas_data_task, device.code) for helmet_code in self.health_device_codes: @@ -182,20 +224,21 @@ self.thread_pool.submit_task(self.alarm_message_center.process_messages) # todo 明火 - self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + # self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + self.model = YOLO('weights/labor-v8-20241114.pt') self.model_classes = { # 0: '三脚架', - 3: '人', + # 3: '人', 4: '作业信息公示牌', 6: '危险告知牌', 9: '反光衣', # 11: '呼吸面罩', # 13: '四合一', - 15: '头', + # 15: '头', 16: '安全告知牌', - 18: '安全帽', + # 18: '安全帽', 20: '安全标识牌', - 24: '工服', + # 24: '工服', 34: '灭火器', 43: '警戒线', 48: '路锥', @@ -204,6 +247,12 @@ self.PERSON_CLASS_IDX = 3 self.HEAD_CLASS_IDX = 15 + self.safe_model = YOLO('weights/yinhuan.pt') + self.safe_model_classes = {0: '人', 1: '头', 2: '安全帽', 3: '工服', 4: '烟头', 5: '电话', 6: '袖标'} + self.PERSON_CLASS_IDX = 0 + self.HEAD_CLASS_IDX = 1 + self.SAFETY_CLASS_IDX = [2,3,6] + self.vid_stride = 3 self.stream_loader = OpenCVStreamLoad(camera_url=device.input_stream_url, camera_code=device.code, @@ -350,20 +399,33 @@ # todo 需要生成报警记录吗 def model_predict(self, frames): + result_boxes = [] + safe_result_boxes = [] + results_generator = self.model.track(frames, save_txt=False, save=False, verbose=False, conf=0.5, classes=list(self.model_classes.keys()), imgsz=640, stream=True) - result_boxes = [] + for r in results_generator: result_boxes.append(r.boxes) - return result_boxes + safe_results_generator = self.safe_model.track(frames,save_txt=False, save=False, verbose=False, conf=0.5, + classes=list(self.safe_model_classes.keys()), + imgsz=640, + stream=True) + for s in safe_results_generator: + safe_result_boxes.append(s.boxes) - def handle_behave_alarm(self, frames, result_boxes): + return result_boxes, safe_result_boxes + + def handle_behave_alarm(self, frames, result_boxes, safe_result_boxes): behave_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 0] for alarm_dict in behave_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + use_safe_model = alarm_dict.get('model_type', 'labor') == 'safe' + boxes = safe_result_boxes if use_safe_model else result_boxes + model_classes = self.safe_model_classes if use_safe_model else self.model_classes + for idx, frame_boxes in enumerate(boxes): frame = frames[idx] object_boxes = [box for box in frame_boxes if int(box.cls) in alarm_dict['class_idx']] if alarm_dict['handelType'] == 0: # 检测到就报警 @@ -371,12 +433,20 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 红色标注目标设备，蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: - box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=box_color, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", - color=box_color, + f"{self.safe_model_classes[box_cls]}", + color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) @@ -391,10 +461,18 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{self.safe_model_classes[box_cls]}", color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, @@ -431,6 +509,7 @@ has_alarm = True if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人 annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) # 已报警，清零，重新计数 person_status[alarm_dict['alarm_name']] = 0 @@ -438,20 +517,85 @@ if has_alarm: self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 for box in frame_boxes: - box_cls = box.cls - if box_cls != self.PERSON_CLASS_IDX and box_cls != self.HEAD_CLASS_IDX: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{model_classes[int(box.cls)]}", color=COLOR_BLUE, rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) + elif alarm_dict['handelType'] == 4: # 人检测到报警：吸烟、打电话 + person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] + has_alarm = False + annotator = None + for person_box in person_boxes: + if person_box.id is None: + continue - def handle_break_in_alarm(self, frames, result_boxes): + person_bbox = person_box.xyxy.cpu().squeeze() + person_id = int(person_box.id) + person_object_box = max( + (box for box in object_boxes if is_overlapping(person_bbox, box.xyxy.cpu().squeeze())), + key=lambda box: box.conf.item(), + default=None + ) + has_object = person_object_box is not None + + + person_status = self.tracking_status[person_id] + if alarm_dict['alarm_name'] not in person_status: + person_status[alarm_dict['alarm_name']] = 0 + if has_object: + person_status[alarm_dict['alarm_name']] += 1 + else: + person_status[alarm_dict['alarm_name']] = 0 + + person_alarm = person_status[alarm_dict['alarm_name']] > self.max_missing_frames + if person_alarm: + has_alarm = True + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人、目标 + annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) + annotator.box_label(person_object_box.xyxy.cpu().squeeze(),'',color=COLOR_RED,rotated=False) + # 已报警，清零，重新计数 + person_status[alarm_dict['alarm_name']] = 0 + + if has_alarm: + self.alarm_message_center.add_message(alarm_dict) + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 + for box in frame_boxes: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, + alarm_np_img=annotator.result()) + + + def handle_break_in_alarm(self, frames, result_boxes, safe_result_boxes): break_in_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 3] for alarm_dict in break_in_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): frame = frames[idx] person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] head_boxes = [box for box in frame_boxes if int(box.cls) == self.HEAD_CLASS_IDX] @@ -524,10 +668,10 @@ t1 = time.time() self.device_status_manager.set_status(device_id=self.device.id) - result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame + result_boxes, safe_result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame t2 = time.time() - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): current_person_ids = {int(box.id) for box in frame_boxes if box.cls is not None and box.id is not None and int( box.cls) == self.PERSON_CLASS_IDX} @@ -542,12 +686,12 @@ if self.tracking_status[person_id]['disappear_frames'] > self.disappear_threshold: self.tracking_status.pop(person_id) - self.handle_behave_alarm(frames, result_boxes) - self.handle_break_in_alarm(frames, result_boxes) + self.handle_behave_alarm(frames, result_boxes, safe_result_boxes) + self.handle_break_in_alarm(frames, result_boxes, safe_result_boxes) t3 = time.time() - # for person_id in self.tracking_status.keys(): - # print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') + for person_id in self.tracking_status.keys(): + print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') for idx, frame in enumerate(frames): annotator = Annotator(frame, None, 18, "Arial.ttf", False, example="人") @@ -557,6 +701,11 @@ f"{self.model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", color=colors(int(s_box.cls)), rotated=False) + for s_box in safe_result_boxes[idx]: + annotator.box_label(s_box.xyxy.cpu().squeeze(), + f"{self.safe_model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", + color=colors(int(s_box.cls)), + rotated=False) self.display_frame_manager.add_frame(self.device.id, annotator.result()) # self.display_frame_manager.add_frame(self.device.id, frames[idx]) diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/main.py b/main.py index 2646193..8da2462 100644 --- a/main.py +++ b/main.py @@ -45,4 +45,4 @@ uvicorn_logger.setLevel(logging.INFO) - uvicorn.run(app, host="0.0.0.0", port=9299, log_config=None) + uvicorn.run(app, host="0.0.0.0", port=9000, log_config=None) diff --git a/scene_handler/block_scene_handler.py b/scene_handler/block_scene_handler.py index c1696af..db0b841 100644 --- a/scene_handler/block_scene_handler.py +++ b/scene_handler/block_scene_handler.py @@ -9,7 +9,7 @@ from scipy.spatial import ConvexHull from algo.stream_loader import OpenCVStreamLoad -from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area +from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area, is_overlapping from common.device_status_manager import DeviceStatusManager from common.display_frame_manager import DisplayFrameManager from common.global_logger import logger @@ -42,6 +42,7 @@ 1 未检测到报警 2 人未穿戴报警 3 其他 +4 人员检测到报警 ''' ALARM_DICT = [ { @@ -90,7 +91,7 @@ }, { 'alarmCategory': 0, - 'alarmType': '2', + 'alarmType': '18', 'handelType': 2, 'category_order': -1, 'class_idx': [18], @@ -98,6 +99,31 @@ 'alarmContent': '未佩戴安全帽', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', 'label': '未佩戴安全帽', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '19', # todo + 'handelType': 4, + 'category_order': -1, + 'class_idx': [4], + 'alarm_name': 'cigarette', + 'alarmContent': '吸烟', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '吸烟', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '2', + 'handelType': 4, # todo + 'category_order': -1, + 'class_idx': [5], + 'alarm_name': 'phone', + 'alarmContent': '打电话', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '打电话', + 'model_type': 'safe', }, # todo 明火 { @@ -138,17 +164,32 @@ 'alarmType': '3', 'handelType': 2, 'category_order': 4, - 'class_idx': [24], + 'class_idx': [3], 'alarm_name': 'break_in_alarm', 'alarmContent': '非法闯入', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x37\x00\xCB', 'label': '非法闯入', + 'model_type': 'safe', }, ] COLOR_RED = (0, 0, 255) COLOR_BLUE = (255, 0, 0) +HEALTH_DEVICE_TYPE = '2' # 安全帽设备类型 +HARMFUL_DEVICE_TYPE = '4' # 四合一设备类型 + +def get_group_device_list(device_code): + health_device_codes = [] + harmful_device_codes = [] + url = f'http://111.198.10.15:22006/v3/device/listGroupDevs?devcode={device_code}' + response = get_request(url) + if response and response.get('code') == 200 and response.get('data'): + data = response.get('data') + for item in data: + health_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HEALTH_DEVICE_TYPE] + harmful_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HARMFUL_DEVICE_TYPE] + return health_device_codes, harmful_device_codes class BlockSceneHandler(BaseSceneHandler): @@ -170,8 +211,9 @@ self.display_frame_manager = DisplayFrameManager() # todo 要改成通过后台接口读取设备编号 - self.health_device_codes = ['HWIH061000056395'] - self.harmful_device_codes = ['862635063168165A'] + # self.health_device_codes = ['HWIH061000056395'] + # self.harmful_device_codes = ['862635063168165A'] + self.health_device_codes, self.harmful_device_codes = get_group_device_list(device.code) self.thread_pool.submit_task(self.gas_data_task, device.code) for helmet_code in self.health_device_codes: @@ -182,20 +224,21 @@ self.thread_pool.submit_task(self.alarm_message_center.process_messages) # todo 明火 - self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + # self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + self.model = YOLO('weights/labor-v8-20241114.pt') self.model_classes = { # 0: '三脚架', - 3: '人', + # 3: '人', 4: '作业信息公示牌', 6: '危险告知牌', 9: '反光衣', # 11: '呼吸面罩', # 13: '四合一', - 15: '头', + # 15: '头', 16: '安全告知牌', - 18: '安全帽', + # 18: '安全帽', 20: '安全标识牌', - 24: '工服', + # 24: '工服', 34: '灭火器', 43: '警戒线', 48: '路锥', @@ -204,6 +247,12 @@ self.PERSON_CLASS_IDX = 3 self.HEAD_CLASS_IDX = 15 + self.safe_model = YOLO('weights/yinhuan.pt') + self.safe_model_classes = {0: '人', 1: '头', 2: '安全帽', 3: '工服', 4: '烟头', 5: '电话', 6: '袖标'} + self.PERSON_CLASS_IDX = 0 + self.HEAD_CLASS_IDX = 1 + self.SAFETY_CLASS_IDX = [2,3,6] + self.vid_stride = 3 self.stream_loader = OpenCVStreamLoad(camera_url=device.input_stream_url, camera_code=device.code, @@ -350,20 +399,33 @@ # todo 需要生成报警记录吗 def model_predict(self, frames): + result_boxes = [] + safe_result_boxes = [] + results_generator = self.model.track(frames, save_txt=False, save=False, verbose=False, conf=0.5, classes=list(self.model_classes.keys()), imgsz=640, stream=True) - result_boxes = [] + for r in results_generator: result_boxes.append(r.boxes) - return result_boxes + safe_results_generator = self.safe_model.track(frames,save_txt=False, save=False, verbose=False, conf=0.5, + classes=list(self.safe_model_classes.keys()), + imgsz=640, + stream=True) + for s in safe_results_generator: + safe_result_boxes.append(s.boxes) - def handle_behave_alarm(self, frames, result_boxes): + return result_boxes, safe_result_boxes + + def handle_behave_alarm(self, frames, result_boxes, safe_result_boxes): behave_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 0] for alarm_dict in behave_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + use_safe_model = alarm_dict.get('model_type', 'labor') == 'safe' + boxes = safe_result_boxes if use_safe_model else result_boxes + model_classes = self.safe_model_classes if use_safe_model else self.model_classes + for idx, frame_boxes in enumerate(boxes): frame = frames[idx] object_boxes = [box for box in frame_boxes if int(box.cls) in alarm_dict['class_idx']] if alarm_dict['handelType'] == 0: # 检测到就报警 @@ -371,12 +433,20 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 红色标注目标设备，蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: - box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=box_color, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", - color=box_color, + f"{self.safe_model_classes[box_cls]}", + color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) @@ -391,10 +461,18 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{self.safe_model_classes[box_cls]}", color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, @@ -431,6 +509,7 @@ has_alarm = True if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人 annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) # 已报警，清零，重新计数 person_status[alarm_dict['alarm_name']] = 0 @@ -438,20 +517,85 @@ if has_alarm: self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 for box in frame_boxes: - box_cls = box.cls - if box_cls != self.PERSON_CLASS_IDX and box_cls != self.HEAD_CLASS_IDX: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{model_classes[int(box.cls)]}", color=COLOR_BLUE, rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) + elif alarm_dict['handelType'] == 4: # 人检测到报警：吸烟、打电话 + person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] + has_alarm = False + annotator = None + for person_box in person_boxes: + if person_box.id is None: + continue - def handle_break_in_alarm(self, frames, result_boxes): + person_bbox = person_box.xyxy.cpu().squeeze() + person_id = int(person_box.id) + person_object_box = max( + (box for box in object_boxes if is_overlapping(person_bbox, box.xyxy.cpu().squeeze())), + key=lambda box: box.conf.item(), + default=None + ) + has_object = person_object_box is not None + + + person_status = self.tracking_status[person_id] + if alarm_dict['alarm_name'] not in person_status: + person_status[alarm_dict['alarm_name']] = 0 + if has_object: + person_status[alarm_dict['alarm_name']] += 1 + else: + person_status[alarm_dict['alarm_name']] = 0 + + person_alarm = person_status[alarm_dict['alarm_name']] > self.max_missing_frames + if person_alarm: + has_alarm = True + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人、目标 + annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) + annotator.box_label(person_object_box.xyxy.cpu().squeeze(),'',color=COLOR_RED,rotated=False) + # 已报警，清零，重新计数 + person_status[alarm_dict['alarm_name']] = 0 + + if has_alarm: + self.alarm_message_center.add_message(alarm_dict) + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 + for box in frame_boxes: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, + alarm_np_img=annotator.result()) + + + def handle_break_in_alarm(self, frames, result_boxes, safe_result_boxes): break_in_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 3] for alarm_dict in break_in_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): frame = frames[idx] person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] head_boxes = [box for box in frame_boxes if int(box.cls) == self.HEAD_CLASS_IDX] @@ -524,10 +668,10 @@ t1 = time.time() self.device_status_manager.set_status(device_id=self.device.id) - result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame + result_boxes, safe_result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame t2 = time.time() - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): current_person_ids = {int(box.id) for box in frame_boxes if box.cls is not None and box.id is not None and int( box.cls) == self.PERSON_CLASS_IDX} @@ -542,12 +686,12 @@ if self.tracking_status[person_id]['disappear_frames'] > self.disappear_threshold: self.tracking_status.pop(person_id) - self.handle_behave_alarm(frames, result_boxes) - self.handle_break_in_alarm(frames, result_boxes) + self.handle_behave_alarm(frames, result_boxes, safe_result_boxes) + self.handle_break_in_alarm(frames, result_boxes, safe_result_boxes) t3 = time.time() - # for person_id in self.tracking_status.keys(): - # print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') + for person_id in self.tracking_status.keys(): + print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') for idx, frame in enumerate(frames): annotator = Annotator(frame, None, 18, "Arial.ttf", False, example="人") @@ -557,6 +701,11 @@ f"{self.model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", color=colors(int(s_box.cls)), rotated=False) + for s_box in safe_result_boxes[idx]: + annotator.box_label(s_box.xyxy.cpu().squeeze(), + f"{self.safe_model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", + color=colors(int(s_box.cls)), + rotated=False) self.display_frame_manager.add_frame(self.device.id, annotator.result()) # self.display_frame_manager.add_frame(self.device.id, frames[idx]) diff --git a/services/schedule_job.py b/services/schedule_job.py index a6e5a2f..b30acf8 100644 --- a/services/schedule_job.py +++ b/services/schedule_job.py @@ -10,8 +10,6 @@ from services.alarm_record_service import AlarmRecordService from services.device_frame_service import DeviceFrameService -from apscheduler.schedulers.asyncio import AsyncIOScheduler -from pytz import timezone import re import gzip diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/main.py b/main.py index 2646193..8da2462 100644 --- a/main.py +++ b/main.py @@ -45,4 +45,4 @@ uvicorn_logger.setLevel(logging.INFO) - uvicorn.run(app, host="0.0.0.0", port=9299, log_config=None) + uvicorn.run(app, host="0.0.0.0", port=9000, log_config=None) diff --git a/scene_handler/block_scene_handler.py b/scene_handler/block_scene_handler.py index c1696af..db0b841 100644 --- a/scene_handler/block_scene_handler.py +++ b/scene_handler/block_scene_handler.py @@ -9,7 +9,7 @@ from scipy.spatial import ConvexHull from algo.stream_loader import OpenCVStreamLoad -from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area +from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area, is_overlapping from common.device_status_manager import DeviceStatusManager from common.display_frame_manager import DisplayFrameManager from common.global_logger import logger @@ -42,6 +42,7 @@ 1 未检测到报警 2 人未穿戴报警 3 其他 +4 人员检测到报警 ''' ALARM_DICT = [ { @@ -90,7 +91,7 @@ }, { 'alarmCategory': 0, - 'alarmType': '2', + 'alarmType': '18', 'handelType': 2, 'category_order': -1, 'class_idx': [18], @@ -98,6 +99,31 @@ 'alarmContent': '未佩戴安全帽', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', 'label': '未佩戴安全帽', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '19', # todo + 'handelType': 4, + 'category_order': -1, + 'class_idx': [4], + 'alarm_name': 'cigarette', + 'alarmContent': '吸烟', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '吸烟', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '2', + 'handelType': 4, # todo + 'category_order': -1, + 'class_idx': [5], + 'alarm_name': 'phone', + 'alarmContent': '打电话', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '打电话', + 'model_type': 'safe', }, # todo 明火 { @@ -138,17 +164,32 @@ 'alarmType': '3', 'handelType': 2, 'category_order': 4, - 'class_idx': [24], + 'class_idx': [3], 'alarm_name': 'break_in_alarm', 'alarmContent': '非法闯入', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x37\x00\xCB', 'label': '非法闯入', + 'model_type': 'safe', }, ] COLOR_RED = (0, 0, 255) COLOR_BLUE = (255, 0, 0) +HEALTH_DEVICE_TYPE = '2' # 安全帽设备类型 +HARMFUL_DEVICE_TYPE = '4' # 四合一设备类型 + +def get_group_device_list(device_code): + health_device_codes = [] + harmful_device_codes = [] + url = f'http://111.198.10.15:22006/v3/device/listGroupDevs?devcode={device_code}' + response = get_request(url) + if response and response.get('code') == 200 and response.get('data'): + data = response.get('data') + for item in data: + health_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HEALTH_DEVICE_TYPE] + harmful_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HARMFUL_DEVICE_TYPE] + return health_device_codes, harmful_device_codes class BlockSceneHandler(BaseSceneHandler): @@ -170,8 +211,9 @@ self.display_frame_manager = DisplayFrameManager() # todo 要改成通过后台接口读取设备编号 - self.health_device_codes = ['HWIH061000056395'] - self.harmful_device_codes = ['862635063168165A'] + # self.health_device_codes = ['HWIH061000056395'] + # self.harmful_device_codes = ['862635063168165A'] + self.health_device_codes, self.harmful_device_codes = get_group_device_list(device.code) self.thread_pool.submit_task(self.gas_data_task, device.code) for helmet_code in self.health_device_codes: @@ -182,20 +224,21 @@ self.thread_pool.submit_task(self.alarm_message_center.process_messages) # todo 明火 - self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + # self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + self.model = YOLO('weights/labor-v8-20241114.pt') self.model_classes = { # 0: '三脚架', - 3: '人', + # 3: '人', 4: '作业信息公示牌', 6: '危险告知牌', 9: '反光衣', # 11: '呼吸面罩', # 13: '四合一', - 15: '头', + # 15: '头', 16: '安全告知牌', - 18: '安全帽', + # 18: '安全帽', 20: '安全标识牌', - 24: '工服', + # 24: '工服', 34: '灭火器', 43: '警戒线', 48: '路锥', @@ -204,6 +247,12 @@ self.PERSON_CLASS_IDX = 3 self.HEAD_CLASS_IDX = 15 + self.safe_model = YOLO('weights/yinhuan.pt') + self.safe_model_classes = {0: '人', 1: '头', 2: '安全帽', 3: '工服', 4: '烟头', 5: '电话', 6: '袖标'} + self.PERSON_CLASS_IDX = 0 + self.HEAD_CLASS_IDX = 1 + self.SAFETY_CLASS_IDX = [2,3,6] + self.vid_stride = 3 self.stream_loader = OpenCVStreamLoad(camera_url=device.input_stream_url, camera_code=device.code, @@ -350,20 +399,33 @@ # todo 需要生成报警记录吗 def model_predict(self, frames): + result_boxes = [] + safe_result_boxes = [] + results_generator = self.model.track(frames, save_txt=False, save=False, verbose=False, conf=0.5, classes=list(self.model_classes.keys()), imgsz=640, stream=True) - result_boxes = [] + for r in results_generator: result_boxes.append(r.boxes) - return result_boxes + safe_results_generator = self.safe_model.track(frames,save_txt=False, save=False, verbose=False, conf=0.5, + classes=list(self.safe_model_classes.keys()), + imgsz=640, + stream=True) + for s in safe_results_generator: + safe_result_boxes.append(s.boxes) - def handle_behave_alarm(self, frames, result_boxes): + return result_boxes, safe_result_boxes + + def handle_behave_alarm(self, frames, result_boxes, safe_result_boxes): behave_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 0] for alarm_dict in behave_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + use_safe_model = alarm_dict.get('model_type', 'labor') == 'safe' + boxes = safe_result_boxes if use_safe_model else result_boxes + model_classes = self.safe_model_classes if use_safe_model else self.model_classes + for idx, frame_boxes in enumerate(boxes): frame = frames[idx] object_boxes = [box for box in frame_boxes if int(box.cls) in alarm_dict['class_idx']] if alarm_dict['handelType'] == 0: # 检测到就报警 @@ -371,12 +433,20 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 红色标注目标设备，蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: - box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=box_color, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", - color=box_color, + f"{self.safe_model_classes[box_cls]}", + color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) @@ -391,10 +461,18 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{self.safe_model_classes[box_cls]}", color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, @@ -431,6 +509,7 @@ has_alarm = True if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人 annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) # 已报警，清零，重新计数 person_status[alarm_dict['alarm_name']] = 0 @@ -438,20 +517,85 @@ if has_alarm: self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 for box in frame_boxes: - box_cls = box.cls - if box_cls != self.PERSON_CLASS_IDX and box_cls != self.HEAD_CLASS_IDX: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{model_classes[int(box.cls)]}", color=COLOR_BLUE, rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) + elif alarm_dict['handelType'] == 4: # 人检测到报警：吸烟、打电话 + person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] + has_alarm = False + annotator = None + for person_box in person_boxes: + if person_box.id is None: + continue - def handle_break_in_alarm(self, frames, result_boxes): + person_bbox = person_box.xyxy.cpu().squeeze() + person_id = int(person_box.id) + person_object_box = max( + (box for box in object_boxes if is_overlapping(person_bbox, box.xyxy.cpu().squeeze())), + key=lambda box: box.conf.item(), + default=None + ) + has_object = person_object_box is not None + + + person_status = self.tracking_status[person_id] + if alarm_dict['alarm_name'] not in person_status: + person_status[alarm_dict['alarm_name']] = 0 + if has_object: + person_status[alarm_dict['alarm_name']] += 1 + else: + person_status[alarm_dict['alarm_name']] = 0 + + person_alarm = person_status[alarm_dict['alarm_name']] > self.max_missing_frames + if person_alarm: + has_alarm = True + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人、目标 + annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) + annotator.box_label(person_object_box.xyxy.cpu().squeeze(),'',color=COLOR_RED,rotated=False) + # 已报警，清零，重新计数 + person_status[alarm_dict['alarm_name']] = 0 + + if has_alarm: + self.alarm_message_center.add_message(alarm_dict) + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 + for box in frame_boxes: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, + alarm_np_img=annotator.result()) + + + def handle_break_in_alarm(self, frames, result_boxes, safe_result_boxes): break_in_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 3] for alarm_dict in break_in_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): frame = frames[idx] person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] head_boxes = [box for box in frame_boxes if int(box.cls) == self.HEAD_CLASS_IDX] @@ -524,10 +668,10 @@ t1 = time.time() self.device_status_manager.set_status(device_id=self.device.id) - result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame + result_boxes, safe_result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame t2 = time.time() - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): current_person_ids = {int(box.id) for box in frame_boxes if box.cls is not None and box.id is not None and int( box.cls) == self.PERSON_CLASS_IDX} @@ -542,12 +686,12 @@ if self.tracking_status[person_id]['disappear_frames'] > self.disappear_threshold: self.tracking_status.pop(person_id) - self.handle_behave_alarm(frames, result_boxes) - self.handle_break_in_alarm(frames, result_boxes) + self.handle_behave_alarm(frames, result_boxes, safe_result_boxes) + self.handle_break_in_alarm(frames, result_boxes, safe_result_boxes) t3 = time.time() - # for person_id in self.tracking_status.keys(): - # print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') + for person_id in self.tracking_status.keys(): + print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') for idx, frame in enumerate(frames): annotator = Annotator(frame, None, 18, "Arial.ttf", False, example="人") @@ -557,6 +701,11 @@ f"{self.model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", color=colors(int(s_box.cls)), rotated=False) + for s_box in safe_result_boxes[idx]: + annotator.box_label(s_box.xyxy.cpu().squeeze(), + f"{self.safe_model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", + color=colors(int(s_box.cls)), + rotated=False) self.display_frame_manager.add_frame(self.device.id, annotator.result()) # self.display_frame_manager.add_frame(self.device.id, frames[idx]) diff --git a/services/schedule_job.py b/services/schedule_job.py index a6e5a2f..b30acf8 100644 --- a/services/schedule_job.py +++ b/services/schedule_job.py @@ -10,8 +10,6 @@ from services.alarm_record_service import AlarmRecordService from services.device_frame_service import DeviceFrameService -from apscheduler.schedulers.asyncio import AsyncIOScheduler -from pytz import timezone import re import gzip diff --git a/tcp/tcp_client_connector.py b/tcp/tcp_client_connector.py index 8dbd86c..6e0b776 100644 --- a/tcp/tcp_client_connector.py +++ b/tcp/tcp_client_connector.py @@ -81,28 +81,49 @@ """连接到设备""" while True: try: + if self.is_connected: + logger.info(f"已连接到 {self.ip}:{self.port}") + return + logger.info(f"正在连接到 {self.ip}:{self.port}...") - # 设置连接超时 + + # 使用 asyncio.open_connection 的超时包装 + connect_task = asyncio.open_connection(self.ip, self.port) self.reader, self.writer = await asyncio.wait_for( - asyncio.open_connection(self.ip, self.port), timeout=self.timeout + connect_task, + timeout=self.timeout ) + + # 验证连接是否真正建立 + if self.writer is None or self.writer.is_closing(): + raise ConnectionError("连接未能成功建立") + self.is_connected = True logger.info(f"已连接到 {self.ip}:{self.port}") if not self.tasks_started: - asyncio.create_task(self.process_message_queue()) # Start processing message queue + asyncio.create_task(self.process_message_queue()) asyncio.create_task(self.start_gas_query()) self.tasks_started = True - break + return # 连接成功后直接返回 - except (asyncio.TimeoutError, ConnectionRefusedError, OSError) as e: - logger.error(f"连接到 {self.ip}:{self.port} 失败，错误: {e}") - logger.info(f"{self.reconnect_interval} 秒后将重连到 {self.ip}:{self.port}") + except asyncio.TimeoutError: + logger.error(f"连接超时: {self.ip}:{self.port}") await asyncio.sleep(self.reconnect_interval) + + except (ConnectionRefusedError, OSError) as e: + logger.error(f"连接被拒绝或网络错误: {self.ip}:{self.port}, 错误: {e}") + await asyncio.sleep(self.reconnect_interval) + except Exception as e: - # 给未知异常一个兜底打印，方便排查 - logger.exception(f"连接到 {self.ip}:{self.port} 遇到未知错误: {e}") + logger.exception(f"连接时发生未知错误: {e}") await asyncio.sleep(self.reconnect_interval) + + finally: + if not self.is_connected: + logger.info(f"{self.reconnect_interval} 秒后将重试连接 {self.ip}:{self.port}") + else: + logger.info(f"已连接到 {self.ip}:{self.port}") async def reconnect(self): """处理断线重连""" @@ -138,6 +159,7 @@ except (ConnectionResetError, BrokenPipeError) as e: logger.exception(f"Error during disconnection") self.reader = self.writer = None + self.is_connected = False logger.info(f"Disconnected from {self.ip}:{self.port}") async def start_gas_query(self): diff --git a/.gitignore b/.gitignore index b82300a..5e9050c 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,7 @@ logs test* *.zip -.idea/* \ No newline at end of file +.idea/* +__pycache__/ +**/__pycache__/ +*.py[cod] \ No newline at end of file diff --git a/common/detect_utils.py b/common/detect_utils.py index 33c340c..f664f04 100644 --- a/common/detect_utils.py +++ b/common/detect_utils.py @@ -48,6 +48,13 @@ return width * height +def is_overlapping(bbox1, bbox2): + # 检查两个坐标框是否重叠 + x1, y1, x2, y2 = bbox1 + x3, y3, x4, y4 = bbox2 + + return not (x2 < x3 or x4 < x1 or y2 < y3 or y4 < y1) + def bbox_area(bbox): # 计算坐标框的面积 diff --git a/db/safe-algo-pro.db b/db/safe-algo-pro.db index 08505c8..966beaa 100644 --- a/db/safe-algo-pro.db +++ b/db/safe-algo-pro.db Binary files differ diff --git a/main.py b/main.py index 2646193..8da2462 100644 --- a/main.py +++ b/main.py @@ -45,4 +45,4 @@ uvicorn_logger.setLevel(logging.INFO) - uvicorn.run(app, host="0.0.0.0", port=9299, log_config=None) + uvicorn.run(app, host="0.0.0.0", port=9000, log_config=None) diff --git a/scene_handler/block_scene_handler.py b/scene_handler/block_scene_handler.py index c1696af..db0b841 100644 --- a/scene_handler/block_scene_handler.py +++ b/scene_handler/block_scene_handler.py @@ -9,7 +9,7 @@ from scipy.spatial import ConvexHull from algo.stream_loader import OpenCVStreamLoad -from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area +from common.detect_utils import is_within_alert_range, get_person_head, intersection_area, bbox_area, is_overlapping from common.device_status_manager import DeviceStatusManager from common.display_frame_manager import DisplayFrameManager from common.global_logger import logger @@ -42,6 +42,7 @@ 1 未检测到报警 2 人未穿戴报警 3 其他 +4 人员检测到报警 ''' ALARM_DICT = [ { @@ -90,7 +91,7 @@ }, { 'alarmCategory': 0, - 'alarmType': '2', + 'alarmType': '18', 'handelType': 2, 'category_order': -1, 'class_idx': [18], @@ -98,6 +99,31 @@ 'alarmContent': '未佩戴安全帽', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', 'label': '未佩戴安全帽', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '19', # todo + 'handelType': 4, + 'category_order': -1, + 'class_idx': [4], + 'alarm_name': 'cigarette', + 'alarmContent': '吸烟', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '吸烟', + 'model_type': 'safe', + }, + { + 'alarmCategory': 0, + 'alarmType': '2', + 'handelType': 4, # todo + 'category_order': -1, + 'class_idx': [5], + 'alarm_name': 'phone', + 'alarmContent': '打电话', + 'alarmSoundMessage': b'\xaa\x01\x00\x93\x01\x00\x95', # todo + 'label': '打电话', + 'model_type': 'safe', }, # todo 明火 { @@ -138,17 +164,32 @@ 'alarmType': '3', 'handelType': 2, 'category_order': 4, - 'class_idx': [24], + 'class_idx': [3], 'alarm_name': 'break_in_alarm', 'alarmContent': '非法闯入', 'alarmSoundMessage': b'\xaa\x01\x00\x93\x37\x00\xCB', 'label': '非法闯入', + 'model_type': 'safe', }, ] COLOR_RED = (0, 0, 255) COLOR_BLUE = (255, 0, 0) +HEALTH_DEVICE_TYPE = '2' # 安全帽设备类型 +HARMFUL_DEVICE_TYPE = '4' # 四合一设备类型 + +def get_group_device_list(device_code): + health_device_codes = [] + harmful_device_codes = [] + url = f'http://111.198.10.15:22006/v3/device/listGroupDevs?devcode={device_code}' + response = get_request(url) + if response and response.get('code') == 200 and response.get('data'): + data = response.get('data') + for item in data: + health_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HEALTH_DEVICE_TYPE] + harmful_device_codes = [item.get('deviceCode', '') for item in data if item.get('deviceType', '') == HARMFUL_DEVICE_TYPE] + return health_device_codes, harmful_device_codes class BlockSceneHandler(BaseSceneHandler): @@ -170,8 +211,9 @@ self.display_frame_manager = DisplayFrameManager() # todo 要改成通过后台接口读取设备编号 - self.health_device_codes = ['HWIH061000056395'] - self.harmful_device_codes = ['862635063168165A'] + # self.health_device_codes = ['HWIH061000056395'] + # self.harmful_device_codes = ['862635063168165A'] + self.health_device_codes, self.harmful_device_codes = get_group_device_list(device.code) self.thread_pool.submit_task(self.gas_data_task, device.code) for helmet_code in self.health_device_codes: @@ -182,20 +224,21 @@ self.thread_pool.submit_task(self.alarm_message_center.process_messages) # todo 明火 - self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + # self.model = YOLO('weights/labor-v8-20250115-fp16.engine') + self.model = YOLO('weights/labor-v8-20241114.pt') self.model_classes = { # 0: '三脚架', - 3: '人', + # 3: '人', 4: '作业信息公示牌', 6: '危险告知牌', 9: '反光衣', # 11: '呼吸面罩', # 13: '四合一', - 15: '头', + # 15: '头', 16: '安全告知牌', - 18: '安全帽', + # 18: '安全帽', 20: '安全标识牌', - 24: '工服', + # 24: '工服', 34: '灭火器', 43: '警戒线', 48: '路锥', @@ -204,6 +247,12 @@ self.PERSON_CLASS_IDX = 3 self.HEAD_CLASS_IDX = 15 + self.safe_model = YOLO('weights/yinhuan.pt') + self.safe_model_classes = {0: '人', 1: '头', 2: '安全帽', 3: '工服', 4: '烟头', 5: '电话', 6: '袖标'} + self.PERSON_CLASS_IDX = 0 + self.HEAD_CLASS_IDX = 1 + self.SAFETY_CLASS_IDX = [2,3,6] + self.vid_stride = 3 self.stream_loader = OpenCVStreamLoad(camera_url=device.input_stream_url, camera_code=device.code, @@ -350,20 +399,33 @@ # todo 需要生成报警记录吗 def model_predict(self, frames): + result_boxes = [] + safe_result_boxes = [] + results_generator = self.model.track(frames, save_txt=False, save=False, verbose=False, conf=0.5, classes=list(self.model_classes.keys()), imgsz=640, stream=True) - result_boxes = [] + for r in results_generator: result_boxes.append(r.boxes) - return result_boxes + safe_results_generator = self.safe_model.track(frames,save_txt=False, save=False, verbose=False, conf=0.5, + classes=list(self.safe_model_classes.keys()), + imgsz=640, + stream=True) + for s in safe_results_generator: + safe_result_boxes.append(s.boxes) - def handle_behave_alarm(self, frames, result_boxes): + return result_boxes, safe_result_boxes + + def handle_behave_alarm(self, frames, result_boxes, safe_result_boxes): behave_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 0] for alarm_dict in behave_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + use_safe_model = alarm_dict.get('model_type', 'labor') == 'safe' + boxes = safe_result_boxes if use_safe_model else result_boxes + model_classes = self.safe_model_classes if use_safe_model else self.model_classes + for idx, frame_boxes in enumerate(boxes): frame = frames[idx] object_boxes = [box for box in frame_boxes if int(box.cls) in alarm_dict['class_idx']] if alarm_dict['handelType'] == 0: # 检测到就报警 @@ -371,12 +433,20 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 红色标注目标设备，蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: - box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + box_color = COLOR_RED if int(box.cls) in alarm_dict['class_idx'] else COLOR_BLUE + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=box_color, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", - color=box_color, + f"{self.safe_model_classes[box_cls]}", + color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) @@ -391,10 +461,18 @@ self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) + # 蓝色标注正常施工设备 for box in frame_boxes: - if int(box.cls) != self.HEAD_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注人和正常穿戴设备 + for box in safe_result_boxes[idx]: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX or box_cls == self.PERSON_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{self.safe_model_classes[box_cls]}", color=COLOR_BLUE, rotated=False) self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, @@ -431,6 +509,7 @@ has_alarm = True if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人 annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) # 已报警，清零，重新计数 person_status[alarm_dict['alarm_name']] = 0 @@ -438,20 +517,85 @@ if has_alarm: self.alarm_message_center.add_message(alarm_dict) if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 for box in frame_boxes: - box_cls = box.cls - if box_cls != self.PERSON_CLASS_IDX and box_cls != self.HEAD_CLASS_IDX: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: annotator.box_label(box.xyxy.cpu().squeeze(), - f"{self.model_classes[int(box.cls)]}", + f"{model_classes[int(box.cls)]}", color=COLOR_BLUE, rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, alarm_np_img=annotator.result()) + elif alarm_dict['handelType'] == 4: # 人检测到报警：吸烟、打电话 + person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] + has_alarm = False + annotator = None + for person_box in person_boxes: + if person_box.id is None: + continue - def handle_break_in_alarm(self, frames, result_boxes): + person_bbox = person_box.xyxy.cpu().squeeze() + person_id = int(person_box.id) + person_object_box = max( + (box for box in object_boxes if is_overlapping(person_bbox, box.xyxy.cpu().squeeze())), + key=lambda box: box.conf.item(), + default=None + ) + has_object = person_object_box is not None + + + person_status = self.tracking_status[person_id] + if alarm_dict['alarm_name'] not in person_status: + person_status[alarm_dict['alarm_name']] = 0 + if has_object: + person_status[alarm_dict['alarm_name']] += 1 + else: + person_status[alarm_dict['alarm_name']] = 0 + + person_alarm = person_status[alarm_dict['alarm_name']] > self.max_missing_frames + if person_alarm: + has_alarm = True + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + annotator = Annotator(deepcopy(frame)) if annotator is None else annotator + # 红色标注人、目标 + annotator.box_label(person_bbox, alarm_dict['label'], color=COLOR_RED, rotated=False) + annotator.box_label(person_object_box.xyxy.cpu().squeeze(),'',color=COLOR_RED,rotated=False) + # 已报警，清零，重新计数 + person_status[alarm_dict['alarm_name']] = 0 + + if has_alarm: + self.alarm_message_center.add_message(alarm_dict) + if self.alarm_record_center.need_alarm(self.device.code, alarm_dict): + # 蓝色标注正常穿戴设备 + for box in frame_boxes: + box_cls = int(box.cls) + if box_cls in self.SAFETY_CLASS_IDX: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + # 蓝色标注正常施工设备 + for box in result_boxes[idx]: + annotator.box_label(box.xyxy.cpu().squeeze(), + f"{self.model_classes[int(box.cls)]}", + color=COLOR_BLUE, + rotated=False) + self.alarm_record_center.upload_alarm_record(self.device.code, alarm_dict, + alarm_np_img=annotator.result()) + + + def handle_break_in_alarm(self, frames, result_boxes, safe_result_boxes): break_in_alarm_dicts = [d for d in ALARM_DICT if d['alarmCategory'] == 3] for alarm_dict in break_in_alarm_dicts: - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): frame = frames[idx] person_boxes = [box for box in frame_boxes if int(box.cls) == self.PERSON_CLASS_IDX] head_boxes = [box for box in frame_boxes if int(box.cls) == self.HEAD_CLASS_IDX] @@ -524,10 +668,10 @@ t1 = time.time() self.device_status_manager.set_status(device_id=self.device.id) - result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame + result_boxes, safe_result_boxes = self.model_predict(frames) # 结果都是二维数组，对应batch中的每个frame t2 = time.time() - for idx, frame_boxes in enumerate(result_boxes): + for idx, frame_boxes in enumerate(safe_result_boxes): current_person_ids = {int(box.id) for box in frame_boxes if box.cls is not None and box.id is not None and int( box.cls) == self.PERSON_CLASS_IDX} @@ -542,12 +686,12 @@ if self.tracking_status[person_id]['disappear_frames'] > self.disappear_threshold: self.tracking_status.pop(person_id) - self.handle_behave_alarm(frames, result_boxes) - self.handle_break_in_alarm(frames, result_boxes) + self.handle_behave_alarm(frames, result_boxes, safe_result_boxes) + self.handle_break_in_alarm(frames, result_boxes, safe_result_boxes) t3 = time.time() - # for person_id in self.tracking_status.keys(): - # print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') + for person_id in self.tracking_status.keys(): + print(f'person_id: {person_id}, status: {self.tracking_status[person_id]}') for idx, frame in enumerate(frames): annotator = Annotator(frame, None, 18, "Arial.ttf", False, example="人") @@ -557,6 +701,11 @@ f"{self.model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", color=colors(int(s_box.cls)), rotated=False) + for s_box in safe_result_boxes[idx]: + annotator.box_label(s_box.xyxy.cpu().squeeze(), + f"{self.safe_model_classes[int(s_box.cls)]} {float(s_box.conf):.2f}", + color=colors(int(s_box.cls)), + rotated=False) self.display_frame_manager.add_frame(self.device.id, annotator.result()) # self.display_frame_manager.add_frame(self.device.id, frames[idx]) diff --git a/services/schedule_job.py b/services/schedule_job.py index a6e5a2f..b30acf8 100644 --- a/services/schedule_job.py +++ b/services/schedule_job.py @@ -10,8 +10,6 @@ from services.alarm_record_service import AlarmRecordService from services.device_frame_service import DeviceFrameService -from apscheduler.schedulers.asyncio import AsyncIOScheduler -from pytz import timezone import re import gzip diff --git a/tcp/tcp_client_connector.py b/tcp/tcp_client_connector.py index 8dbd86c..6e0b776 100644 --- a/tcp/tcp_client_connector.py +++ b/tcp/tcp_client_connector.py @@ -81,28 +81,49 @@ """连接到设备""" while True: try: + if self.is_connected: + logger.info(f"已连接到 {self.ip}:{self.port}") + return + logger.info(f"正在连接到 {self.ip}:{self.port}...") - # 设置连接超时 + + # 使用 asyncio.open_connection 的超时包装 + connect_task = asyncio.open_connection(self.ip, self.port) self.reader, self.writer = await asyncio.wait_for( - asyncio.open_connection(self.ip, self.port), timeout=self.timeout + connect_task, + timeout=self.timeout ) + + # 验证连接是否真正建立 + if self.writer is None or self.writer.is_closing(): + raise ConnectionError("连接未能成功建立") + self.is_connected = True logger.info(f"已连接到 {self.ip}:{self.port}") if not self.tasks_started: - asyncio.create_task(self.process_message_queue()) # Start processing message queue + asyncio.create_task(self.process_message_queue()) asyncio.create_task(self.start_gas_query()) self.tasks_started = True - break + return # 连接成功后直接返回 - except (asyncio.TimeoutError, ConnectionRefusedError, OSError) as e: - logger.error(f"连接到 {self.ip}:{self.port} 失败，错误: {e}") - logger.info(f"{self.reconnect_interval} 秒后将重连到 {self.ip}:{self.port}") + except asyncio.TimeoutError: + logger.error(f"连接超时: {self.ip}:{self.port}") await asyncio.sleep(self.reconnect_interval) + + except (ConnectionRefusedError, OSError) as e: + logger.error(f"连接被拒绝或网络错误: {self.ip}:{self.port}, 错误: {e}") + await asyncio.sleep(self.reconnect_interval) + except Exception as e: - # 给未知异常一个兜底打印，方便排查 - logger.exception(f"连接到 {self.ip}:{self.port} 遇到未知错误: {e}") + logger.exception(f"连接时发生未知错误: {e}") await asyncio.sleep(self.reconnect_interval) + + finally: + if not self.is_connected: + logger.info(f"{self.reconnect_interval} 秒后将重试连接 {self.ip}:{self.port}") + else: + logger.info(f"已连接到 {self.ip}:{self.port}") async def reconnect(self): """处理断线重连""" @@ -138,6 +159,7 @@ except (ConnectionResetError, BrokenPipeError) as e: logger.exception(f"Error during disconnection") self.reader = self.writer = None + self.is_connected = False logger.info(f"Disconnected from {self.ip}:{self.port}") async def start_gas_query(self): diff --git a/tcp/tcp_server.py b/tcp/tcp_server.py index 366949d..57011b5 100644 --- a/tcp/tcp_server.py +++ b/tcp/tcp_server.py @@ -1,4 +1,5 @@ import asyncio +import base64 from datetime import datetime import re import json @@ -8,7 +9,7 @@ from services.global_config import GlobalConfig HOST = '0.0.0.0' -PORT = 12345 +PORT = 9001 harmful_gas_manager = HarmfulGasManager() push_ts_dict = {} @@ -95,14 +96,22 @@ async def data_push(device_code, message): global_config = GlobalConfig() - harmful_push_config = global_config.get_gas_push_config() + harmful_push_config = global_config.get_harmful_gas_push_config() if harmful_push_config and harmful_push_config.push_url: last_ts = push_ts_dict.get(device_code) current_time = datetime.now() # 检查是否需要推送数据 if last_ts is None or (current_time - last_ts).total_seconds() > harmful_push_config.push_interval: - asyncio.create_task(send_request_async(harmful_push_config.push_url, message)) + # 将字符串编码为字节类型 + encoded_bytes = base64.b64encode(message.encode('utf-8')) + # 将字节编码结果转换为字符串 + encoded_string = encoded_bytes.decode('utf-8') + print(f'before encode: {message}') + print(f'after encode: {encoded_string}') + push_message = {"content": encoded_string} + print(f'body: {push_message}') + asyncio.create_task(send_request_async(harmful_push_config.push_url, push_message)) push_ts_dict[device_code] = current_time # 更新推送时间戳 else: print('no harmful push config') @@ -142,8 +151,14 @@ print(f"连接关闭: {client_address}") break - message = data.decode('utf-8') - print(f"收到数据({client_address}): {message}") + try: + message = data.decode('utf-8') + print(f"收到数据({client_address}): {repr(message)}") + await handle_message(message) + except UnicodeDecodeError as e: + print(f"无法解析消息来自 {client_address}: {e}") + # You can either log the raw data or take other actions + print(f"原始数据: {data}") await handle_message(message)