91 lines
3.4 KiB
C++
91 lines
3.4 KiB
C++
// 薄膜缺陷检测的 C++ 推理参考实现（OpenVINO）。
|
||
//
|
||
// 与 Python 侧（engine/tiling.py、engine/detection.py、scripts/openvino_predict.py）
|
||
// 逐步骤对应，两边结果已逐框比对过。
|
||
//
|
||
// 模型只接受固定 N×N 的输入且**禁止缩放**，所以任意尺寸的图要先切片：
|
||
// - 边长 < N：镜像填充（reflect，不重复边缘像素），原图居中；
|
||
// - 边长 >= N：按步长裁剪，末片贴边，保证整图被覆盖。
|
||
// 推理后把每片的框平移回原图，跨片做一次 NMS 去重。
|
||
//
|
||
// 产线如果本来就按 N×N 送图，只需要 Detector::infer_tile 这一条路径。
|
||
|
||
#pragma once
|
||
|
||
#include <cstdint>
|
||
#include <memory>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
namespace film {
|
||
|
||
// 归一化到 [0,1] 的 xyxy 框。
|
||
struct Box {
|
||
int class_id = 0;
|
||
float x1 = 0.f, y1 = 0.f, x2 = 0.f, y2 = 0.f;
|
||
float score = 0.f;
|
||
|
||
float area() const { return (x2 > x1 && y2 > y1) ? (x2 - x1) * (y2 - y1) : 0.f; }
|
||
};
|
||
|
||
// 一个片的取法：先从原图取 [y, y+crop_height) x [x, x+crop_width)，
|
||
// 再镜像填充成 size×size，原图内容落在 (pad_left, pad_top)。
|
||
struct Tile {
|
||
int x = 0, y = 0;
|
||
int crop_width = 0, crop_height = 0;
|
||
int pad_left = 0, pad_top = 0;
|
||
};
|
||
|
||
// 灰度或 RGB 的连续内存图像。channels 为 1 或 3。
|
||
struct ImageView {
|
||
const uint8_t* data = nullptr;
|
||
int width = 0;
|
||
int height = 0;
|
||
int channels = 1;
|
||
};
|
||
|
||
// 一个维度上的裁剪起点。末片强制贴边，末尾不足一步的区域不会被漏掉。
|
||
std::vector<int> axis_starts(int length, int size, int overlap);
|
||
|
||
// 把 width×height 切成若干 size×size 的片。
|
||
std::vector<Tile> plan_tiles(int width, int height, int size, int overlap = 0);
|
||
|
||
// 取出一片并镜像填充到 size×size，输出 size*size*channels 的连续缓冲。
|
||
void cut_tile(const ImageView& image, const Tile& tile, int size, std::vector<uint8_t>& out);
|
||
|
||
// 交并比。
|
||
float iou(const Box& a, const Box& b);
|
||
|
||
// inner 落在 outer 里的面积占比。面积悬殊时 IoU 很小但这个值接近 1。
|
||
float containment(const Box& inner, const Box& outer);
|
||
|
||
// 非极大值抑制。contained_threshold > 0 时额外按包含比例抑制：
|
||
// 缺陷被切割线切开时，相邻片只看到一角，给出的小残框与完整框的 IoU 很小，
|
||
// 光靠 IoU 压不掉，但它几乎整个落在完整框里。
|
||
std::vector<Box> nms(std::vector<Box> boxes, float iou_threshold = 0.5f,
|
||
bool per_class = true, float contained_threshold = 0.f);
|
||
|
||
class Detector {
|
||
public:
|
||
// model_path 指向导出目录里的 .xml；device 可用 "CPU" / "GPU" / "AUTO"。
|
||
Detector(const std::string& model_path, const std::string& device = "CPU");
|
||
~Detector();
|
||
|
||
int input_size() const; // 网络输入边长 N
|
||
int num_classes() const; // 输出通道数 - 4
|
||
|
||
// 对一张已经是 N×N 的图推理，返回该片内的归一化框（已做 NMS）。
|
||
std::vector<Box> infer_tile(const ImageView& tile, float conf = 0.25f,
|
||
float iou_threshold = 0.5f) const;
|
||
|
||
// 对任意尺寸的图：切片 -> 逐片推理 -> 平移回原图 -> 跨片 NMS。
|
||
std::vector<Box> infer(const ImageView& image, float conf = 0.25f,
|
||
float iou_threshold = 0.5f, int overlap = 32) const;
|
||
|
||
private:
|
||
struct Impl;
|
||
std::unique_ptr<Impl> impl_;
|
||
};
|
||
|
||
} // namespace film
|