用YOLOv8+边缘计算构建铁路道口违规检测原型

# generate_synthetic_data.py
import cv2
import numpy as np
from PIL import Image, ImageDraw

# 在背景图上随机放置行人/车辆/火车贴图，并生成yolo格式标签
def generate_sample():
    bg = np.zeros((720,1280,3), np.uint8)
    bg[:] = (100,100,100) # 灰色背景，实际可用实景
    # 随机绘制矩形作为“车辆”
    x, y, w, h = np.random.randint(0,800,4)
    cv2.rectangle(bg, (x,y), (x+w,y+h), (0,255,0), -1)
    # 标签文件
    with open('labels/train.txt','a') as f:
        # class_id(0=person,1=car,2=train) cx cy w h (归一化)
        cx = (x+w/2)/1280
        cy = (y+h/2)/720
        nw = w/1280
        nh = h/720
        f.write(f'1 {cx:.4f} {cy:.4f} {nw:.4f} {nh:.4f}\n')
    cv2.imwrite('images/train.jpg', bg)

# data.yaml
train: ./datasets/railway/images/train
val: ./datasets/railway/images/val
nc: 3
names: ['person','car','train']

# train.py
from ultralytics import YOLO

model = YOLO('yolov8n.pt')  # 从预训练模型开始
results = model.train(
    data='data.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    device='cuda',
    lr0=0.01,
    augment=True
)
model.export(format='onnx', imgsz=640)  # 导出为ONNX

# main.py
import cv2
import paho.mqtt.client as mqtt
import numpy as np
from ultralytics import YOLO

# 加载TensorRT模型
model = YOLO('best.engine', task='detect')

class CrossingFSM:
    def __init__(self):
        self.state = 'IDLE'  # IDLE | BARRIER_DOWN | ALERT
        self.barrier_down = False
        self.alert_sent = False

    def process_detections(self, dets, bridge=True):  # 假设bridge()获取栏杆状态
        if self.barrier_down and not self.alert_sent:
            # 只关注person和car（class 0,1）
            intruders = [d for d in dets if d[5] in [0,1] and d[4] > 0.5]
            if len(intruders) > 0:
                self.state = 'ALERT'
                self.alert_sent = True
                return intruders
        # 如果栏杆抬起，reset
        if not self.barrier_down:
            self.state = 'IDLE'
            self.alert_sent = False
        return []

fsm = CrossingFSM()
client = mqtt.Client()
client.connect('localhost', 1883)

cap = cv2.VideoCapture(0)  # 或RTSP地址
while True:
    ret, frame = cap.read()
    if not ret:
        break
    # 模拟栏杆状态（实际可接GPIO或物体检测）
    fsm.barrier_down = check_barrier(frame)  # 你需要实现这个函数
    
    results = model(frame, imgsz=640, device='cuda')
    dets = results[0].boxes.data.cpu().numpy()  # [x1,y1,x2,y2,conf,cls]
    intruders = fsm.process_detections(dets)
    if intruders:
        # 画框并发送告警
        for box in intruders:
            x1,y1,x2,y2 = map(int, box[:4])
            cv2.rectangle(frame, (x1,y1), (x2,y2), (0,0,255), 2)
        cv2.imwrite('alert_snapshot.jpg', frame)
        client.publish('railway/alert', 'Intruder at crossing!')
        # 避免重复告警，状态机已处理
    
    cv2.imshow('Crossing Monitor', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()

railway-crossing-ai/
├── data/
│   ├── synthetic/          # 合成数据生成脚本
│   │   └── generate.py
│   ├── datasets/            # 训练/验证集
│   │   ├── images/
│   │   └── labels/
│   └── data.yaml
├── models/
│   ├── yolov8n.pt
│   ├── best.onnx
│   └── best.engine          # TensorRT引擎 (INT8)
├── src/
│   ├── main.py              # 实时推理主程序
│   ├── fsm.py               # 状态机
│   ├── mqtt_client.py       # MQTT封装
│   └── barrier_sensor.py    # 栏杆状态检测（占位）
├── requirements.txt
├── Dockerfile               # 用于Jetson或x86的容器
└── README.md