ai_special_exam/welding_reset_detect.py

import cv2
import torch
import threading
from shapely.geometry import box, Polygon

from datetime import datetime
from ultralytics import YOLO

from config import WELDING_MODEL_PATHS,WELDING_VIDEO_SOURCES
from config import WELDING_CH1_RTSP,WELDING_CH2_RTSP,WELDING_CH3_RTSP,WELDING_CH4_RTSP,WELDING_CH5_RTSP
from config import SAVE_IMG_PATH,POST_IMG_PATH2,WELDING_REGION2,WELDING_REGION3
from globals import oil_barrel_flag,main_switch_flag,ground_wire_flag,welding_components_flag,welding_machine_switch_flag
from globals import lock,redis_client,stop_event


def init_rest_detection():
    redis_client.delete("welding_reset_post_path")#删除该列表welding_reset_post_path
    
    redis_client.set("welding_main_switch_save_img",'False')
    redis_client.set("welding_oil_barrel_save_img",'False')
    redis_client.set("welding_ground_wire_save_img",'False')
    redis_client.set("welding_components_save_img",'False')
    redis_client.set("welding_machine_switch_save_img",'False')


def start_reset_detection(start_events):
        # Create threads for each video stream and model
    threads = []
    for model_path, video_source in zip(WELDING_MODEL_PATHS, WELDING_VIDEO_SOURCES):
        event = threading.Event()
        start_events.append(event)
        thread = threading.Thread(target=process_video, args=(model_path, video_source,event))
        threads.append(thread)
        thread.daemon=True
        thread.start()


    # Wait for all threads to complete
    for thread in threads:
        thread.join()
        print("复位检测子线程结束")

# Function to process video with YOLO model
def process_video(model_path, video_source, start_event):
    # Load YOLO model
    model = YOLO(model_path)
    #results = model.predict(video_source,stream=True,verbose=False,conf=0.4,device='0')#这里的results是一个生成器
    cap = cv2.VideoCapture(video_source)
    # Loop through the video frames
    while cap.isOpened():
        if stop_event.is_set():#控制停止推理
            break
        # Read a frame from the video
        success, frame = cap.read()

        if success:
            if cap.get(cv2.CAP_PROP_POS_FRAMES) % 25 != 0:#跳帧检测，
                continue

            if video_source == WELDING_CH2_RTSP or video_source == WELDING_CH4_RTSP:#这两个视频流用的分类模型，因为分类模型预处理较慢，需要手动resize
                frame=cv2.resize(frame,(640,640))

            # Run YOLOv8 inference on the frame
            results = model.predict(frame,verbose=False,conf=0.4)

            global oil_barrel_flag,main_switch_flag,ground_wire_flag,welding_components_flag,welding_machine_switch_flag

    #with lock:
            for r in results:


                #因为搭铁线不好识别，所以当没有检测到时，ground_wire_flag=False
                if video_source == WELDING_CH2_RTSP:
                    # annotated_frame = results[0].plot()
                    # cv2.namedWindow('main_switch', cv2.WINDOW_NORMAL)
                    # cv2.imshow('main_switch', annotated_frame)
                    # if cv2.waitKey(1) & 0xFF == ord("q"):
                    #     break  

                    if r.probs.top1conf>0.6:
                        label=model.names[r.probs.top1]
                        
                        welding_components_flag=True if label == "component" else False
                    else:
                        continue

                if video_source == WELDING_CH4_RTSP:
                    if r.probs.top1conf>0.6:
                        label=model.names[r.probs.top1]#获取最大概率的类别的label
                        
                        main_switch_flag = True if label == "open" else False
                    else:
                        continue

                if video_source == WELDING_CH5_RTSP or video_source == WELDING_CH3_RTSP or video_source == WELDING_CH1_RTSP:
                    ##下面这些都是tensor类型
                    boxes = r.boxes.xyxy  # 提取所有检测到的边界框坐标
                    confidences = r.boxes.conf  # 提取所有检测到的置信度
                    classes = r.boxes.cls  # 提取所有检测到的类别索引

                    if video_source == WELDING_CH5_RTSP:
                        ground_wire_flag=False
                        #welding_components_flag=False
                    if video_source == WELDING_CH3_RTSP:
                        oil_barrel_flag=True

                    for i in range(len(boxes)):
                        x1, y1, x2, y2 = boxes[i].tolist()
                        confidence = confidences[i].item()
                        cls = int(classes[i].item())
                        label = model.names[cls]

                        if label=="dump":#检测油桶
                            x_center = (x1 + x2) / 2
                            y_center = (y1 + y2) / 2
                            center_point = (int(x_center), int(y_center))
                            is_inside = cv2.pointPolygonTest(WELDING_REGION2.reshape((-1, 1, 2)), center_point, False)
                            #print(is_inside)
                            
                            if is_inside>=0 :
                                oil_barrel_flag=False #表示油桶在危险区域
                            else:
                                oil_barrel_flag=True 

                        
                        if label== "open":#检测焊机开关
                            welding_machine_switch_flag = True



                        if label=="close":#检测焊机开关
                            welding_machine_switch_flag = False


                        if label=="grounding_wire" :
                            rect_shapely = box(x1,y1, x2, y2)#使用shapely库创建的矩形
                            WELDING_REGION3_shapely = Polygon(WELDING_REGION3.tolist()) #shapely计算矩形检测框和多边形的iou使用
                            intersection = rect_shapely.intersection(WELDING_REGION3_shapely)
                                    # 计算并集
                            union = rect_shapely.union(WELDING_REGION3_shapely)
                                    # 计算 IoU
                            iou = intersection.area / union.area


                            if iou>0 :
                                ground_wire_flag=True #表示搭铁线连接在焊台上
                            else:
                                ground_wire_flag=False #表示未连接上

                flag_count = sum([oil_barrel_flag, main_switch_flag, ground_wire_flag, welding_components_flag, welding_machine_switch_flag])
                redis_client.set("welding_reset_flag",flag_count)
                #print("主开关",main_switch_flag)

                if video_source == WELDING_CH4_RTSP :#检测到总开关
                        if main_switch_flag and redis_client.get("welding_main_switch_save_img")=='False':
                            print("当前总开关没有复位")##焊前检查只保存一次
                            redis_client.set("welding_main_switch_save_img",'True')
                            #save_time=datetime.now().strftime('%Y%m%d_%H%M')
                            save_time=datetime.now().strftime('%Y%m%d_%H')
                            imgpath = f"{SAVE_IMG_PATH}/welding_resetStep2_{save_time}.jpg"
                            post_path = f"{POST_IMG_PATH2}/welding_resetStep2_{save_time}.jpg"
                            redis_client.rpush("welding_reset_post_path",post_path)#welding_reset_post_path为一个列表，存储需要发送的图片路径，rpush为从右侧加入
                            annotated_frame = results[0].plot()
                            cv2.imwrite(imgpath, annotated_frame)



                                
                if video_source == WELDING_CH3_RTSP :##检测油桶{0: 'dump'},在安全区域时保存图片一张
                    if oil_barrel_flag and redis_client.get("welding_oil_barrel_save_img")=='False':
                        print("当前油桶没有复位")
                        redis_client.set("welding_oil_barrel_save_img",'True')
                        #save_time=datetime.now().strftime('%Y%m%d_%H%M')
                        save_time=datetime.now().strftime('%Y%m%d_%H')
                        imgpath = f"{SAVE_IMG_PATH}/welding_resetStep1_{save_time}.jpg"
                        post_path = f"{POST_IMG_PATH2}/welding_resetStep1_{save_time}.jpg"
                        redis_client.rpush("welding_reset_post_path",post_path)
                        annotated_frame = results[0].plot()
                        cv2.imwrite(imgpath, annotated_frame)
                        #reset_post(welding_resetStep='1',path=post_path)
                        #post(step='2',path=post_path)

                if video_source == WELDING_CH1_RTSP:#检测到焊机开关
                        if welding_machine_switch_flag and redis_client.get("welding_machine_switch_save_img")=='False':
                            print("当前焊机开关没有复位")##焊前检查只保存一次
                            redis_client.set("welding_machine_switch_save_img",'True')
                            #save_time=datetime.now().strftime('%Y%m%d_%H%M')
                            save_time=datetime.now().strftime('%Y%m%d_%H')
                            imgpath = f"{SAVE_IMG_PATH}/welding_resetStep5_{save_time}.jpg"
                            post_path = f"{POST_IMG_PATH2}/welding_resetStep5_{save_time}.jpg"
                            redis_client.rpush("welding_reset_post_path",post_path)
                            #cv2.imwrite(imgpath, annotator.result())
                            annotated_frame = results[0].plot()
                            cv2.imwrite(imgpath, annotated_frame)
                            #reset_post(welding_resetStep='5',path=post_path)
                            #post(step='2',path=post_path)
            

                if video_source == WELDING_CH5_RTSP:
                    if ground_wire_flag and redis_client.get("welding_ground_wire_save_img")=='False':
                        print("搭铁线没有复位")
                        redis_client.set("welding_ground_wire_save_img",'True')
                        #save_time=datetime.now().strftime('%Y%m%d_%H%M')  
                        save_time=datetime.now().strftime('%Y%m%d_%H')
                        imgpath = f"{SAVE_IMG_PATH}/welding_resetStep3_{save_time}.jpg"
                        post_path = f"{POST_IMG_PATH2}/welding_resetStep3_{save_time}.jpg"
                        redis_client.rpush("welding_reset_post_path",post_path)
                        # result_image = annotator.result()
                        # cv2.polylines(result_image, [WELDING_REGION3.reshape(-1, 1, 2)], isClosed=True, color=(0, 255, 0), thickness=4)
                        # cv2.imwrite(imgpath, result_image)
                        annotated_frame = results[0].plot()
                        cv2.imwrite(imgpath, annotated_frame)
                        #reset_post(welding_resetStep="3",path=post_path)
                        #time.sleep(1)
                        #post(step='4',path=post_path)

                if video_source == WELDING_CH2_RTSP:
                    if welding_components_flag and redis_client.get("welding_components_save_img")=='False':
                        print("焊件没有复位")
                        redis_client.set("welding_components_save_img",'True')
                        #save_time=datetime.now().strftime('%Y%m%d_%H%M')  
                        save_time=datetime.now().strftime('%Y%m%d_%H')
                        imgpath = f"{SAVE_IMG_PATH}/welding_resetStep4_{save_time}.jpg"
                        post_path = f"{POST_IMG_PATH2}/welding_resetStep4_{save_time}.jpg"
                        redis_client.rpush("welding_reset_post_path",post_path)
                        # result_image = annotator.result()
                        # cv2.polylines(result_image, [REGION4.reshape(-1, 1, 2)], isClosed=True, color=(0, 255, 0), thickness=4)
                        # cv2.imwrite(imgpath, result_image)
                        annotated_frame = results[0].plot()
                        cv2.imwrite(imgpath, annotated_frame)

                        #reset_post(welding_resetStep='4',path=post_path)
                        #post(step='4',path=post_path)

                #运行到这里表示一个线程检测完毕
                start_event.set()
            


        else:
            # Break the loop if the end of the video is reached
            break

    # Release the video capture object and close the display window
    cap.release()
    # 释放模型资源（如果使用GPU）
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    del model
    #cv2.destroyAllWindows()