使用Onnruntime实现Padim/Patchcore推理

Train.py

from pathlib import Path
from anomalib.data import Folder
from anomalib.models import Patchcore
from anomalib.engine import Engine
def main():
    # 1. 配置数据集
    datamodule = Folder(
        name="bottledataset",
        root=Path("D:/bottle"),  # 数据集根目录
        normal_dir="good",  # 正常样本文件夹名
        abnormal_dir="defect",  # 异常样本文件夹名
        train_batch_size=32,
        eval_batch_size=32,
        num_workers=4,
    )
    # 2. 定义 PatchCore 模型
    model = Patchcore(
        backbone="wide_resnet50_2",  # 特征提取骨干网络
        layers=["layer2", "layer3"],  # 提取的层
        pre_trained=True,  # 使用预训练权重
        num_neighbors=9  # 核心参数：邻居数量
    )
    # 3. 配置引擎并开始训练
    engine = Engine(
        max_epochs=1,  # PatchCore 通常 1 个 epoch 即可收敛
        accelerator="auto",  # 自动检测并使用 GPU/CPU
        devices=1
    )
    engine.fit(model=model, datamodule=datamodule)
    # 4. 训练结束后，保存模型权重
    engine.trainer.save_checkpoint("patchcore_model.ckpt", weights_only=False)
if __name__ == '__main__':
    main()

Predict.py

from anomalib.engine import Engine
from anomalib.models import Patchcore
import numpy as np
model = Patchcore(
    backbone="wide_resnet50_2"
)
engine = Engine()
predictions = engine.predict(
    model=model,
    ckpt_path="./results/Patchcore/bottledataset/latest/weights/lightning/model.ckpt",
    data_path="./test.png"
)
for p in predictions:
    print(
        "score:",
        p.pred_score
    )
    print(
        "label:",
        p.pred_label
    )
    print(
        "mask:",
        p.anomaly_map
    )

Export

from anomalib.engine import Engine
from anomalib.deploy import ExportType
from anomalib.models import Patchcore
model = Patchcore()
engine = Engine()
engine.export(
    model=model,
    ckpt_path="./results/Patchcore/bottledataset/latest/weights/lightning/model.ckpt",
    export_type=ExportType.ONNX,
    input_size=(256,256)
)

Deploy.py

import cv2
import numpy as np
import onnxruntime as ort
import time
# 1. 配置路径
image_path = "test.png"
model_path = "model.onnx"
# 2. 初始化 ONNX Runtime 会话
options = ort.SessionOptions()
options.log_severity_level = 3  # ORT_LOGGING_LEVEL_INFO
session = ort.InferenceSession(model_path, options, providers=['CPUExecutionProvider'])
# 3. 读取图像并预处理
original_image = cv2.imread(image_path)
if original_image is None:
    raise ValueError(f"Error: Could not read image at {image_path}")
original_height, original_width = original_image.shape[:2]
inp_height, inp_width = 256, 256
# BGR 转 RGB
rgb_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
# Resize 并归一化到 [0, 1]
resized = cv2.resize(rgb_image, (inp_width, inp_height))
resized = resized.astype(np.float32) / 255.0
# HWC 转 CHW (NCHW 格式)
input_data = np.transpose(resized, (2, 0, 1))  # Shape: (3, H, W)
input_data = np.expand_dims(input_data, axis=0)  # Shape: (1, 3, H, W)
# 4. 执行推理并计时
input_name = session.get_inputs()[0].name
start_time = time.time()
results = session.run(None, {input_name: input_data})
end_time = time.time()
inference_time_ms = (end_time - start_time) * 1000
# 5. 解析输出结果
# 假设输出顺序为: pred_score, pred_label, anomaly_map, pred_mask
# 为了更稳健，也可以通过遍历 session.get_outputs() 匹配名称
output_names = [out.name for out in session.get_outputs()]
result_dict = dict(zip(output_names, results))
score = result_dict.get("pred_score", np.array([0.0]))[0]
label = bool(result_dict.get("pred_label", np.array([False]))[0])
anomaly_map = result_dict.get("anomaly_map")
pred_mask = result_dict.get("pred_mask")
# 6. 解析 anomaly_map 形状
if anomaly_map is not None:
    map_dims = anomaly_map.shape
    map_h, map_w = 0, 0
    if len(map_dims) == 4 and map_dims[0] == 1 and map_dims[1] == 1:
        map_h, map_w = map_dims[2], map_dims[3]
    elif len(map_dims) == 3 and map_dims[0] == 1:
        map_h, map_w = map_dims[1], map_dims[2]
    else:
        total = anomaly_map.size
        map_h = int(np.sqrt(total))
        map_w = map_h
        if map_h * map_w != total:
            map_w = total // map_h
    anomaly_map = anomaly_map.reshape(map_h, map_w)
else:
    raise ValueError("anomaly_map not found in model outputs")
# 7. 后处理：生成热力图并叠加
# Resize 到原始图像尺寸
anomaly_resized = cv2.resize(anomaly_map, (original_width, original_height), interpolation=cv2.INTER_LINEAR)
# Min-Max 归一化到 [0, 1]
min_val, max_val = np.min(anomaly_resized), np.max(anomaly_resized)
if max_val - min_val > 1e-6:
    anomaly_norm = (anomaly_resized - min_val) / (max_val - min_val)
else:
    anomaly_norm = anomaly_resized.copy()
# 转为 8-bit 灰度图并应用 JET 伪彩色
anomaly_gray = (anomaly_norm * 255).astype(np.uint8)
heatmap = cv2.applyColorMap(anomaly_gray, cv2.COLORMAP_JET)
# 热力图与原图融合 (权重 0.5)
overlay = cv2.addWeighted(heatmap, 0.5, original_image, 0.5, 0)
# 8. 叠加 pred_mask 轮廓
if pred_mask is not None:
    # 将 bool 数组转为 uint8 掩膜
    mask_mat = (pred_mask.reshape(map_h, map_w).astype(np.uint8)) * 255
    mask_resized = cv2.resize(mask_mat, (original_width, original_height))
    # 创建红色半透明图层并只在 mask 区域叠加
    mask_color = np.zeros_like(original_image, dtype=np.uint8)
    mask_color[:, :, 2] = 255  # 红色通道
    blended = cv2.addWeighted(overlay, 1.0, mask_color, 0.3, 0)
    # 使用 mask 作为掩膜，将 blended 覆盖到 overlay 上
    overlay = np.where(mask_resized[:, :, np.newaxis] == 255, blended, overlay)
# 9. 拼接图像并输出结果
result_overlay_image = overlay.copy()
img_gray_to_color = cv2.cvtColor(mask_resized, cv2.COLOR_GRAY2BGR) if pred_mask is not None else np.zeros_like(
    original_image)
result_image = np.hstack([original_image, result_overlay_image, img_gray_to_color])
# 打印结果
print(f"推理耗时: {inference_time_ms:.2f} ms")
print(f"异常分数: {score:.4f}")
print(f"异常判定: {'异常' if label else '正常'}")
# 保存并显示结果
cv2.imwrite("result.png", result_image)
cv2.imshow("demo", result_image)
cv2.waitKey(5000)
cv2.destroyAllWindows()

Deploy.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.Eventing.Reader;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using static System.Net.Mime.MediaTypeNames;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using OpenCvSharp;
namespace ConsoleApp1
{
    internal class Program
    {
        static void Main(string[] args)
        {
            string image_path = "test.png";
            string model_path = "padim.onnx";
            DateTime dt1 = DateTime.Now;
            DateTime dt2 = DateTime.Now;
            Mat resultOverlayImage;         
            SessionOptions options = new SessionOptions();
            options.LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_INFO;
            options.AppendExecutionProvider_CPU(0); ;
            InferenceSession onnx_session = new InferenceSession(model_path, options); 
            Tensor<float> input_tensor;
            List<NamedOnnxValue> input_container = new List<NamedOnnxValue>();
            IDisposableReadOnlyCollection<DisposableNamedOnnxValue> result_infer;
            int inpHeight = 256, inpWidth = 256;
            Mat originalImage = new Mat(image_path);
            int originalWidth = originalImage.Cols;
            int originalHeight = originalImage.Rows;
            Mat rgbImage = new Mat();
            Cv2.CvtColor(originalImage, rgbImage, ColorConversionCodes.BGR2RGB);
            Mat resized = new Mat();
            Cv2.Resize(rgbImage, resized, new OpenCvSharp.Size(inpWidth, inpHeight));
            resized.ConvertTo(resized, MatType.CV_32FC3, 1.0 / 255.0);
            int height = inpHeight;
            int width = inpWidth;
            Mat[] channels = Cv2.Split(resized);  
            List<float> dataList = new List<float>();
            for (int c = 0; c < 3; c++)
            {
                float[] channelData = new float[height * width];
                System.Runtime.InteropServices.Marshal.Copy(channels[c].Data, channelData, 0, height * width);
                dataList.AddRange(channelData);
            }
            float[] inputData = dataList.ToArray();
            input_tensor = new DenseTensor<float>(inputData, new[] { 1, 3, height, width });
            input_container.Clear();
            input_container.Add(NamedOnnxValue.CreateFromTensor("input", input_tensor));
            dt1 = DateTime.Now;
            result_infer = onnx_session.Run(input_container);
            dt2 = DateTime.Now;
            var pred_score_tensor = result_infer.FirstOrDefault(x => x.Name == "pred_score")?.AsTensor<float>();
            var pred_label_tensor = result_infer.FirstOrDefault(x => x.Name == "pred_label")?.AsTensor<bool>();
            var anomaly_map_tensor = result_infer.FirstOrDefault(x => x.Name == "anomaly_map")?.AsTensor<float>();
            var pred_mask_tensor = result_infer.FirstOrDefault(x => x.Name == "pred_mask")?.AsTensor<bool>();
            float score = pred_score_tensor.First(); 
            bool label = pred_label_tensor?.First() ?? false;
            // 解析 anomaly_map：形状通常为 [1, 1, H, W] 或 [1, H, W]
            var dimensions = anomaly_map_tensor.Dimensions.ToArray();
            int mapH = dimensions.Length >= 3 ? dimensions[dimensions.Length - 2] : 0;
            int mapW = dimensions.Length >= 2 ? dimensions[dimensions.Length - 1] : 0;
            float[] mapData = anomaly_map_tensor.ToArray();
            // 确保是二维数据，如果有多余维度则 reshape
            if (dimensions.Length == 4 && dimensions[0] == 1 && dimensions[1] == 1)
            {
                // 已经是 [1,1,H,W] 直接使用
                mapH = dimensions[2];
                mapW = dimensions[3];
            }
            else if (dimensions.Length == 3 && dimensions[0] == 1)
            {
                // [1,H,W] 的情况
                mapH = dimensions[1];
                mapW = dimensions[2];
            }
            else
            {
                // 若形状不符合预期，尝试平铺后重新组织
                int total = mapData.Length;
                mapH = (int)Math.Sqrt(total);
                mapW = mapH;
                if (mapH * mapW != total) mapW = total / mapH;
            }
            // 将 anomaly_map 转为 Mat (CV_32FC1)
            Mat anomalyMat = Mat.FromPixelData(mapH, mapW, MatType.CV_32FC1, mapData);
            // ------------------ 后处理 ------------------
            // 1. 将异常图 resize 到原始图像尺寸
            Mat anomalyResized = new Mat();
            Cv2.Resize(anomalyMat, anomalyResized, new OpenCvSharp.Size(originalWidth, originalHeight), interpolation: InterpolationFlags.Linear);
            // 2. Min-Max 归一化到 [0,1] 范围
            double minVal, maxVal;
            Cv2.MinMaxLoc(anomalyResized, out minVal, out maxVal);
            Mat anomalyNorm = new Mat();
            if (maxVal - minVal > 1e-6)
            {
                anomalyResized.ConvertTo(anomalyNorm, MatType.CV_32FC1, 1.0 / (maxVal - minVal), -minVal / (maxVal - minVal));
            }
            else
            {
                anomalyNorm = anomalyResized.Clone();
            }
            Mat anomalyGray = new Mat();
            anomalyNorm.ConvertTo(anomalyGray, MatType.CV_8UC1, 255.0);
            // 4. 应用 JET 伪彩色生成热力图
            Mat heatmap = new Mat();
            Cv2.ApplyColorMap(anomalyGray, heatmap, ColormapTypes.Jet);
            // 5. 热力图与原图融合（权重 0.5 热力图 + 0.5 原图）
            Mat originalBGR = originalImage.Clone();
            Mat overlay = new Mat();
            Cv2.AddWeighted(heatmap, 0.5, originalBGR, 0.5, 0, overlay);
            Mat maskResized = new Mat();
            // 叠加 pred_mask 轮廓（二值掩膜）
            if (pred_mask_tensor != null)
            {
                bool[] maskData = pred_mask_tensor.ToArray();
                // 假设 mask 形状与 anomaly_map 相同，同样 resize 到原图大小
                Mat maskMat = new Mat(mapH, mapW, MatType.CV_8UC1);
                for (int i = 0; i < maskData.Length; i++)
                    maskMat.Set<byte>(i / mapW, i % mapW, maskData[i] ? (byte)255 : (byte)0);
                maskResized = new Mat();
                Cv2.Resize(maskMat, maskResized, new OpenCvSharp.Size(originalWidth, originalHeight));
                // 创建红色半透明图层
                Mat maskColor = new Mat(originalHeight, originalWidth, MatType.CV_8UC3, new Scalar(0, 0, 255));
                // 先计算全图加权（无 mask）
                Mat blended = new Mat();
                Cv2.AddWeighted(overlay, 1.0, maskColor, 0.3, 0, blended);
                // 将 mask 区域从 blended 复制到 overlay 中
                blended.CopyTo(overlay, maskResized);
            }
            resultOverlayImage = overlay.Clone();
            string resultText = $"推理耗时: {(dt2 - dt1).TotalMilliseconds:F2} ms\r\n";
            resultText += $"异常分数: {score:F4}\r\n";
            resultText += $"异常判定: {(label ? "异常" : "正常")}";
            Mat imgGrayToColor = new Mat();
            Cv2.CvtColor(maskResized, imgGrayToColor, ColorConversionCodes.GRAY2BGR);
            Mat[] imagesToConcat = { originalImage, resultOverlayImage, imgGrayToColor};
            Mat resultImage = new Mat();
            Cv2.HConcat(imagesToConcat, resultImage);
            Console.WriteLine(resultText);
            Cv2.ImWrite("result.png",resultImage);
            Cv2.ImShow("demo", resultImage);
            Cv2.WaitKey(5000);
        }
    }
}

Deplpy.cpp

#include <iostream>
#include <vector>
#include <string>
#include <chrono>
#include <cmath>
#include <algorithm>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
int main() {
    std::string image_path = "test.png";
    std::string model_path = "model.onnx";
    // 1. 初始化 ONNX Runtime 会话
    Ort::Env env(ORT_LOGGING_LEVEL_INFO, "ConsoleApp");
    Ort::SessionOptions session_options;
    std::wstring modelPath = std::wstring(model_path.begin(), model_path.end());
    session_options.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
    Ort::Session session_(env, modelPath.c_str(), session_options);
    // 2. 读取图像并预处理
    cv::Mat originalImage = cv::imread(image_path);
    if (originalImage.empty()) {
        std::cerr << "Error: Could not read image." << std::endl;
        return -1;
    }
    int originalWidth = originalImage.cols;
    int originalHeight = originalImage.rows;
    int inpHeight = 256, inpWidth = 256;
    // BGR 转 RGB
    cv::Mat rgbImage;
    cv::cvtColor(originalImage, rgbImage, cv::COLOR_BGR2RGB);
    // Resize 并归一化到 [0, 1]
    cv::Mat resized;
    cv::resize(rgbImage, resized, cv::Size(inpWidth, inpHeight));
    resized.convertTo(resized, CV_32FC3, 1.0 / 255.0);
    // HWC 转 CHW (NCHW 格式)
    std::vector<float> inputData(3 * inpHeight * inpWidth);
    std::vector<cv::Mat> channels(3);
    cv::split(resized, channels);
    for (int c = 0; c < 3; ++c) {
        std::memcpy(inputData.data() + c * inpHeight * inpWidth,
            channels[c].data,
            inpHeight * inpWidth * sizeof(float));
    }
    // 3. 创建输入 Tensor 并推理
    std::array<int64_t, 4> input_shape = { 1, 3, inpHeight, inpWidth };
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, inputData.data(), inputData.size(),
        input_shape.data(), input_shape.size()
    );
    const char* input_names[] = { "input" };
    const char* output_names[] = { "pred_score", "pred_label", "anomaly_map", "pred_mask" };
    auto start_time = std::chrono::high_resolution_clock::now();
    auto output_tensors = session_.Run(
        Ort::RunOptions{ nullptr },
        input_names, &input_tensor, 1,
        output_names, 4
    );
    auto end_time = std::chrono::high_resolution_clock::now();
    double inference_time_ms = std::chrono::duration<double, std::milli>(end_time - start_time).count();
    // 4. 解析输出结果
    float* pred_score_data = output_tensors[0].GetTensorMutableData<float>();
    float score = pred_score_data[0];
    bool* pred_label_data = output_tensors[1].GetTensorMutableData<bool>();
    bool label = pred_label_data[0];
    // 解析 anomaly_map
    float* map_data = output_tensors[2].GetTensorMutableData<float>();
    auto map_dims = output_tensors[2].GetTensorTypeAndShapeInfo().GetShape();
    int mapH = 0, mapW = 0;
    if (map_dims.size() == 4 && map_dims[0] == 1 && map_dims[1] == 1) {
        mapH = static_cast<int>(map_dims[2]);
        mapW = static_cast<int>(map_dims[3]);
    }
    else if (map_dims.size() == 3 && map_dims[0] == 1) {
        mapH = static_cast<int>(map_dims[1]);
        mapW = static_cast<int>(map_dims[2]);
    }
    else {
        int total = static_cast<int>(map_dims.back()); // fallback
        mapH = static_cast<int>(std::sqrt(total));
        mapW = mapH;
        if (mapH * mapW != total) mapW = total / mapH;
    }
    // 将 float 数据包装为 OpenCV Mat
    cv::Mat anomalyMat(mapH, mapW, CV_32FC1, map_data);
    // 5. 后处理：生成热力图并叠加
    cv::Mat anomalyResized;
    cv::resize(anomalyMat, anomalyResized, cv::Size(originalWidth, originalHeight), 0, 0, cv::INTER_LINEAR);
    // Min-Max 归一化
    double minVal, maxVal;
    cv::minMaxLoc(anomalyResized, &minVal, &maxVal);
    cv::Mat anomalyNorm;
    if (maxVal - minVal > 1e-6) {
        anomalyResized.convertTo(anomalyNorm, CV_32FC1, 1.0 / (maxVal - minVal), -minVal / (maxVal - minVal));
    }
    else {
        anomalyNorm = anomalyResized.clone();
    }
    cv::Mat anomalyGray;
    anomalyNorm.convertTo(anomalyGray, CV_8UC1, 255.0);
    // 应用 JET 伪彩色
    cv::Mat heatmap;
    cv::applyColorMap(anomalyGray, heatmap, cv::COLORMAP_JET);
    // 热力图与原图融合 (0.5, 0.5)
    cv::Mat overlay;
    cv::addWeighted(heatmap, 0.5, originalImage, 0.5, 0, overlay);
    // 6. 叠加 pred_mask 轮廓
    cv::Mat maskResized;
    bool* mask_data = output_tensors[3].GetTensorMutableData<bool>();
    cv::Mat maskMat(mapH, mapW, CV_8UC1);
    for (int i = 0; i < mapH * mapW; ++i) {
        maskMat.data[i] = mask_data[i] ? 255 : 0;
    }
    cv::resize(maskMat, maskResized, cv::Size(originalWidth, originalHeight));
    // 创建红色半透明图层并只在 mask 区域叠加
    cv::Mat maskColor(originalHeight, originalWidth, CV_8UC3, cv::Scalar(0, 0, 255));
    cv::Mat blended;
    cv::addWeighted(overlay, 1.0, maskColor, 0.3, 0, blended);
    blended.copyTo(overlay, maskResized);
    // 7. 拼接图像并输出结果
    cv::Mat resultOverlayImage = overlay.clone();
    cv::Mat imgGrayToColor;
    cv::cvtColor(maskResized, imgGrayToColor, cv::COLOR_GRAY2BGR);
    std::vector<cv::Mat> imagesToConcat = { originalImage, resultOverlayImage, imgGrayToColor };
    cv::Mat resultImage;
    cv::hconcat(imagesToConcat, resultImage);
    std::cout << "推理耗时: " << std::fixed << std::setprecision(2) << inference_time_ms << " ms\n";
    std::cout << "异常分数: " << std::fixed << std::setprecision(4) << score << "\n";
    std::cout << "异常判定: " << (label ? "异常" : "正常") << std::endl;
    cv::imwrite("result.png", resultImage);
    cv::imshow("demo", resultImage);
    cv::waitKey(5000);
    return 0;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.18)
project(Patch)
set("OpenCV_DIR" "E:\\Opencv\\opencv_vs\\build")
set("ONNXRUNTIME_DIR" "E:\\Onnxruntime\\cpu\\1.15")
set(OpenCV_INCLUDE_DIRS ${OpenCV_DIR}\\include)
set(OpenCV_LIB_DIRS ${OpenCV_DIR}\\x64\\vc16\\lib) 
set(OpenCV_LIBS "opencv_world480d.lib" "opencv_world480.lib")    
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
include_directories(${OpenCV_INCLUDE_DIRS}) 
link_directories(${OpenCV_LIB_DIRS})  
find_package(OpenCV QUIET)	
link_libraries(${OpenCV_LIBS})
add_executable(Patch Deploy.cpp)
target_compile_features(Patch PRIVATE cxx_std_14)
find_library(PATH ${ONNXRUNTIME_DIR})
target_include_directories(Patch PRIVATE "${ONNXRUNTIME_DIR}/include")
target_link_libraries(Patch "${ONNXRUNTIME_DIR}/lib/onnxruntime.lib")