import numpy as np
import onnxruntime as ort
from PIL import Image
import torchvision.transforms as T


# 读取并预处理图片
# def process_image(image_path):
#     image = Image.open(image_path).convert("RGB")
#     preprocess = T.Compose([
#         T.Resize((224, 224)),
#         T.ToTensor(),
#         T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
#     ])
#     return preprocess(image)
def process_image(image_path):
    # 打开图像并转换为RGB
    image = Image.open(image_path).convert("RGB")

    # 调整图像大小
    image = image.resize((224, 224))

    # 转换为numpy数组并归一化
    image_array = np.array(image) / 255.0  # 将像素值缩放到[0, 1]

    # 进行标准化
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image_array = (image_array - mean) / std
    image_array = image_array.transpose((2, 0, 1)).copy()

    return image_array.astype(np.float32)

# 推理函数
def infer(session, image):
    input_name = session.get_inputs()[0].name
    # ONNX 模型需要输入为 numpy 数组
    # image = image.numpy()  # 转换为 numpy 数组
    image = np.expand_dims(image, axis=0)  # 增加批次维度
    output = session.run(None, {input_name: image})
    return output

if __name__ == '__main__':
    # 加载模型
    model_path = 'blackbox_models/vgg16/vgg16.onnx'
    session = ort.InferenceSession(model_path)

    # 指定要推理的图片路径
    image_path = 'blackbox_models/vgg16/trigger/images/0/ILSVRC2012_val_00000537.JPEG'
    # 处理图像
    processed_image = process_image(image_path)

    # 进行推理
    predictions = infer(session, processed_image)

    # 输出预测结果
    print(predictions)

    cls = np.argmax(predictions[0])

    print(cls)