棋子检测方案升级：YOLO 目标检测替代传统 ROI 阈值

path: ./dataset
train: images/train
val: images/val

nc: 2
names:
  0: black
  1: white

pip install labelImg
labelImg  # 启动 GUI

#!/usr/bin/env python3
# auto_label_pipeline.py
"""
半自动标注流水线：
1. 用现有 ROI 检测器跑一遍 → 生成 YOLO 标注
2. 低置信度样本标记为需人工复核
3. 人工微调后直接可用
"""

import cv2
import os
from pathlib import Path

def generate_yolo_labels(board_image, grid_points, stone_state, 
                         conf_threshold=0.8):
    """
    board_image: 矫正后的棋盘图
    grid_points: 361 交点坐标 (由 BoardRectifier 输出)
    stone_state: ROI 检测器输出的局面字符串
    conf_threshold: 低于此值的 ROI 特征记为 low_confidence
    
    返回: (labels_list, low_confidence_indices)
    """
    h, w = board_image.shape[:2]
    cell_size = w / 19
    
    labels = []
    low_conf = []
    
    for i, (pt, stone) in enumerate(zip(grid_points, stone_state)):
        if stone == '.':
            continue
        
        class_id = 0 if stone == 'B' else 1
        radius_px = cell_size * 0.35
        
        # 归一化
        cx = pt[0] / w
        cy = pt[1] / h
        bw = (2 * radius_px) / w
        bh = (2 * radius_px) / h
        
        labels.append(f"{class_id} {cx:.6f} {cy:.6f} {bw:.6f} {bh:.6f}")
        
        # 检查 ROI 置信度（需从 StoneDetector 获取）
        # 这里简化：周围 8 邻域有冲突标记则低置信
        pass
    
    return labels, low_conf


def batch_auto_label(image_dir, grid_file, stone_state_file, output_dir):
    """
    批量处理所有矫正后的棋盘图
    grid_file: 每行一张图的 361 点坐标 (JSON/CSV)
    stone_state_file: ROI 检测器输出
    """
    import json
    
    with open(grid_file) as f:
        grids = json.load(f)  # {filename: [[x,y], ...]}
    with open(stone_state_file) as f:
        states = json.load(f)  # {filename: "BWB...W."}
    
    os.makedirs(output_dir, exist_ok=True)
    review_list = []
    
    for fname, grid in grids.items():
        img_path = os.path.join(image_dir, fname)
        if not os.path.exists(img_path):
            continue
        
        img = cv2.imread(img_path)
        state = states.get(fname, '.' * 361)
        
        labels, low_conf = generate_yolo_labels(img, grid, state)
        
        # 写入 YOLO 标注
        label_path = os.path.join(output_dir, 
                                   Path(fname).stem + '.txt')
        with open(label_path, 'w') as f:
            f.write('\n'.join(labels))
        
        if low_conf:
            review_list.append(fname)
    
    # 输出需人工复核的图片列表
    print(f"Auto-labeled: {len(grids)} images")
    print(f"Needs review: {len(review_list)} images")
    with open('review_list.txt', 'w') as f:
        f.write('\n'.join(review_list))


if __name__ == '__main__':
    batch_auto_label(
        image_dir='dataset/images/train/',
        grid_file='grid_coords.json',
        stone_state_file='roi_states.json',
        output_dir='dataset/labels/train/'
    )

# 仅打开需复核的图片
labelImg dataset/images/train/ --labels dataset/labels/train/ \
  --review-list review_list.txt

# validate_labels.py — 检查标注格式与一致性
import os

def validate_dataset(labels_dir, images_dir):
    errors = []
    for label_file in os.listdir(labels_dir):
        if not label_file.endswith('.txt'):
            continue
        
        name = label_file.replace('.txt', '')
        img_file = name + '.jpg'
        if img_file not in os.listdir(images_dir):
            errors.append(f"Missing image: {img_file}")
            continue
        
        with open(os.path.join(labels_dir, label_file)) as f:
            lines = f.readlines()
        
        for i, line in enumerate(lines):
            parts = line.strip().split()
            if len(parts) != 5:
                errors.append(f"{label_file}:{i} bad format")
                continue
            
            cls_id = int(parts[0])
            if cls_id not in (0, 1):
                errors.append(f"{label_file}:{i} bad class {cls_id}")
            
            cx, cy, bw, bh = map(float, parts[1:5])
            if not (0 <= cx <= 1 and 0 <= cy <= 1 and 
                    0 < bw <= 1 and 0 < bh <= 1):
                errors.append(
                    f"{label_file}:{i} out of bounds: {cx:.3f},{cy:.3f}"
                )
    
    if errors:
        print(f"⚠ Found {len(errors)} errors:")
        for e in errors[:20]:  # 最多显示 20 个
            print(f"  {e}")
    else:
        print("✅ All labels valid")

validate_dataset('dataset/labels/train/', 'dataset/images/train/')

# Ultralytics 内置增强，在 data.yaml 或训练参数中配置
augmentation:
  hsv_h: 0.015      # 色调抖动（棋子颜色不变，不需要大 H 变化）
  hsv_s: 0.3        # 饱和度
  hsv_v: 0.3        # 亮度（关键！模拟光照变化）
  degrees: 5.0      # 旋转 ±5°
  translate: 0.05   # 平移
  scale: 0.2        # 缩放（模拟棋子到摄像头距离变化）
  shear: 2.0        # 错切
  flipud: 0.0       # 不上下翻转（棋盘有方向）
  fliplr: 0.5       # 左右翻转
  mosaic: 0.0       # Mosaic 对棋子检测意义不大，关闭
  erasing: 0.1      # 随机擦除（模拟遮挡）

# 创建虚拟环境
python3 -m venv yolo_env
source yolo_env/bin/activate

# 安装依赖
pip install ultralytics torch torchvision onnx onnxruntime
pip install opencv-python matplotlib seaborn

#!/usr/bin/env python3
# train_stone_detector.py
"""
围棋棋子 YOLO 检测器训练脚本
"""

from ultralytics import YOLO
import torch

def main():
    # 检查设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"[INFO] Training on: {device}")

    # ---- 方式 1: 从头训练 ----
    model = YOLO('yolov8n.pt')  # 预训练权重作为起点

    results = model.train(
        data='dataset/data.yaml',
        epochs=100,
        imgsz=640,
        batch=16 if device == 'cuda' else 8,
        device=device,
        workers=4,
        
        # 优化器
        optimizer='AdamW',
        lr0=0.001,
        lrf=0.01,
        momentum=0.937,
        weight_decay=0.0005,
        
        # 数据增强（围棋棋子场景定制）
        hsv_h=0.0,       # 棋子颜色固定（黑/白），不调色调
        hsv_s=0.2,       # 轻微饱和度变化
        hsv_v=0.4,       # 亮度变化加大（模拟光照）
        degrees=5.0,     # 小角度旋转
        translate=0.05,
        scale=0.3,       # 缩放范围
        flipud=0.0,      # 不上下翻转
        fliplr=0.5,
        mosaic=0.0,      # 围棋棋盘是规则网格，Mosaic 破坏空间结构
        erasing=0.1,
        
        # 训练策略
        warmup_epochs=3,
        cos_lr=True,
        close_mosaic=0,  # mosaic=0 不需要这个
        
        # 保存
        save=True,
        save_period=10,
        project='runs/stone_detector',
        name='yolov8n_stone',
        exist_ok=True,
        
        # 验证
        val=True,
        
        # 早停
        patience=20,
    )

    print(f"[INFO] Best model saved at: {results.save_dir}")
    
    # 验证最佳模型
    metrics = model.val()
    print(f"[INFO] Validation metrics:")
    print(f"  mAP50:  {metrics.box.map50:.4f}")
    print(f"  mAP50-95: {metrics.box.map:.4f}")
    print(f"  Precision: {metrics.box.p[0]:.4f}")
    print(f"  Recall: {metrics.box.r[0]:.4f}")


if __name__ == '__main__':
    main()

# 启动训练 + TensorBoard
tensorboard --logdir runs/stone_detector --port 6006

# eval_model.py
from ultralytics import YOLO

model = YOLO('runs/stone_detector/yolov8n_stone/weights/best.pt')

# 在验证集上评估
results = model.val()

# 在单张图上预测并可视化
results = model.predict(
    source='test_images/',
    save=True,
    conf=0.25,       # 置信度阈值
    iou=0.45,        # NMS IoU 阈值
    show_labels=True,
    show_conf=True,
)

# 导出混淆矩阵
results = model.val(plots=True)

# export_to_onnx.py
from ultralytics import YOLO

model = YOLO('runs/stone_detector/yolov8n_stone/weights/best.pt')

# 导出 ONNX（静态输入尺寸）
model.export(
    format='onnx',
    imgsz=640,
    opset=12,           # ONNX opset 版本，兼容 OpenCV 4.x
    simplify=True,       # onnx-simplifier 简化计算图
    dynamic=False,       # 静态尺寸，推理更快
    half=False,          # FP32，CPU 推理
)

print("[INFO] Exported to: runs/stone_detector/yolov8n_stone/weights/best.onnx")

# 可选：导出 TensorRT（如果有 Jetson 或 GPU 部署）
# model.export(format='engine', imgsz=640, half=True, device=0)

# Ubuntu/Debian
sudo apt install libonnxruntime-dev

# 或从源码/Microsoft 下载预编译包
# https://github.com/microsoft/onnxruntime/releases
wget https://github.com/microsoft/onnxruntime/releases/download/v1.18.0/onnxruntime-linux-x64-1.18.0.tgz
tar xzf onnxruntime-linux-x64-1.18.0.tgz -C /opt/

# ONNX Runtime
set(ONNXRUNTIME_ROOT /opt/onnxruntime-linux-x64-1.18.0)
find_library(ONNXRUNTIME_LIB onnxruntime
    PATHS ${ONNXRUNTIME_ROOT}/lib
    NO_DEFAULT_PATH)
add_library(onnxruntime SHARED IMPORTED)
set_target_properties(onnxruntime PROPERTIES
    IMPORTED_LOCATION ${ONNXRUNTIME_LIB}
    INTERFACE_INCLUDE_DIRECTORIES ${ONNXRUNTIME_ROOT}/include)

// detect/YoloStoneDetector.h
#pragma once

#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <string>

/// YOLO 检测到的单个棋子
struct YoloDetection {
    int classId;          // 0 = 黑子, 1 = 白子
    float confidence;     // 置信度 [0, 1]
    cv::Rect2f box;       // 检测框 (左上角, 宽, 高)
    cv::Point2f center;   // 框中心 = 棋子中心
};

/// YOLO ONNX 棋子检测器
class YoloStoneDetector {
public:
    struct Config {
        std::string modelPath;      // ONNX 模型路径
        float confThreshold = 0.25; // 置信度阈值
        float iouThreshold = 0.45;  // NMS IoU 阈值
        int inputWidth  = 640;      // 模型输入宽度
        int inputHeight = 640;      // 模型输入高度
        bool useCUDA = false;       // ONNX CUDA EP
    };

    explicit YoloStoneDetector(const Config &cfg);
    ~YoloStoneDetector();

    /// 检测棋子，返回所有检测结果
    std::vector<YoloDetection> detect(const cv::Mat &bgrImage);

    /// 与棋盘网格对齐：将检测结果映射到 19x19 局面
    std::string alignToBoard(
        const std::vector<YoloDetection> &detections,
        const std::vector<cv::Point2f> &gridPoints,
        float snapThreshold = 0.4   // 占 cell 比例阈值
    );

private:
    void preprocess(const cv::Mat &src, cv::Mat &dst);
    std::vector<YoloDetection> postprocess(
        const float *output,       // [1, 84, 8400] or similar
        const cv::Size &originalSize);

    Ort::Env m_env;
    Ort::Session m_session;
    Ort::AllocatorWithDefaultOptions m_allocator;

    // 输入/输出节点名（因模型而异，需要查询）
    std::string m_inputName;
    std::string m_outputName;
    
    Config m_config;
    int m_numClasses;     // 类别数, 2
    int m_numOutputs;     // 每行输出 = 4 + numClasses (xywh + scores)
};

// detect/YoloStoneDetector.cpp
#include "YoloStoneDetector.h"
#include <onnxruntime_cxx_api.h>
#include <algorithm>
#include <cmath>

YoloStoneDetector::YoloStoneDetector(const Config &cfg)
    : m_env(ORT_LOGGING_LEVEL_WARNING, "YoloStone")
    , m_session(nullptr)
    , m_config(cfg)
    , m_numClasses(2)
    , m_numOutputs(4 + 2)  // xywh + black, white
{
    Ort::SessionOptions opts;
    opts.SetIntraOpNumThreads(4);
    opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    
    if (m_config.useCUDA) {
        OrtCUDAProviderOptions cudaOpts;
        cudaOpts.device_id = 0;
        opts.AppendExecutionProvider_CUDA(cudaOpts);
    }

    m_session = Ort::Session(m_env, m_config.modelPath.c_str(), opts);

    // 获取输入输出名称
    auto inputName = m_session.GetInputNameAllocated(0, m_allocator);
    m_inputName = inputName.get();
    auto outputName = m_session.GetOutputNameAllocated(0, m_allocator);
    m_outputName = outputName.get();
}

YoloStoneDetector::~YoloStoneDetector() = default;

void YoloStoneDetector::preprocess(const cv::Mat &src, cv::Mat &dst) {
    // 1. Letterbox resize → 640x640
    cv::Mat letterbox;
    float scale = std::min(
        (float)m_config.inputWidth  / src.cols,
        (float)m_config.inputHeight / src.rows
    );
    int newW = (int)(src.cols * scale);
    int newH = (int)(src.rows * scale);
    cv::resize(src, letterbox, cv::Size(newW, newH));
    
    int padW = m_config.inputWidth - newW;
    int padH = m_config.inputHeight - newH;
    int top = padH / 2, bottom = padH - top;
    int left = padW / 2, right = padW - left;
    cv::copyMakeBorder(letterbox, letterbox,
        top, bottom, left, right, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));

    // 2. BGR → RGB + HWC → CHW + normalize to [0, 1]
    letterbox.convertTo(dst, CV_32FC3, 1.0 / 255.0);
}

std::vector<YoloDetection> YoloStoneDetector::detect(const cv::Mat &bgrImage) {
    // ---- Preprocess ----
    cv::Mat blob;
    preprocess(bgrImage, blob);
    
    int64_t inputShape[] = {1, 3, m_config.inputHeight, m_config.inputWidth};
    auto memoryInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    
    Ort::Value inputTensor = Ort::Value::CreateTensor<float>(
        memoryInfo,
        (float*)blob.data,
        3 * m_config.inputHeight * m_config.inputWidth,
        inputShape, 4
    );

    // ---- Inference ----
    const char *inputNames[]  = { m_inputName.c_str() };
    const char *outputNames[] = { m_outputName.c_str() };
    
    auto outputs = m_session.Run(
        Ort::RunOptions{nullptr},
        inputNames, &inputTensor, 1,
        outputNames, 1
    );

    // ---- Postprocess ----
    auto *rawOutput = outputs[0].GetTensorMutableData<float>();
    auto outputShape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
    
    return postprocess(rawOutput, bgrImage.size());
}

std::vector<YoloDetection> YoloStoneDetector::postprocess(
    const float *output,
    const cv::Size &originalSize)
{
    // YOLOv8 ONNX 输出: [1, 6, 8400] or [1, 84, 8400] (取决于导出方式)
    // 这里假设简化版输出 [num_dets, 6] → [x1, y1, x2, y2, conf, class]
    // 参考 ultralytics 默认导出格式
    
    std::vector<YoloDetection> detections;
    
    // TODO: 根据实际 ONNX 输出格式调整解析逻辑
    // 典型 YOLOv8 输出: [1, 84, 8400]
    //   rows 0-3: 中心点 xywh (归一化)
    //   rows 4-83: 80 类 score
    // 围棋简化后: [1, 6, 8400]
    //   rows 0-3: xywh
    //   row 4: black score
    //   row 5: white score
    
    int numDets = 8400;  // 典型值，需要从实际输出 shape 获取
    
    for (int i = 0; i < numDets; ++i) {
        // 解析两类 score
        float scoreBlack = output[4 * numDets + i];
        float scoreWhite = output[5 * numDets + i];
        
        int bestClass;
        float maxScore;
        if (scoreBlack > scoreWhite) {
            bestClass = 0;
            maxScore = scoreBlack;
        } else {
            bestClass = 1;
            maxScore = scoreWhite;
        }
        
        if (maxScore < m_config.confThreshold) continue;
        
        // xywh 归一化 → 像素坐标
        float cx = output[0 * numDets + i];
        float cy = output[1 * numDets + i];
        float w  = output[2 * numDets + i];
        float h  = output[3 * numDets + i];
        
        float x1 = (cx - w / 2) * originalSize.width;
        float y1 = (cy - h / 2) * originalSize.height;
        float bw = w * originalSize.width;
        float bh = h * originalSize.height;
        
        detections.push_back({
            bestClass,
            maxScore,
            cv::Rect2f(x1, y1, bw, bh),
            cv::Point2f(cx * originalSize.width, cy * originalSize.height)
        });
    }
    
    // NMS（可选，YOLOv8 内置 NMS 时不需要）
    return detections;
}

std::string YoloStoneDetector::alignToBoard(
    const std::vector<YoloDetection> &detections,
    const std::vector<cv::Point2f> &gridPoints,
    float snapThreshold)
{
    // 361 个交叉点初始化为空
    std::string board(361, '.');
    
    // 计算 cell size
    float cellSize = cv::norm(gridPoints[0] - gridPoints[1]); // 相邻交点距离
    
    for (const auto &det : detections) {
        // 找最近的网格交点
        float minDist = FLT_MAX;
        int nearestIdx = -1;
        
        for (int i = 0; i < 361; ++i) {
            float dist = cv::norm(det.center - gridPoints[i]);
            if (dist < minDist) {
                minDist = dist;
                nearestIdx = i;
            }
        }
        
        // 阈值检查：离网格交点的距离 < cell 的一半
        if (minDist < cellSize * 0.5 && nearestIdx >= 0) {
            board[nearestIdx] = (det.classId == 0) ? 'B' : 'W';
        }
    }
    
    return board;
}

// detect/CvDnnYoloDetector.h
#pragma once
#include <opencv2/dnn.hpp>
#include <opencv2/opencv.hpp>

class CvDnnYoloDetector {
public:
    explicit CvDnnYoloDetector(const std::string &onnxPath) {
        m_net = cv::dnn::readNetFromONNX(onnxPath);
        // 尝试 OpenCL / CUDA 后端
        m_net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
        m_net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
    }

    std::vector<YoloDetection> detect(const cv::Mat &bgr) {
        // 构建 blob
        cv::Mat blob = cv::dnn::blobFromImage(
            bgr, 1.0/255.0, cv::Size(640, 640),
            cv::Scalar(), true, false  // swapRB=true (BGR→RGB)
        );

        m_net.setInput(blob);
        cv::Mat output = m_net.forward();

        return parseOutput(output, bgr.size());
    }

private:
    cv::dnn::Net m_net;
    // parseOutput() 逻辑同 ONNX Runtime 版本
};

// 在 GoBoardDetector 中增加 YOLO 路径

class GoBoardDetector : public IBoardDetector {
public:
    enum class StoneMode {
        Threshold,  // 原有 ROI 双阈值方案
        YOLO        // YOLO ONNX 推理
    };

    void setStoneMode(StoneMode mode, const std::string &modelPath = "");

    // 原有接口不变
    std::string recognizeStones(
        const cv::Mat &dewarpedBgr,
        const std::vector<cv::Point2f> &gridPixel) override
    {
        if (m_stoneMode == StoneMode::YOLO && m_yoloDetector) {
            auto dets = m_yoloDetector->detect(dewarpedBgr);
            return m_yoloDetector->alignToBoard(dets, gridPixel);
        }
        // fallback: 原有阈值方案
        return m_stoneDetector.classify(dewarpedBgr, gridPixel);
    }

private:
    StoneMode m_stoneMode = StoneMode::Threshold;
    std::unique_ptr<YoloStoneDetector> m_yoloDetector;
    StoneDetector m_stoneDetector;
};

# 用 YOLOv8s 教师模型蒸馏 YOLOv8n
from ultralytics import YOLO

# 训练大模型
teacher = YOLO('yolov8s.pt')
teacher.train(data='dataset/data.yaml', epochs=100, imgsz=640)

# 蒸馏训练小模型
student = YOLO('yolov8n.pt')
student.train(
    data='dataset/data.yaml',
    epochs=100,
    # 蒸馏参数（Ultralytics 目前需要手动实现蒸馏逻辑）
)

# 收集困难样本
# 1. 运行时检测：ROI 方案与 YOLO 方案结果不一致的帧 → 保存
# 2. 人工确认标注
# 3. 加入训练集重新训练

# INT8 量化（CPU 推理提速 ~2x，模型缩小 ~4x）
from ultralytics import YOLO

model = YOLO('best.pt')
model.export(
    format='onnx',
    imgsz=640,
    int8=True,          # INT8 量化
    data='dataset/data.yaml',  # 校准数据集
)