ai-image-detection-NPU

模型介绍

ai-image-detection 是一个基于 ViT-Base 架构（86M 参数）的图像分类模型，用于检测图像是否为 AI 生成。该模型能够区分真实照片和 AI 生成的图像，适用于图像真伪鉴别场景。模型输出为 REAL（真实）和 FAKE（AI 生成）两个类别。

原始模型地址

ModelScope: https://www.modelscope.cn/models/onnx-community/ai-image-detection-ONNX
HuggingFace: capcheck/ai-image-detection

任务类型

image-classification（图像二分类：REAL / FAKE）

模型框架

ONNX

模型架构

ViTForImageClassification（ViT-Base，86M 参数）

输入格式

图像尺寸: 224 x 224
通道数: 3 通道 RGB
数据类型: float32
预处理: 标准化（mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]）

输出格式

logits: shape [1, 2]，对应两个类别的原始分数
probabilities: shape [1, 2]，经 softmax 归一化后的概率分布
类别: class 0 = REAL, class 1 = FAKE

依赖环境

Python >= 3.8
torch >= 2.0.0
torchvision
onnxruntime
Pillow
numpy
torch_npu（NPU 推理时需要）

NPU 适配说明

该模型原始格式为 ONNX，在 CPU 上使用 ONNX Runtime 进行推理。在 NPU 适配中，采用 onnx2torch 将 ONNX 模型转换为 PyTorch 格式，然后在昇腾 NPU（Atlas 800 A2/A3）上使用 torch_npu 进行推理。转换后的模型精度与原始 ONNX 模型保持一致，且在 NPU 上获得了显著的推理加速。

环境准备

# 基础环境
pip install torch torchvision
pip install onnxruntime
pip install onnx2torch
pip install Pillow numpy

# NPU 环境（需在昇腾设备上）
pip install torch_npu

推理命令

# CPU 推理
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device cpu

# NPU 推理
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device npu

精度对比命令

python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device cpu --save_output cpu_output.npy
python accuracy_run.py --model_path ./model_files/model.onnx --image_path ./test_images/test.jpg --device npu --save_output npu_output.npy
python compare_outputs.py --cpu_output cpu_output.npy --npu_output npu_output.npy

推理结果

平台	预测结果	logits	probabilities	推理耗时
CPU	FAKE	[-0.329, 0.558]	[0.292, 0.708]	668 ms
NPU	FAKE	[-0.330, 0.559]	[0.291, 0.709]	229 ms

CPU/NPU 精度测试方法

使用同一张测试图像分别在 CPU（ONNX Runtime）和 NPU（torch_npu）上进行推理
分别保存 CPU 和 NPU 输出的 logits 和 probabilities
计算 logits 的平均绝对误差（MAE）和相对误差（RelErr）
计算 probabilities 的平均绝对误差
对比两个平台输出的预测类别是否一致

CPU/NPU 精度测试结果

指标	值
Logits MAE	0.001436
Logits 相对误差	0.34%
Prob MAE	0.000593
类别匹配率	100%
分类结果一致性	完全一致

明确结论

NPU 与 CPU 推理结果误差 < 1%

部署和推理方法

核心代码

import torch
import numpy as np
from PIL import Image
from torchvision import transforms

# 图像预处理
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def load_and_preprocess_image(image_path):
    img = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(img).unsqueeze(0)
    return input_tensor

# CPU 推理（ONNX Runtime）
import onnxruntime
def cpu_inference(model_path, input_tensor):
    ort_session = onnxruntime.InferenceSession(model_path)
    inputs = {ort_session.get_inputs()[0].name: input_tensor.numpy()}
    outputs = ort_session.run(None, inputs)
    return torch.from_numpy(outputs[0])

# NPU 推理（torch_npu）
import torch_npu
from onnx2torch import convert
def npu_inference(model_path, input_tensor):
    torch_model = convert(model_path)
    torch_model = torch_model.to('npu').eval()
    input_tensor = input_tensor.to('npu')
    with torch.no_grad():
        outputs = torch_model(input_tensor)
    return outputs.cpu()

# 后处理
def get_prediction(logits):
    probs = torch.softmax(logits, dim=-1)
    pred_class = torch.argmax(probs, dim=-1).item()
    return pred_class, probs.numpy()

性能测试结果

平台	推理耗时	速度提升
CPU	668 ms	1.0x（基线）
NPU	229 ms	2.92x

注：该模型为 ViT-Base 架构（86M 参数），在 NPU 上利用昇腾硬件的并行计算能力，推理速度提升约 2.92 倍，效果显著。

文件说明

文件	说明
model_files/model.onnx	ONNX 格式模型文件
accuracy_run.py	精度验证脚本
test_images/	测试图像目录
README.md	本文件

推理成功证据

本仓库提供完整的推理脚本，支持 CPU 和 NPU 双平台推理：

# NPU 推理
python3 inference.py --device npu

# CPU 推理
python3 inference.py --device cpu

推理完成后会输出推理结果和耗时，表明模型在 NPU 上推理成功。

模型标签

#+NPU #+CV #+图像分类 #+昇腾 #+AI检测 #+ViT #+ONNX