import torch
from ultralytics import YOLO
import onnx
import onnxsim

# 加载模型
model = YOLO('yolov8n.pt')

# 导出默认参数设置
input_shape = (1, 3, 640, 640)

# 导出 ONNX
model.export(
    format='onnx',
    opset=12,
    simplify=False,  # 先不简化，手动处理
    dynamic=False,
    batch=1,
    imgsz=640,
)

# 手动简化模型
onnx_model = onnx.load('yolov8n.onnx')

# 检查模型
onnx.checker.check_model(onnx_model)

# 简化
model_simp, check = onnxsim.simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"

# 保存简化后的模型
onnx.save(model_simp, 'yolov8n-sim.onnx')

# 克隆源码
git clone https://github.com/Tencent/ncnn.git
cd ncnn
mkdir build && cd build

# 编译（开启 ONNX 支持）
cmake -DNCNN_BUILD_TOOLS=ON -DNCNN_VULKAN=OFF ..
make -j8

# 编译好的工具在 tools/onnx/ 目录下

# 树莓派上直接编译
cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/pi3.toolchain.cmake ..

./onnx2ncnn yolov8n-sim.onnx yolov8n.param yolov8n.bin

model.export(format='onnx', opset=11, ...)

# 网络结构优化（去除冗余层）
./ncnnoptimize yolov8n.param yolov8n.bin yolov8n-opt.param yolov8n-opt.bin 65536

// 输出特征图处理
for (int i = 0; i < feat.w; i++) {
    float* ptr = feat.row(0) + i;
    
    // 中心点坐标
    float x = ptr[0 * feat.w];
    float y = ptr[1 * feat.w];
    float w = ptr[2 * feat.w];
    float h = ptr[3 * feat.w];
    
    // 找到最大置信度类别
    float max_score = 0;
    int max_class = -1;
    for (int c = 0; c < 80; c++) {
        float score = ptr[(4 + c) * feat.w];
        if (score > max_score) {
            max_score = score;
            max_class = c;
        }
    }
    
    if (max_score > conf_thresh) {
        // 转换到原图坐标
        float x1 = (x - w / 2) * scale;
        float y1 = (y - h / 2) * scale;
        float x2 = (x + w / 2) * scale;
        float y2 = (y + h / 2) * scale;
        
        objects.push_back({x1, y1, x2, y2, max_class, max_score});
    }
}

# 准备 100-500 张和业务场景相似的图片，放在 images/ 目录下
ls images/ | head -10

./ncnn2table yolov8n-opt.param yolov8n-opt.bin yolov8n.table images/ mean norm shape

./ncnn2table yolov8n-opt.param yolov8n-opt.bin yolov8n.table \
    images/ \
    0.0,0.0,0.0 \
    0.003921568627451,0.003921568627451,0.003921568627451 \
    640,640,3

./ncnn2int8 yolov8n-opt.param yolov8n-opt.bin yolov8n-int8.param yolov8n-int8.bin yolov8n.table

// Conv3x3_s1_neon
void conv3x3s1_neon(const Mat& bottom_blob, Mat& top_blob, ...)
{
    // 16 通道同时计算
    int nn = 16;
    // ...
    // NEON 内联汇编
    asm volatile (
        "vld1.8   {d0-d3}, [%[img]]       \n"
        "vmul.s8  q0, q0, %[weight]       \n"
        "vst1.8   {d0-d3}, [%[out]]       \n"
        :
        : [img]"r"(img_ptr), [weight]"r"(w_ptr), [out]"r"(out_ptr)
        : "memory", "q0", "q1"
    );
}

# CMakeLists.txt 关键配置
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -march=armv8-a+crc+simd")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=armv8-a+crc+simd")
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -O3 -ffast-math")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -ffast-math")

#include <opencv2/opencv.hpp>
#include "net.h"

class YOLOv8Detector {
public:
    YOLOv8Detector(const std::string& param_path, 
                   const std::string& bin_path,
                   bool use_int8 = false);
    ~YOLOv8Detector();
    
    std::vector<Object> detect(const cv::Mat& image, 
                               float conf_thresh = 0.25,
                               float nms_thresh = 0.45);
    
private:
    ncnn::Net net_;
    int input_size_ = 640;
    int num_classes_ = 80;
    
    void preprocess(const cv::Mat& image, ncnn::Mat& in);
    void postprocess(const ncnn::Mat& out, std::vector<Object>& objects,
                     float scale, float conf_thresh, float nms_thresh);
};

YOLOv8Detector::YOLOv8Detector(const std::string& param_path,
                               const std::string& bin_path,
                               bool use_int8) {
    // 加载模型
    net_.load_param(param_path.c_str());
    net_.load_model(bin_path.c_str());
    
    // 设置线程数（树莓派 4 核，设为 4）
    net_.opt.num_threads = 4;
    
    // 开启 NEON 优化
    net_.opt.use_vulkan_compute = false;
    net_.opt.use_fp16_packed = true;
    net_.opt.use_fp16_storage = true;
    net_.opt.use_fp16_arithmetic = true;
    net_.opt.use_int8_storage = use_int8;
    net_.opt.use_int8_arithmetic = use_int8;
}

void YOLOv8Detector::preprocess(const cv::Mat& image, ncnn::Mat& in) {
    // 按比例缩放，保持宽高比
    float scale = std::min(input_size_ * 1.0f / image.cols,
                           input_size_ * 1.0f / image.rows);
    
    int w = image.cols * scale;
    int h = image.rows * scale;
    
    // Letterbox 填充
    cv::Mat resized;
    cv::resize(image, resized, cv::Size(w, h));
    
    // 居中填充到 640x640
    cv::Mat padded = cv::Mat::zeros(input_size_, input_size_, CV_8UC3);
    int x_offset = (input_size_ - w) / 2;
    int y_offset = (input_size_ - h) / 2;
    resized.copyTo(padded(cv::Rect(x_offset, y_offset, w, h)));
    
    // BGR -> RGB，归一化
    in = ncnn::Mat::from_pixels(padded.data, ncnn::Mat::PIXEL_BGR2RGB,
                                input_size_, input_size_);
    
    // YOLOv8 预处理：除以 255
    float norm[3] = {1/255.0f, 1/255.0f, 1/255.0f};
    in.substract_mean_normalize(0, norm);
}

std::vector<Object> YOLOv8Detector::detect(const cv::Mat& image,
                                            float conf_thresh,
                                            float nms_thresh) {
    // 计算缩放比例
    float scale = std::min(input_size_ * 1.0f / image.cols,
                           input_size_ * 1.0f / image.rows);
    
    // 预处理
    ncnn::Mat in;
    preprocess(image, in);
    
    // 推理
    ncnn::Extractor ex = net_.create_extractor();
    ex.input("images", in);
    
    ncnn::Mat out;
    ex.extract("output0", out);
    
    // 后处理
    std::vector<Object> objects;
    postprocess(out, objects, scale, conf_thresh, nms_thresh);
    
    return objects;
}

void YOLOv8Detector::postprocess(const ncnn::Mat& out,
                                 std::vector<Object>& objects,
                                 float scale, float conf_thresh,
                                 float nms_thresh) {
    // out shape: [1, 84, 8400]，在 NCNN 中存储为 w=8400, h=84
    const int num_points = out.w;
    const int stride_offset = input_size_ / 64;
    
    for (int i = 0; i < num_points; i++) {
        // 找到最大类别得分
        float max_score = 0;
        int max_class = -1;
        
        for (int c = 0; c < num_classes_; c++) {
            float score = out.row(4 + c)[i];
            if (score > max_score) {
                max_score = score;
                max_class = c;
            }
        }
        
        if (max_score > conf_thresh) {
            float cx = out.row(0)[i];
            float cy = out.row(1)[i];
            float w = out.row(2)[i];
            float h = out.row(3)[i];
            
            // 转换回原图坐标（减去 letterbox 偏移）
            float x1 = (cx - w / 2) / scale;
            float y1 = (cy - h / 2) / scale;
            float x2 = (cx + w / 2) / scale;
            float y2 = (cy + h / 2) / scale;
            
            objects.push_back({x1, y1, x2, y2, max_class, max_score});
        }
    }
    
    // NMS
    qsort_descent_inplace(objects);
    std::vector<int> picked;
    nms_sorted_bboxes(objects, picked, nms_thresh);
}

// 在 param 文件中，给指定层加 flag=1
Convolution      conv_out  1  1  ...  256 84 1 1 1 1 0=1

from ultralytics import YOLO

# 教师模型
teacher = YOLO('yolov8l.pt')

# 学生模型
student = YOLO('yolov8n.yaml')

# 蒸馏训练
student.train(
    data='coco.yaml',
    epochs=100,
    distill='yolov8l.pt',
    distill_ratio=0.5
)