backend-duoji-monitor/scripts/box_edge_inspector.py

"""
箱子堆侧面竖线检测工具
用途：检测箱子堆侧面是否有箱子倾斜或突出
原理：SobelX梯度 → 逐行扫描 + 连续性跟踪 → 聚类成竖线 → 分析异常
"""

import cv2
import numpy as np
from pathlib import Path
from collections import Counter


# ============================================================
# 参数配置区 — 在这里调整所有控制参数
# ============================================================

class Config:
    """所有可调参数，集中管理"""

    # ---- 输入/输出 ----
    IMAGE_PATH = r""           # 输入图片路径（必填）
    OUTPUT_PATH = None         # 输出路径，None=自动生成 _inspected 后缀
    SHOW_WINDOW = True         # 是否弹窗显示结果

    # ---- ROI 裁剪区域（在旋转之后应用）----
    # 设为 None 表示不裁剪，使用整张图
    ROI_X1 = None              # 左边界(px)，例如 100
    ROI_Y1 = None              # 上边界(px)，例如 50
    ROI_X2 = None              # 右边界(px)，例如 800
    ROI_Y2 = None              # 下边界(px)，例如 900

    # ---- 旋转（在裁剪之前应用）----
    ROTATE_DEG = 0.0           # 旋转角度(度)，正=逆时针。例如图片歪了3度就填 3.0

    # ---- 边缘检测方向 ----
    # "dark_to_light" = 从右到左，深→浅（SobelX 负值）
    # "light_to_dark" = 从右到左，浅→深（SobelX 正值）
    GRADIENT_DIR = "dark_to_light"

    # ---- 扫描参数 ----
    SCAN_STEP = 20             # 垂直扫描步长(px)，每 N px 扫一行
    BORDER_MARGIN = 30         # 排除左右边缘各 N px（避免边框噪点）
    X_CLUSTER_GAP = 40        # X方向聚类间距(px)，同一组内的最大间距
    MIN_BRIGHTNESS = 25        # 行平均亮度阈值，低于此值的暗行直接跳过(0-255)
    MAX_GROUPS = 1             # 最多保留几个竖线组(通常只关注主边框=1)

    # ---- 梯度参数 ----
    GAUSSIAN_KERNEL = 5        # 高斯模糊核大小(奇数)
    MIN_GRADIENT_FIRST = 10.0  # 第一遍粗定位的最低梯度强度
    MIN_GRADIENT_TRACK = 6.0   # 第二遍跟踪的最低梯度强度
    TRACK_SMOOTH = 0.7         # 跟踪平滑系数(0-1)，越大越稳定但响应慢

    # ---- 异常检测阈值 ----
    PROTRUDE_PX = 15.0          # 突出异常阈值(px)：偏离中线的距离
    PROTRUDE_MIN_CONSEC = 3     # 突出：最少连续 N 行都突出才算有效
    PROTRUDE_PROXIMITY_PX = 10  # 突出：连续突出的各行 X 间距在此范围内才算同一处突出

    TILT_DEG = 3.0             # 局部倾斜异常阈值(度)：相邻两行之间的角度变化
    TILT_MIN_CONSEC = 3        # 倾倒：最少连续 N 段都超过倾斜阈值才算有效
    TILT_SIMILARITY_DEG = 2.0  # 倾倒：连续各段角度最大差异在此范围内才认为"相似"


def _build_gradient_map(gray: np.ndarray, direction: str) -> np.ndarray:
    """
    SobelX 水平梯度图。
    direction:
      "dark_to_light" → 右暗左亮，保留 SobelX 负值
      "light_to_dark" → 右亮左暗，保留 SobelX 正值
    """
    k = Config.GAUSSIAN_KERNEL
    blurred = cv2.GaussianBlur(gray, (k, k), 1.0)
    sobel_x = cv2.Sobel(blurred, cv2.CV_32F, dx=1, dy=0, ksize=3)

    if direction == "dark_to_light":
        grad = np.maximum(-sobel_x, 0)   # 右暗左亮
    else:  # light_to_dark
        grad = np.maximum(sobel_x, 0)    # 右亮左暗
    return grad


def _row_brightness(gray_row: np.ndarray) -> float:
    """返回一行灰度的平均值(0-255)。"""
    return float(np.mean(gray_row))


def extract_edges(gray: np.ndarray,
                  scan_step: int,
                  border_margin: int,
                  x_cluster_gap: int,
                  min_brightness: float,
                  gradient_dir: str,
                  max_groups: int = 1):
    """
    从灰度图逐行扫描，使用连续性跟踪策略。
    
    新增：跳过太暗的行(min_brightness)。
    """
    h, w = gray.shape
    grad = _build_gradient_map(gray, gradient_dir)

    scan_rows = list(range(0, h, scan_step))
    n_scan = len(scan_rows)

    # ---- 第一遍：粗略定位主区域中心 X（跳过暗行） ----
    rough_points = []
    for y in scan_rows:
        # 跳过太暗的行
        if _row_brightness(gray[y, :]) < min_brightness:
            continue
        row_grad = grad[y, :]
        lo = border_margin
        hi = w - border_margin
        if hi <= lo:
            continue
        # 找窗口内最大值位置
        best_val = -1
        best_x = -1
        for i in range(lo, hi):
            if row_grad[i] > best_val:
                best_val = row_grad[i]
                best_x = i
        if best_x >= 0 and best_val >= Config.MIN_GRADIENT_FIRST:
            rough_points.append(best_x)

    if not rough_points:
        return [], []

    main_center = int(np.median(rough_points))

    # ---- 第二遍：连续性跟踪（同样跳过暗行） ----
    all_points = []
    track_x = main_center
    smooth = Config.TRACK_SMOOTH

    for idx, y in enumerate(scan_rows):
        # 跳过太暗的行
        if _row_brightness(gray[y, :]) < min_brightness:
            continue

        row = grad[y, :]

        # 策略A：在跟踪位置附近 ±x_cluster_gap 搜索
        search_lo = max(border_margin, track_x - x_cluster_gap)
        search_hi = min(w - border_margin, track_x + x_cluster_gap)

        best_x = None
        best_val = -1

        for i in range(search_lo, search_hi):
            if row[i] > best_val:
                best_val = row[i]
                best_x = i

        # 窗口内太弱，尝试全局搜索
        if best_x is not None and best_val < Config.MIN_GRADIENT_TRACK:
            glo_lo = border_margin
            glo_hi = w - border_margin
            g_best_val = -1
            g_best_x = None
            for i in range(glo_lo, glo_hi):
                if row[i] > g_best_val:
                    g_best_val = row[i]
                    g_best_x = i
            if g_best_x is not None and g_best_val >= Config.MIN_GRADIENT_TRACK * 0.5:
                best_x = g_best_x
                best_val = g_best_val
            else:
                best_x = None

        if best_x is not None:
            all_points.append({
                'x': int(best_x),
                'y': y,
                'row_idx': idx,
            })
            track_x = int(track_x * smooth + best_x * (1 - smooth))

    if not all_points:
        return [], []

    # ---- 按 X 聚类 ----
    all_points.sort(key=lambda p: p['x'])
    clusters = []
    cluster = [all_points[0]]
    for i in range(1, len(all_points)):
        if all_points[i]['x'] - cluster[-1]['x'] <= x_cluster_gap:
            cluster.append(all_points[i])
        else:
            clusters.append(cluster)
            cluster = [all_points[i]]
    clusters.append(cluster)

    # 过滤小簇
    min_count = max(3, int(n_scan * 0.08))
    clusters = [c for c in clusters if len(c) >= min_count]

    # 只保留最显著的组
    clusters.sort(key=len, reverse=True)
    clusters = clusters[:max_groups]

    # ---- 构建结果 ----
    segments = []
    groups_info = []

    for gid, cl in enumerate(clusters):
        cl_sorted = sorted(cl, key=lambda p: p['y'])
        xs_arr = np.array([p['x'] for p in cl_sorted], dtype=np.float64)
        ys_arr = np.array([p['y'] for p in cl_sorted], dtype=np.float64)

        if len(xs_arr) >= 3:
            coeffs = np.polyfit(ys_arr, xs_arr, 1)
            angle_h = np.degrees(np.arctan(coeffs[0]))
        else:
            angle_h = 0.0
        angle_v = round((90.0 - angle_h) % 180, 2)
        median_x = int(np.median(xs_arr))

        groups_info.append({
            'group_id': gid,
            'median_x': median_x,
            'angle': angle_v,
            'count': len(cl_sorted),
            'y_min': int(ys_arr.min()) if len(ys_arr) else 0,
            'y_max': int(ys_arr.max()) if len(ys_arr) else 0,
        })

        for p in cl_sorted:
            segments.append({
                'x': p['x'],
                'y': p['y'],
                'row_idx': p['row_idx'],
                'group_id': gid,
                'group_angle': angle_v,
            })

    segments.sort(key=lambda s: (s['group_id'], s['y']))
    return segments, groups_info


# ============================================================
# 图像预处理：旋转 + ROI裁剪
# ============================================================

def preprocess_image(img: np.ndarray,
                     rotate_deg: float,
                     roi: tuple | None) -> tuple[np.ndarray, np.ndarray]:
    """
    预处理：旋转 → ROI裁剪 → 返回 (处理后彩色图, 处理后灰度图)
    
    roi = (x1, y1, x2, y2) 或 None
    """
    result = img.copy()

    # ---- 旋转 ----
    if abs(rotate_deg) > 0.001:
        h, w = result.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, rotate_deg, 1.0)
        # 计算新图像尺寸以包含全部内容
        cos_a = abs(M[0, 0])
        sin_a = abs(M[0, 1])
        new_w = int(h * sin_a + w * cos_a)
        new_h = int(h * cos_a + w * sin_a)
        M[0, 2] += (new_w - w) / 2
        M[1, 2] += (new_h - h) / 2
        result = cv2.warpAffine(result, M, (new_w, new_h),
                                flags=cv2.INTER_LINEAR,
                                borderValue=(0, 0, 0))
        print(f"  [预处理] 旋转 {rotate_deg}deg → {new_w}x{new_h}")

    # ---- ROI 裁剪 ----
    if roi is not None:
        x1, y1, x2, y2 = roi
        rh, rw = result.shape[:2]
        x1 = max(0, min(x1, rw))
        y1 = max(0, min(y1, rh))
        x2 = max(0, min(x2, rw))
        y2 = max(0, min(y2, rh))
        if x2 > x1 and y2 > y1:
            result = result[y1:y2, x1:x2]
            print(f"  [预处理] ROI裁剪 ({x1},{y1})-({x2},{y2}) → {result.shape[1]}x{result.shape[0]}")

    gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    return result, gray


# ============================================================
# 异常分析
# ============================================================

def analyze_segments(segments: list,
                     groups_info: list,
                     protrude_px: float,
                     protrude_min_consec: int,
                     protrude_proximity_px: float,
                     tilt_deg: float,
                     tilt_min_consec: int,
                     tilt_similarity_deg: float) -> dict:
    """
    异常分析（连续段判定）：
    
    【突出】同一组内连续 N 行都突出，且相邻突出行 X 间距 < proximity_px → 报异常
    【倾倒】同一组内连续 N 段的局部倾斜都 > tilt_deg，且各段角度相近(< similarity_deg) → 报异常
    """
    if not segments or not groups_info:
        return {'alerts': [], 'stats': {'total_segments': 0, 'x_range': (0,0), 'groups': []}}

    all_x = [s['x'] for s in segments]
    x_low = float(np.percentile(all_x, 5))
    x_high = float(np.percentile(all_x, 95))
    stats = {'total_segments': len(segments), 'x_range': (x_low, x_high), 'groups': groups_info}
    alerts = []

    for ginfo in groups_info:
        gid = ginfo['group_id']
        median_x = float(ginfo['median_x'])
        group_segs = [s for s in segments if s['group_id'] == gid]
        if len(group_segs) < 2:
            continue

        # ============================================
        # 1. 突出检测：连续 N 行突出 + 位置相近
        # ============================================
        # 标记每行是否突出
        protruding_indices = []  # 在 group_segs 中的索引
        for idx, seg in enumerate(group_segs):
            off = abs(seg['x'] - median_x)
            if off > protrude_px:
                protruding_indices.append(idx)

        # 找连续段
        if protruding_indices:
            runs = []
            run_start = 0
            for i in range(1, len(protruding_indices)):
                gap = protruding_indices[i] - protruding_indices[i - 1]
                # 索引连续(gap==1) 且 X 位置相近
                prev_seg = group_segs[protruding_indices[i - 1]]
                curr_seg = group_segs[protruding_indices[i]]
                x_dist = abs(curr_seg['x'] - prev_seg['x'])
                if gap <= 2 and x_dist <= protrude_proximity_px:
                    continue  # 同一段
                else:
                    run_len = i - run_start
                    if run_len >= protrude_min_consec:
                        runs.append(protruding_indices[run_start:i])
                    run_start = i
            # 最后一段
            run_len = len(protruding_indices) - run_start
            if run_len >= protrude_min_consec:
                runs.append(protruding_indices[run_start:])

            for run in runs:
                segs_in_run = [group_segs[i] for i in run]
                offsets = [round(abs(s['x'] - median_x), 1) for s in segs_in_run]
                avg_off = round(np.mean(offsets), 1)
                sides = ['right' if s['x'] > median_x else 'left' for s in segs_in_run]
                # 多数侧
                from collections import Counter
                main_side = Counter(sides).most_common(1)[0][0]

                alerts.append({
                    'type': 'protrude',
                    'group_id': gid,
                    'side': main_side,
                    'y_start': segs_in_run[0]['y'],
                    'y_end': segs_in_run[-1]['y'],
                    'count': len(segs_in_run),
                    'avg_offset_px': avg_off,
                    'segments': segs_in_run,
                })

        # ============================================
        # 2. 倾倒检测：连续 N 段局部倾斜都超标 + 角度相似
        # ============================================
        local_tilts = []  # [(seg_index, tilt_angle), ...]
        for i in range(1, len(group_segs)):
            p, c = group_segs[i - 1], group_segs[i]
            dy = c['y'] - p['y']
            if dy == 0:
                continue
            lt = abs(np.degrees(np.arctan((c['x'] - p['x']) / dy)))
            local_tilts.append((i, round(lt, 2)))

        # 找连续段：每段倾斜都超标 + 角度差异在 similarity 内
        tilted_indices = [t for t in local_tilts if t[1] > tilt_deg]

        if tilted_indices:
            runs = []
            run_start = 0
            for i in range(1, len(tilted_indices)):
                gap = tilted_indices[i][0] - tilted_indices[i - 1][0]
                angle_diff = abs(tilted_indices[i][1] - tilted_indices[i - 1][1])
                if gap <= 2 and angle_diff <= tilt_similarity_deg:
                    continue
                else:
                    run_len = i - run_start
                    if run_len >= tilt_min_consec:
                        runs.append(tilted_indices[run_start:i])
                    run_start = i
            run_len = len(tilted_indices) - run_start
            if run_len >= tilt_min_consec:
                runs.append(tilted_indices[run_start:])

            for run in runs:
                angles = [t[1] for t in run]
                avg_angle = round(np.mean(angles), 2)
                # 涉及的 segments
                start_idx = run[0][0] - 1  # 局部倾斜的起点段索引
                end_idx = run[-1][0]
                involved = group_segs[start_idx:end_idx + 1]

                alerts.append({
                    'type': 'tilt',
                    'group_id': gid,
                    'y_start': involved[0]['y'],
                    'y_end': involved[-1]['y'],
                    'count': len(run),
                    'avg_tilt_deg': avg_angle,
                    'max_tilt_deg': round(max(angles), 2),
                    'segments': involved,
                })

    return {'alerts': alerts, 'stats': stats}


# ============================================================
# 可视化 — 只画竖线，绿正常红异常
# ============================================================

def draw_results(img: np.ndarray, segments: list,
                 groups_info: list, alerts: list) -> np.ndarray:
    """
    画竖线：绿色=正常，红色=异常。
    异常按连续段标记（alert 里的 'segments' 字段给出具体行）。
    """
    result = img.copy()
    h, _ = result.shape[:2]

    # 收集所有异常行（x, y）
    alert_rows = set()
    for a in alerts:
        if 'segments' in a:
            for s in a['segments']:
                alert_rows.add((int(s['x']), int(s['y'])))

    for ginfo in groups_info:
        gid = ginfo['group_id']
        for s in segments:
            if s['group_id'] != gid:
                continue
            sx, sy = s['x'], s['y']
            bad = (sx, sy) in alert_rows
            cv2.line(result, (sx, max(0, sy - 10)), (sx, min(h, sy + 10)),
                     (0, 0, 255) if bad else (0, 220, 0), 3 if bad else 2)

    return result


# ============================================================
# 主流程
# ============================================================

def inspect(image_path: str,
            output_path: str = None,
            show: bool = False,
            rotate_deg: float = 0.0,
            roi: tuple | None = None,
            gradient_dir: str = "dark_to_light",
            scan_step: int = 20,
            border_margin: int = 30,
            x_cluster_gap: int = 40,
            min_brightness: float = 25.0,
            max_groups: int = 1,
            protrude_px: float = 15.0,
            protrude_min_consec: int = 3,
            protrude_proximity_px: float = 10.0,
            tilt_deg: float = 3.0,
            tilt_min_consec: int = 3,
            tilt_similarity_deg: float = 2.0) -> dict:

    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"无法读取图片: {image_path}")
    print(f"[输入] {image_path}  原始尺寸: {img.shape[1]}x{img.shape[0]}")

    # ---- 预处理：旋转 + ROI裁剪 ----
    processed_img, gray = preprocess_image(img, rotate_deg, roi)

    # ---- 边缘提取 ----
    segments, groups_info = extract_edges(
        gray,
        scan_step=scan_step,
        border_margin=border_margin,
        x_cluster_gap=x_cluster_gap,
        min_brightness=min_brightness,
        gradient_dir=gradient_dir,
        max_groups=max_groups,
    )

    # ---- 异常分析 ----
    result = analyze_segments(
        segments, groups_info,
        protrude_px=protrude_px,
        protrude_min_consec=protrude_min_consec,
        protrude_proximity_px=protrude_proximity_px,
        tilt_deg=tilt_deg,
        tilt_min_consec=tilt_min_consec,
        tilt_similarity_deg=tilt_similarity_deg,
    )

    # ---- 可视化（在预处理后的图上画） ----
    annotated = draw_results(processed_img, segments, groups_info, result['alerts'])

    # ---- 保存 ----
    if output_path is None:
        p = Path(image_path)
        output_path = str(p.parent / f"{p.stem}_inspected{p.suffix}")
    cv2.imwrite(output_path, annotated)
    print(f"[OK] 结果已保存: {output_path}")

    if show:
        cv2.imshow("Box Edge Inspector", annotated)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    # ---- 报告 ----
    dir_label = "深→浅" if gradient_dir == "dark_to_light" else "浅→深"
    print(f"\n{'='*50}")
    print(f"检测完成: {image_path}")
    print(f"  方向: {dir_label}  |  旋转: {rotate_deg}deg")
    if roi:
        print(f"  ROI: {roi}")
    print(f"  扫描步长: {scan_step}px  |  最小亮度: {min_brightness}")
    print(f"  竖线段数: {result['stats']['total_segments']}")
    print(f"  轮廓组数: {len(result['stats']['groups'])}")
    print(f"  X范围: [{result['stats']['x_range'][0]:.0f}, {result['stats']['x_range'][1]:.0f}]")
    print(f"  突出阈值: {protrude_px}px 连续{protrude_min_consec}行  |  倾倒阈值: {tilt_deg}deg 连续{tilt_min_consec}段")

    if result['alerts']:
        print(f"  ! 发现{len(result['alerts'])}处异常:")
        for i, a in enumerate(result['alerts'], 1):
            t = a['type']
            if t == 'tilt':
                print(f"    {i}. [倾倒] Y={a['y_start']}~{a['y_end']}  "
                      f"{a['count']}段  平均{a['avg_tilt_deg']}deg  最大{a['max_tilt_deg']}deg")
            elif t == 'protrude':
                print(f"    {i}. [突出] Y={a['y_start']}~{a['y_end']}  "
                      f"向{a['side']}{a['avg_offset_px']}px  {a['count']}行")
    else:
        print(f"  OK 无异常")
    print('='*50)

    return {
        'ok': len(result['alerts'])==0,
        'alerts': result['alerts'],
        'stats': result['stats'],
        'output_path': output_path,
    }


# ============================================================
if __name__ == '__main__':

    # ─── 所有参数在这里改 ──────────────────────────────
    C = Config

    inspect(
        image_path=C.IMAGE_PATH,
        output_path=C.OUTPUT_PATH,
        show=C.SHOW_WINDOW,

        # 几何变换
        rotate_deg=C.ROTATE_DEG,
        roi=(C.ROI_X1, C.ROI_Y1, C.ROI_X2, C.ROI_Y2) if C.ROI_X1 is not None else None,

        # 检测方向
        gradient_dir=C.GRADIENT_DIR,     # "dark_to_light" 或 "light_to_dark"

        # 扫描参数
        scan_step=C.SCAN_STEP,           # 每 N px 扫一行
        border_margin=C.BORDER_MARGIN,   # 排除边缘 N px
        x_cluster_gap=C.X_CLUSTER_GAP,   # 组内最大间距
        min_brightness=C.MIN_BRIGHTNESS, # 暗行阈值，低于此跳过
        max_groups=C.MAX_GROUPS,         # 保留几个竖线组

        # 异常阈值
        protrude_px=C.PROTRUDE_PX,
        protrude_min_consec=C.PROTRUDE_MIN_CONSEC,
        protrude_proximity_px=C.PROTRUDE_PROXIMITY_PX,
        tilt_deg=C.TILT_DEG,
        tilt_min_consec=C.TILT_MIN_CONSEC,
        tilt_similarity_deg=C.TILT_SIMILARITY_DEG,
    )