import cv2
from ultralytics import YOLO
from opcua import Client
import dmPython
import pandas as pd
from datetime import datetime, timedelta

# pip install openpyxl -i https://pypi.tuna.tsinghua.edu.cn/simple
# pip install cryptography -i https://pypi.tuna.tsinghua.edu.cn/simple

class ConveyorBeltMonitor:
    def __init__(self, model_path, plc_config, cameras_config):
        # 初始化 YOLO model
        self.model = YOLO(model_path)

        # PLC 连接 设置
        self.plc_url = plc_config['url']
        self.plc_username = plc_config['username']
        self.plc_password = plc_config['password']
        self.client = None

        # 摄像头 plc node 设置
        self.cameras = {}

        # 每个摄像头状态
        self.camera_states = {}
        self.skip_frames = 20

        self.setup_cameras(cameras_config)

    def setup_cameras(self, cameras_config):
        """Initialize multiple camera connections with their PLC nodes"""
        for camera_id, config in cameras_config.items():
            cap = cv2.VideoCapture(config['rtsp_url'])
            if not cap.isOpened():
                print(f"Failed to open camera {camera_id}")
                continue
            self.cameras[camera_id] = {
                'capture': cap,
                'frame_count': 0,
                'plc_nodes': config['plc_nodes']  # Now accepts multiple PLC nodes
            }
            self.camera_states[camera_id] = {
                'last_plc_values': {node: None for node in config['plc_nodes']},  # Track last value for each node
                'has_queried_database': False
            }

    def connect_plc(self):
        """建立 PLC 连接"""
        self.client = Client(self.plc_url)
        self.client.set_user(self.plc_username)
        self.client.set_password(self.plc_password)
        self.client.connect()
        print("Connected to OPC UA server")

    def fetch_alarm_data(self):
        """查询 达梦 报警信息"""
        db_config = {
            'host': "10.248.56.208",
            'port': "5236",
            'user': "XMSW",
            'password': "7m*d7Zf04#G%4je"
        }

        connection = dmPython.Connection(
            user=db_config['user'],
            password=db_config['password'],
            server=db_config['host'],
            port=db_config['port']
        )

        cursor = connection.cursor()
        now = datetime.now()
        two_minutes_ago = now - timedelta(minutes=2)

        query = f"""
            SELECT TAG_NAME, TAG_COMMENT, ALARM_TIME, EVENT_TIME, ALARM_TEXT
            FROM XMSW.T_ALARM
            WHERE ALARM_TIME BETWEEN '{two_minutes_ago}' AND '{now}'
            ORDER BY ALARM_TIME DESC
        """

        cursor.execute(query)
        rows = cursor.fetchall()
        columns = ['TAG_NAME', 'TAG_COMMENT', 'ALARM_TIME', 'EVENT_TIME', 'ALARM_TEXT']
        df = pd.DataFrame(rows, columns=columns)

        cursor.close()
        connection.close()
        return df

    def process_frame(self, frame):
        """帧 缩放"""
        dim = (int(frame.shape[1] * 0.5), int(frame.shape[0] * 0.5))
        resized_frame = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)

        results = self.model(resized_frame)
        annotated_frame = results[0].plot()
        detections = results[0].boxes

        return annotated_frame, detections

    def check_plc_conditions(self, camera_id):
        """Check PLC conditions for a specific camera"""
        plc_nodes = self.cameras[camera_id]['plc_nodes']
        current_values = {}
        trigger_condition = False

        # 获取 plc node 值
        for node_id in plc_nodes:
            node = self.client.get_node(node_id)
            current_values[node_id] = node.get_value()
            print(f"Camera {camera_id}: PLC node {node_id} value: {current_values[node_id]}")

        last_values = self.camera_states[camera_id]['last_plc_values']

        # 判断 运行状态
        trigger_condition = any(
            last_values[node_id] == 3 and current_values[node_id] != 3
            for node_id in plc_nodes
        )

        # Update last values
        self.camera_states[camera_id]['last_plc_values'] = current_values

        return trigger_condition

    def handle_detections(self, camera_id, detections):
        """Handle detection results and PLC communication for specific camera"""
        if detections is not None and detections.shape[0] > 0:
            print(f"Camera {camera_id}: Objects detected: {len(detections)}")

            trigger_condition = self.check_plc_conditions(camera_id)
            has_queried_database = self.camera_states[camera_id]['has_queried_database']

            if trigger_condition and not has_queried_database:
                current_time = datetime.now().strftime("%Y-%m-%d %H-%M-%S")
                alarm_data = self.fetch_alarm_data()
                excel_file_name = f"胶带机_{camera_id}_{current_time}.xlsx"
                alarm_data.to_excel(excel_file_name, index=False, sheet_name='QueryResults', engine='openpyxl')
                print(f"Camera {camera_id}: Alarm data written to {excel_file_name}")
                self.camera_states[camera_id]['has_queried_database'] = True
            elif not trigger_condition:
                self.camera_states[camera_id]['has_queried_database'] = False
        else:
            print(f"Camera {camera_id}: No objects detected")

    def run(self):
        """Main monitoring loop"""
        try:
            self.connect_plc()

            while True:
                for camera_id, camera_info in self.cameras.items():
                    cap = camera_info['capture']
                    success, frame = cap.read()

                    if not success:
                        print(f"Failed to read frame from camera {camera_id}")
                        continue

                    camera_info['frame_count'] += 1
                    if camera_info['frame_count'] % self.skip_frames != 0:
                        continue

                    annotated_frame, detections = self.process_frame(frame)
                    self.handle_detections(camera_id, detections)

                    cv2.imshow(f"Camera {camera_id}", annotated_frame)

                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        finally:
            self.cleanup()

    def cleanup(self):
        """释放资源"""
        if self.client:
            self.client.disconnect()
        for camera_info in self.cameras.values():
            camera_info['capture'].release()
        cv2.destroyAllWindows()


# Configuration
plc_config = {
    'url': "opc.tcp://10.248.56.206:37801",
    'username': "swxmc",
    'password': "swxmc@147!"
}

cameras_config = {
    'cam218': {
        'rtsp_url': 'rtsp://admin:hik12345+@10.248.53.20',
        'plc_nodes': ["ns=2;s=XMSW_Raw_CON218_RunState"]  # 单 PLC node
    },
    'cam201': {
        'rtsp_url': 'rtsp://admin:hik12345+@10.248.53.96',
        'plc_nodes': [
            "ns=2;s=XMSW_Raw_CON201_RunState1",
            "ns=2;s=XMSW_Raw_CON201_RunState2"
        ]  # 多 PLC nodes
    }
}

# Create and run monitor
monitor = ConveyorBeltMonitor("best.pt", plc_config, cameras_config)
monitor.run()


上述代码有时候获取plc值报TimeoutError，如何解决plc连接超时或读取值超时问题