license_plate_recognition/detect_plate.py

import argparse
import os
import platform
import re
import threading
import time
import json
from collections import deque, defaultdict
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from threading import Lock

import cv2
import numpy as np
import redis
import torch

PLATFORM = platform.system()
import sys

# ===================== 模块导入 =====================
try:
    from models.experimental import attempt_load
    from modules.audio.speaker import IpCast
    from modules.display.screen import Screen, FlashFile
    from modules.radar.radar import RadarData, DeviceInitData, parse_radar_frame, open_serial
    from plate_recognition.double_plate_split_merge import get_split_merge
    from plate_recognition.plate_rec import (
        allFilePath, cv_imread, get_plate_result, init_model,
    )
    from utils.datasets import letterbox
    from utils.general import check_img_size, non_max_suppression_face, scale_coords
except ImportError as e:
    print(f"❌ 导入模块失败：{e}")
    sys.exit(1)

# ===================== 配置参数 =====================
# 本地省份权重
LOCAL_PROVINCE = "湘"
LOCAL_PROVINCE_WEIGHT = 1.25

# 投票配置
VOTE_WINDOW_SECONDS = 0.8
VOTE_THRESHOLD = 1.2
FAST_OUTPUT_THRESHOLD = 1.8
ENABLE_FAST_OUTPUT = True

# 性能配置
IMG_SIZE = 320
FRAME_SKIP = 3
ENABLE_CLAHE = True
ENABLE_MULTI_SCALE = False

# 【修复】阈值配置 - 降低颜色阈值，兼容绿牌/黄牌
DETECT_THRESH = 0.50  # 从0.55降到0.50
COLOR_THRESH = 0.50  # 从0.65降到0.50（绿牌/黄牌颜色识别通常较低）
REC_THRESH = 0.55  # 从0.60降到0.55
PLATE_ASPECT_RATIO = 1.2

# Redis配置
REDIS_HOST = 'localhost'
REDIS_PORT = 6379
REDIS_DB = 0
REDIS_PASSWORD = None
REDIS_KEY = 'plate_results'
WINDOW_SIZE = 5

# 硬件配置
SCREEN_HOST = '192.168.110.200'
SCREEN_PORT = 5005
RADAR_SCREEN_HOST = '192.168.110.199'
RADAR_SCREEN_PORT = 5005

if PLATFORM == 'Windows':
    RADAR_PORT = 'COM3'
    SPEAKER_PORT = 'COM4'
elif PLATFORM == 'Linux':
    RADAR_PORT = '/dev/ttyACM0'
    SPEAKER_PORT = '/dev/ttyUSB0'
else:
    RADAR_PORT = '/dev/ttyACM0'
    SPEAKER_PORT = '/dev/ttyUSB0'

DEVICE_LOW_SPEED = 15
MAX_RECONNECT_ATTEMPTS = 10
RECONNECT_DELAY = 3
MAX_CONSECUTIVE_FAILURES = 5

output_lock = Lock()
cache_lock = Lock()

# 性能配置
BATCH_REDIS_WRITE = True
REDIS_CLEAN_INTERVAL = 30
ASYNC_REDIS = True
INFERENCE_HALF = True
JIT_COMPILE = False
THREAD_POOL_SIZE = 8

# 【修复】扩展车牌正则 - 兼容绿牌(8位)、黄牌、警牌、使领馆牌等
LICENSE_PLATE_PATTERN = re.compile(
    r'^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼使领A-Z]{1}'  # 省份
    r'[A-Z0-9]{1}'  # 第二位（黄牌可能是数字）
    r'[A-Z0-9DF]{5,7}$'  # 剩余位
)

# 【修复】绿牌专用正则（8位新能源）
GREEN_PLATE_PATTERN = re.compile(
    r'^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼]{1}'
    r'[A-Z]{1}'
    r'[A-Z0-9]{1}'
    r'[DF]{1}'  # 新能源标识
    r'[A-Z0-9]{4}$'
)

# 【修复】黄牌正则（大型车、教练车等）
YELLOW_PLATE_PATTERN = re.compile(
    r'^[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼]{1}'
    r'[A-Z0-9]{1}'
    r'[A-Z0-9]{5}$'
)

# ===================== 全局变量 =====================
redis_client = None
redis_pool = None
redis_lock = Lock()
redis_write_queue = deque(maxlen=100)
executor = ThreadPoolExecutor(max_workers=THREAD_POOL_SIZE)
clean_error_count = 0
redis_read_error = 0
cap = None
radar_serial = None
screen = None
radar_screen = None
speaker = None

plate_vote_cache = defaultdict(list)
last_output_dict = {}
displayed_plates = {}

# 【修复】车牌类型统计
plate_type_stats = {"蓝": 0, "绿": 0, "黄": 0, "其他": 0}


# ===================== 图像增强 =====================
def apply_clahe_fast(img: np.ndarray) -> np.ndarray:
    if img is None or img.size == 0:
        return img

    if len(img.shape) == 3:
        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
        cl = clahe.apply(l)
        enhanced = cv2.merge((cl, a, b))
        return cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
    else:
        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
        return clahe.apply(img)


# ===================== 坐标变换 =====================
def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
    if ratio_pad is None:
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
        pad = ((img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2)
    else:
        gain = ratio_pad[0][0]
        pad = ratio_pad[1]
    coords[:, [0, 2, 4, 6]] -= pad[0]
    coords[:, [1, 3, 5, 7]] -= pad[1]
    coords[:, :8] /= gain
    coords[:, :8] = np.clip(coords[:, :8], 0, [img0_shape[1], img0_shape[0]] * 4)
    return coords


def order_points(pts):
    if isinstance(pts, np.ndarray) and pts.size == 0:
        return np.zeros((4, 2), dtype="float32")
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect


def four_point_transform(image, pts):
    if not isinstance(pts, np.ndarray) or pts.shape != (4, 2) or pts.size == 0:
        return image
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB), 1)
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB), 1)
    dst = np.array([[0, 0], [maxWidth - 1, 0], [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype="float32")
    M = cv2.getPerspectiveTransform(rect, dst)
    return cv2.warpPerspective(image, M, (maxWidth, maxHeight))


# ===================== 【修复】字符校验 - 兼容多种车牌 =====================
def validate_plate_characters(plate_str: str, plate_color: str) -> tuple:
    if not plate_str or len(plate_str) < 5:
        return plate_str, False

    clean = plate_str.strip().upper().replace(' ', '')

    # 【修复】放宽第二位校验（黄牌可能是数字）
    if len(clean) >= 2 and not (clean[1].isalpha() or clean[1].isdigit()):
        return plate_str, False

    corrected = list(clean)
    for i in range(2, len(corrected)):
        if corrected[i] in ['O', 'Q']:
            corrected[i] = '0'
        elif corrected[i] in ['I', 'L']:
            corrected[i] = '1'

    corrected_str = ''.join(corrected)

    # 【修复】多种正则匹配
    if (LICENSE_PLATE_PATTERN.match(corrected_str) or
            GREEN_PLATE_PATTERN.match(corrected_str) or
            YELLOW_PLATE_PATTERN.match(corrected_str)):
        return corrected_str, True

    # 【修复】如果颜色是绿/黄，放宽长度要求
    if "绿" in plate_color and len(corrected_str) == 8:
        return corrected_str, True
    elif "黄" in plate_color and len(corrected_str) in [7, 8]:
        return corrected_str, True

    return plate_str, False


# ===================== 【修复】方向判断 - 兼容绿牌/黄牌 =====================
def is_valid_forward_plate(plate_str, bbox, plate_color=""):
    plate_clean = plate_str.strip().upper().replace(' ', '')

    # 【修复】放宽长度校验
    if len(plate_clean) < 7 or len(plate_clean) > 9:
        return False

    x1, y1, x2, y2 = bbox
    width = x2 - x1
    height = y2 - y1
    if height == 0:
        return False

    aspect_ratio = width / height

    # 【修复】绿牌判断逻辑优化
    is_green = ("绿" in plate_color or
                len(plate_clean) == 8 or
                any(c in plate_clean for c in ['D', 'F']) and len(plate_clean) >= 7)

    # 【修复】黄牌判断逻辑
    is_yellow = "黄" in plate_color

    if is_green:
        # 绿牌宽高比放宽
        if aspect_ratio < 1.3:
            return False
    elif is_yellow:
        # 黄牌通常是大型车，宽高比可能不同
        if aspect_ratio < 1.1:
            return False
    else:
        # 蓝牌
        if aspect_ratio < PLATE_ASPECT_RATIO:
            return False

    # 【修复】放宽正则校验
    if not (LICENSE_PLATE_PATTERN.match(plate_clean) or
            GREEN_PLATE_PATTERN.match(plate_clean) or
            len(plate_clean) in [7, 8, 9]):
        return False

    return True


# ===================== 【修复】加权置信度 - 兼容所有车牌类型 =====================
def calculate_weighted_confidence(plate_str: str, base_rec_probs: list, detect_conf: float,
                                  color_conf: float, plate_color: str = "") -> float:
    if not base_rec_probs:
        return 0.0

    rec_avg = np.mean(base_rec_probs)
    base_score = rec_avg * detect_conf * max(color_conf, 0.4)  # 【修复】降低颜色权重下限

    # 本地省份权重（兼容所有颜色）
    if plate_str.startswith(LOCAL_PROVINCE):
        base_score *= LOCAL_PROVINCE_WEIGHT

    # 【修复】各种车牌类型加分
    clean_len = len(plate_str.replace(' ', ''))

    if "绿" in plate_color or clean_len == 8:
        base_score *= 1.08  # 绿牌加分
    elif "黄" in plate_color:
        base_score *= 1.08  # 黄牌加分
    elif "蓝" in plate_color and clean_len == 7:
        base_score *= 1.05

    return base_score


# ===================== 【修复】车牌颜色识别增强 =====================
def enhance_plate_color_detection(plate_color: str, color_conf: float, plate_no: str) -> tuple:
    """
    增强车牌颜色判断
    当颜色置信度低时，根据车牌号特征推断颜色
    """
    clean = plate_no.strip().upper().replace(' ', '')

    # 置信度太低时，根据特征推断
    if color_conf < 0.5:
        # 8位且含D/F = 新能源绿牌
        if len(clean) == 8 and any(c in clean for c in ['D', 'F']):
            return "绿牌", 0.75
        # 7位且第二位是字母 = 可能是蓝牌
        elif len(clean) == 7 and len(clean) >= 2 and clean[1].isalpha():
            return "蓝牌", 0.65
        # 其他情况 = 黄牌可能性
        elif len(clean) in [7, 8]:
            return "黄牌", 0.60

    # 颜色名称标准化
    color_map = {
        "绿": "绿牌", "绿色": "绿牌", "green": "绿牌",
        "蓝": "蓝牌", "蓝色": "蓝牌", "blue": "蓝牌",
        "黄": "黄牌", "黄色": "黄牌", "yellow": "黄牌",
        "白": "白牌", "白色": "白牌",
        "黑": "黑牌", "黑色": "黑牌"
    }

    for key, value in color_map.items():
        if key in plate_color.lower():
            return value, color_conf

    return plate_color if plate_color else "未知", color_conf


# ===================== 快速识别 =====================
def get_plate_rec_landmark_fast(img, xyxy, conf, landmarks, class_num, device,
                                plate_rec_model, is_color=False):
    x1, y1, x2, y2 = map(int, np.ravel(xyxy))
    landmarks_np = np.array(np.ravel(landmarks)).reshape(4, 2).astype(int)
    rect = [x1, y1, x2, y2]

    roi_img = four_point_transform(img, landmarks_np)
    if int(class_num):
        roi_img = get_split_merge(roi_img)

    roi_resized = cv2.resize(roi_img, (128, 32), interpolation=cv2.INTER_LINEAR)
    if ENABLE_CLAHE:
        roi_enhanced = apply_clahe_fast(roi_resized)
    else:
        roi_enhanced = roi_resized

    if not is_color:
        plate_number, rec_prob = get_plate_result(roi_enhanced, device, plate_rec_model, is_color=False)
        plate_color = ""
        color_conf = 0.0
    else:
        plate_number, rec_prob, plate_color, color_conf = get_plate_result(
            roi_enhanced, device, plate_rec_model, is_color=True)

    if isinstance(rec_prob, np.ndarray):
        rec_prob = rec_prob.tolist()

    # 【修复】增强颜色识别
    plate_color, color_conf = enhance_plate_color_detection(plate_color, color_conf, plate_number)

    corrected_plate, is_valid = validate_plate_characters(plate_number, plate_color)
    if is_valid:
        plate_number = corrected_plate

    # 【修复】传入颜色参数
    final_weight = calculate_weighted_confidence(plate_number, rec_prob, conf,
                                                 color_conf if color_conf else 0.4, plate_color)

    is_forward = is_valid_forward_plate(plate_number, rect, plate_color)

    return {
        "rect": rect,
        "detect_conf": conf,
        "landmarks": landmarks_np.tolist(),
        "plate_no": plate_number,
        "rec_conf_raw": rec_prob,
        "rec_conf_weighted": final_weight,
        "plate_color": plate_color,
        "color_conf": color_conf,
        "plate_type": class_num,
        "is_forward": is_forward
    }


# ===================== 快速投票 =====================
def vote_and_filter_plate_fast(plate_no: str, weight: float, current_time: float,
                               plate_color: str = "") -> tuple:
    clean_plate = plate_no.replace(' ', '').upper()
    if len(clean_plate) < 5:
        return False, None, False

    vote_key = clean_plate[:5]

    with cache_lock:
        cache = plate_vote_cache[vote_key]
        cache[:] = [item for item in cache if current_time - item['time'] <= VOTE_WINDOW_SECONDS]

        new_entry = {'plate': clean_plate, 'weight': weight, 'time': current_time, 'color': plate_color}
        cache.append(new_entry)

        plate_scores = defaultdict(float)
        for item in cache:
            plate_scores[item['plate']] += item['weight']

        if not plate_scores:
            return False, None, False

        best_plate, best_score = max(plate_scores.items(), key=lambda x: x[1])

        # 【修复】绿牌/黄牌降低输出阈值
        threshold = VOTE_THRESHOLD
        if "绿" in plate_color or len(best_plate) == 8:
            threshold *= 0.75  # 绿牌更容易输出
        elif "黄" in plate_color:
            threshold *= 0.80  # 黄牌更容易输出
        elif best_plate.startswith(LOCAL_PROVINCE):
            threshold *= 0.85

        if ENABLE_FAST_OUTPUT and best_score >= FAST_OUTPUT_THRESHOLD:
            if best_plate not in displayed_plates or current_time - displayed_plates[best_plate] > 2.0:
                displayed_plates[best_plate] = current_time
                return True, best_plate, True

        if best_score >= threshold:
            if best_plate not in displayed_plates or current_time - displayed_plates[best_plate] > 2.0:
                displayed_plates[best_plate] = current_time
                return True, best_plate, False

    return False, None, False


def check_duplicate_and_update(plate_no: str, current_time: float) -> bool:
    clean_plate = plate_no.replace(' ', '').upper()

    with output_lock:
        if clean_plate in last_output_dict:
            last_time = last_output_dict[clean_plate]
            cooldown = 1.5 if clean_plate.startswith(LOCAL_PROVINCE) else 2.0
            if current_time - last_time < cooldown:
                return False

        last_output_dict[clean_plate] = current_time
        return True


# ===================== Redis 工具 =====================
def get_current_time():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]


def get_current_timestamp():
    return int(time.time())


def get_redis_client():
    global redis_client, redis_pool
    try:
        if redis_pool is None:
            redis_pool = redis.ConnectionPool(
                host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB,
                password=REDIS_PASSWORD, decode_responses=True,
                max_connections=50, socket_timeout=2,
                socket_connect_timeout=2, retry_on_timeout=True
            )
        if redis_client is None:
            redis_client = redis.Redis(connection_pool=redis_pool)
        redis_client.ping()
        return redis_client
    except Exception as e:
        redis_client = None
        return None


def clean_expired_data_batch():
    client = get_redis_client()
    if not client:
        return 0
    try:
        with redis_lock:
            current_ts = get_current_timestamp()
            cutoff_ts = current_ts - WINDOW_SIZE
            zset_key = f"{REDIS_KEY}:sorted"
            expired = client.zrangebyscore(zset_key, 0, cutoff_ts)
            if expired:
                pipe = client.pipeline(transaction=False)
                pipe.hdel(f"{REDIS_KEY}:data", *expired)
                pipe.zremrangebyscore(zset_key, 0, cutoff_ts)
                pipe.execute()
                return len(expired)
    except Exception as e:
        global clean_error_count
        clean_error_count += 1
    return 0


def save_to_redis_async(plate_no, plate_color, detect_conf, color_conf, rec_avg, direction="incoming"):
    try:
        client = get_redis_client()
        if not client:
            return False, "No Redis"

        timestamp = get_current_timestamp()
        entry = {
            "plate_no": plate_no.strip(),
            "plate_color": plate_color,
            "detect_conf": f"{detect_conf:.3f}",
            "color_conf": f"{color_conf:.3f}",
            "rec_avg": f"{rec_avg:.3f}",
            "timestamp": str(timestamp),
            "datetime": get_current_time(),
            "source": "rtsp_fast",
            "direction": direction
        }

        with redis_lock:
            if BATCH_REDIS_WRITE:
                redis_write_queue.append((timestamp, entry))
                if len(redis_write_queue) >= 10:
                    flush_redis_queue(client)
            else:
                pipe = client.pipeline(transaction=False)
                pipe.hset(f"{REDIS_KEY}:data", timestamp, json.dumps(entry))
                pipe.zadd(f"{REDIS_KEY}:sorted", {timestamp: timestamp})
                pipe.execute()
        return True, "OK"
    except Exception as e:
        return False, str(e)


def flush_redis_queue(client):
    if not redis_write_queue:
        return
    try:
        with redis_lock:
            if not redis_write_queue:
                return
            pipe = client.pipeline(transaction=False)
            for ts, data in redis_write_queue:
                pipe.hset(f"{REDIS_KEY}:data", ts, json.dumps(data))
                pipe.zadd(f"{REDIS_KEY}:sorted", {ts: ts})
            pipe.execute()
            redis_write_queue.clear()
    except Exception as e:
        print(f"Redis批量写入失败：{e}")


# ===================== 初始化函数 =====================
def connect_stream(stream_url, cap_options=""):
    global cap
    attempt = 0
    while attempt < MAX_RECONNECT_ATTEMPTS:
        try:
            print(f"[{get_current_time()}] 尝试连接视频流：{stream_url} (第{attempt + 1}次)")
            cap = cv2.VideoCapture()
            os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = cap_options
            success = cap.open(stream_url, cv2.CAP_FFMPEG)
            if success:
                cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                cap.set(cv2.CAP_PROP_FPS, 30)
                ret, frame = cap.read()
                if ret:
                    print(f"[{get_current_time()}] ✅ 视频流连接成功")
                    return cap, True
                cap.release()
            print(f"[{get_current_time()}] 连接失败，{RECONNECT_DELAY}秒后重试...")
            time.sleep(RECONNECT_DELAY)
            attempt += 1
        except Exception as e:
            print(f"[{get_current_time()}] 连接异常：{str(e)}")
            time.sleep(RECONNECT_DELAY)
            attempt += 1
    return None, False


def init_speaker(port: str):
    attempts = 0
    while attempts < MAX_RECONNECT_ATTEMPTS:
        try:
            sp = IpCast(port=port)
            print(f"✅ 语音模块初始化成功 ({port})")
            return sp
        except Exception as e:
            attempts += 1
            if attempts < MAX_RECONNECT_ATTEMPTS:
                time.sleep(RECONNECT_DELAY)
    return None


def init_screen(name: str, ip: str, port: int):
    sc = Screen(name=name, ip=ip, port=str(port))
    attempts = 0
    while attempts < MAX_RECONNECT_ATTEMPTS:
        if sc.get_live_state():
            print(f"✅ {name} 连接成功 ({ip}:{port})")
            return sc
        time.sleep(RECONNECT_DELAY)
        sc.reconnect()
        attempts += 1
    return None


def init_screen_async(name: str, ip: str, port: int, result_dict: dict):
    screen = init_screen(name, ip, port)
    result_dict[name] = screen


def load_model_optimized(weights, device):
    model = attempt_load(weights, map_location=device)
    if JIT_COMPILE and device.type != 'cpu':
        try:
            dummy = torch.rand(1, 3, 640, 640).to(device)
            if INFERENCE_HALF:
                dummy = dummy.half()
            model = torch.jit.trace(model, dummy)
            print("✅ 模型JIT编译成功")
        except Exception as e:
            print(f"JIT编译失败：{e}")
    if INFERENCE_HALF and device.type != 'cpu':
        model.half()
    model.eval()
    for param in model.parameters():
        param.requires_grad = False
    return model


def detect_Recognition_plate_fast(model, orgimg, device, plate_rec_model, img_size, is_color=False):
    conf_thres = 0.25
    iou_thres = 0.5
    dict_list = []

    h0, w0 = orgimg.shape[:2]
    r = img_size / max(h0, w0)
    if abs(r - 1) > 0.1:
        interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
        img0 = cv2.resize(orgimg, (int(w0 * r), int(h0 * r)), interpolation=interp)
    else:
        img0 = orgimg

    imgsz = check_img_size(img_size, s=model.stride.max())
    img = letterbox(img0, new_shape=imgsz)[0]
    img = img[:, :, ::-1].transpose(2, 0, 1).copy()
    img = torch.from_numpy(img).to(device)
    img = img.float() / 255.0
    if INFERENCE_HALF and device.type != 'cpu':
        img = img.half()
    if img.ndim == 3:
        img = img.unsqueeze(0)

    with torch.no_grad():
        pred = model(img)[0]
        pred = non_max_suppression_face(pred, conf_thres, iou_thres)

    for det in pred:
        if len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], orgimg.shape).round()
            det[:, 5:13] = scale_coords_landmarks(img.shape[2:], det[:, 5:13], orgimg.shape).round()

            for j in range(det.size(0)):
                xyxy = det[j, :4].tolist()
                conf = det[j, 4].cpu().item()
                landmarks = det[j, 5:13].tolist()
                class_num = det[j, 13].cpu().item()

                if conf >= DETECT_THRESH:
                    res = get_plate_rec_landmark_fast(
                        orgimg, xyxy, conf, landmarks, class_num,
                        device, plate_rec_model, is_color
                    )
                    dict_list.append(res)
                    break
        break

    return dict_list


# ===================== 主函数 =====================
def start(image_path="imgs"):
    global screen, radar_screen, speaker, redis_client
    global REDIS_HOST, REDIS_PORT, REDIS_KEY, WINDOW_SIZE

    parser = argparse.ArgumentParser()
    parser.add_argument("--detect_model", nargs="+", type=str, default="weights/plate_detect.pt",
                        help="检测模型路径")
    parser.add_argument("--rec_model", type=str, default="weights/plate_rec_color.pth",
                        help="识别模型路径")
    parser.add_argument("--is_color", type=bool, default=True,
                        help="是否识别车牌颜色")
    parser.add_argument("--img_size", type=int, default=IMG_SIZE,
                        help="推理尺寸")
    parser.add_argument("--stream", type=str, default="",
                        help="RTSP/RTMP流地址")
    parser.add_argument("--redis_host", type=str, default="localhost",
                        help="Redis主机")
    parser.add_argument("--redis_port", type=int, default=6379,
                        help="Redis端口")
    parser.add_argument("--redis_key", type=str, default="plate_results",
                        help="Redis键名")
    parser.add_argument("--window_size", type=int, default=5,
                        help="滑动窗口秒数")
    opt = parser.parse_args()

    REDIS_HOST = opt.redis_host
    REDIS_PORT = opt.redis_port
    REDIS_KEY = opt.redis_key
    WINDOW_SIZE = opt.window_size

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == 'cuda':
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True

    redis_client = get_redis_client()
    if redis_client:
        print("✅ Redis连接成功")

    try:
        detect_model = load_model_optimized(opt.detect_model, device)
        plate_rec_model = init_model(device, opt.rec_model, is_color=opt.is_color)
        print(f"✅ 模型加载成功")
    except Exception as e:
        print(f"❌ 模型加载失败：{e}")
        return

    print(f"\n🚀 【速度优化模式 - 全车牌兼容】")
    print(f"   推理尺寸：{opt.img_size} | 帧跳过：{FRAME_SKIP} | 线程数：{THREAD_POOL_SIZE}")
    print(f"   本地省份：{LOCAL_PROVINCE} (权重x{LOCAL_PROVINCE_WEIGHT})")
    print(f"   阈值：检测={DETECT_THRESH} 颜色={COLOR_THRESH} 识别={REC_THRESH}")
    print(f"   支持：蓝牌✅ 绿牌✅ 黄牌✅ 其他✅\n")

    speaker = init_speaker(SPEAKER_PORT)

    screen_init_results = {}
    screen_threads = [
        threading.Thread(target=init_screen_async, args=("主屏幕", SCREEN_HOST, SCREEN_PORT, screen_init_results),
                         daemon=True),
        threading.Thread(target=init_screen_async,
                         args=("雷达屏幕", RADAR_SCREEN_HOST, RADAR_SCREEN_PORT, screen_init_results), daemon=True)
    ]
    for t in screen_threads:
        t.start()
    for t in screen_threads:
        t.join(timeout=30)

    screen = screen_init_results.get("主屏幕")
    radar_screen = screen_init_results.get("雷达屏幕")

    if radar_screen is not None:
        if not hasattr(radar_screen, 'dyna_area_num'):
            radar_screen.dyna_area_num = 0

    DeviceInitData.LowSpeed = DEVICE_LOW_SPEED
    try:
        threading.Thread(target=open_serial, args=(RADAR_PORT, speaker, radar_screen), daemon=True).start()
        print("✅ 雷达线程启动")
    except Exception as e:
        print(f"⚠️ 雷达启动失败：{e}")

    if opt.stream:
        cap_options = "rtsp_transport=tcp;buffer_size=32000;probesize=32;analyzeduration=0;fflags=nobuffer;flags=low_delay"
        cap, connected = connect_stream(opt.stream, cap_options)
        if not connected:
            print(f"[{get_current_time()}] 初始连接失败，退出程序")
            return

        consecutive_failures = 0
        reconnect_count = 0
        frame_count = 0
        processed_count = 0
        last_print_time = time.time()
        inference_times = deque(maxlen=30)
        last_output_dict.clear()
        plate_vote_cache.clear()
        displayed_plates.clear()
        plate_type_stats.clear()
        plate_type_stats.update({"蓝": 0, "绿": 0, "黄": 0, "其他": 0})

        stats = {"in": 0, "out": 0, "total": 0, "fast_output": 0, "vote_output": 0}

        try:
            while True:
                frame_count += 1
                ret, frame = cap.read()

                if not ret:
                    consecutive_failures += 1
                    if consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
                        cap, reconnected = connect_stream(opt.stream, cap_options)
                        if reconnected:
                            reconnect_count += 1
                            consecutive_failures = 0
                            frame_count = 0
                        else:
                            break
                    continue

                consecutive_failures = 0

                if frame_count % FRAME_SKIP != 0:
                    continue

                processed_count += 1
                inference_start = time.time()

                try:
                    dict_list = detect_Recognition_plate_fast(
                        detect_model, frame, device, plate_rec_model, opt.img_size, is_color=opt.is_color
                    )
                    inference_times.append(time.time() - inference_start)
                    current_time = time.time()

                    for res in dict_list:
                        plate_no = res['plate_no'].strip()
                        if len(plate_no) < 4:
                            continue

                        weight = res['rec_conf_weighted']
                        detect_conf = res['detect_conf']
                        color_conf = res.get('color_conf', 0.0)
                        plate_color = res.get('plate_color', '未知')

                        # 【修复】降低阈值，兼容绿牌/黄牌
                        if weight < 0.25:
                            continue

                        should_output, final_plate, is_fast = vote_and_filter_plate_fast(
                            plate_no, weight, current_time, plate_color)

                        if should_output and final_plate:
                            if is_fast:
                                stats["fast_output"] += 1
                            else:
                                stats["vote_output"] += 1

                            if check_duplicate_and_update(final_plate, current_time):
                                is_in = res.get('is_forward', False)
                                direction = "incoming" if is_in else "outgoing"
                                if is_in:
                                    stats["in"] += 1
                                else:
                                    stats["out"] += 1
                                stats["total"] += 1

                                # 【修复】统计车牌类型
                                if "绿" in plate_color or len(final_plate) == 8:
                                    plate_type_stats["绿"] += 1
                                    color_display = "🟢绿牌"
                                elif "黄" in plate_color:
                                    plate_type_stats["黄"] += 1
                                    color_display = "🟡黄牌"
                                elif "蓝" in plate_color:
                                    plate_type_stats["蓝"] += 1
                                    color_display = "🔵蓝牌"
                                else:
                                    plate_type_stats["其他"] += 1
                                    color_display = f"⚪{plate_color}"

                                print(f"[{get_current_time()}] {'⚡' if is_fast else '✅'} {final_plate} | "
                                      f"{color_display} | W:{weight:.2f} | {'IN' if is_in else 'OUT'}")

                                if redis_client:
                                    executor.submit(save_to_redis_async, final_plate, plate_color,
                                                    detect_conf, color_conf, weight / LOCAL_PROVINCE_WEIGHT, direction)

                                if screen:
                                    try:
                                        ff = FlashFile()
                                        ff.set_msg(final_plate, 1)
                                        ff.set_mode(4, 1)
                                        ff.set_origin(0, True, 0)
                                        ff.set_area(128, True, 32)
                                        screen.text_ram(ff, True)
                                    except Exception as e:
                                        pass

                except Exception as e:
                    continue

                if frame_count % REDIS_CLEAN_INTERVAL == 0 and redis_client:
                    executor.submit(clean_expired_data_batch)

                if frame_count % 15 == 0 and redis_client:
                    executor.submit(flush_redis_queue, get_redis_client())

                if time.time() - last_print_time >= 15:
                    avg_inference = sum(inference_times) / len(inference_times) if inference_times else 0
                    fps = 1 / avg_inference if avg_inference > 0 else 0
                    print(f"\n--- 状态 --- 帧:{frame_count} | 识别:{stats['total']} | "
                          f"蓝:{plate_type_stats['蓝']} 绿:{plate_type_stats['绿']} 黄:{plate_type_stats['黄']} 其他:{plate_type_stats['其他']} | "
                          f"耗时:{avg_inference * 1000:.1f}ms | FPS:{fps:.1f}")
                    last_print_time = time.time()

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

        except KeyboardInterrupt:
            print(f"\n[{get_current_time()}] 用户中断")
        finally:
            if cap:
                cap.release()
            cv2.destroyAllWindows()
            executor.shutdown(wait=True)
            if redis_client:
                flush_redis_queue(get_redis_client())
                clean_expired_data_batch()
            print(f"\n结束。总识别：{stats['total']}")
            print(
                f"车牌类型：蓝:{plate_type_stats['蓝']} 绿:{plate_type_stats['绿']} 黄:{plate_type_stats['黄']} 其他:{plate_type_stats['其他']}")


if __name__ == '__main__':
    start()