"""
定义图像分类推理流程
"""

import numpy as np
from PIL import Image
from mindx.sdk import Tensor # mxVision 中的 Tensor 数据结构
from mindx.sdk import base # mxVision 推理接口


class ClassificationInference:
    def __init__(self, model_path, input_size=(224, 224), swap=(2, 0, 1)):
        """
        初始化图像分类模型推理流程
        :param model_path: 图像分类模型onnx文件路径
        :param input_size: 模型输入大小
        :param swap: 变换方式，pytorch需要进行轴变换(默认参数)，tensorflow无需进行轴变换
        """
        self.model_path = model_path
        self.input_size = input_size
        self.swap = swap

    def input_processing(self, image_path):
        """
        对单个图像输入进行处理
        :param image_path: 图像路径
        :return: 处理后输出
        """
        # 打开图像并转换为RGB
        image = Image.open(image_path).convert("RGB")

        # 调整图像大小
        image = image.resize(self.input_size)

        # 转换为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(self.swap).copy()

        return image_array.astype(np.float32)

    def predict(self, image_path):
        """
        对单张图片进行推理
        :param image_path: 图片路径
        :return: 推理结果
        """
        image = self.input_processing(image_path)
        # 执行推理
        base.mx_init()  # 初始化 mxVision 资源
        img = Tensor(image[None, :, :, :])
        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
        outputs = model.infer([img])[0]  # 执行推理。输入数据类型：List[base.Tensor]， 返回模型推理输出的List[base.Tensor]
        outputs.to_host()
        outputs = np.array(outputs)  # 将outputs移动到内存中，并且将base.Tensor类转为numpy array
        return outputs

    def predict_batch(self, image_paths):
        """
        对指定图片列表进行批量推理
        :param image_paths: 待推理的图片路径列表
        :return: 批量推理结果
        """
        batch_images = []

        for image_path in image_paths:
            image = self.input_processing(image_path)
            batch_images.append(image)

        # 将批次图片堆叠成 (batch_size, 3, 224, 224) 维度
        batch_images = np.stack(batch_images)
        batch_images = Tensor(batch_images)
        # 执行推理
        base.mx_init()  # 初始化 mxVision 资源
        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
        outputs = model.infer([batch_images])[0]  # 执行推理。输入数据类型：List[base.Tensor]， 返回模型推理输出的List[base.Tensor]
        outputs.to_host()
        outputs = np.array(outputs)  # 将outputs移动到内存中，并且将base.Tensor类转为numpy array
        return outputs