import torch  
from net import FullyConnectedNet  
from PIL import Image  
import numpy as np  
from torchvision import transforms  
import torch.nn.functional as F 


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = FullyConnectedNet()  #模型结构，参数是随机初始化的，和训练好的模型权重没有关系
model_path = "./mlp_mse/model/mnist_net_best.pth"  

# 加载训练好的权重到模型中：
# 1. torch.load(model_path, map_location=device) 从磁盘读取权重字典（state_dict），
#    map_location=device 保证权重被映射到当前使用的设备上（防止在 GPU 上训练的模型在纯 CPU 机器上加载失败）
# 2. model.load_state_dict(...) 把读取到的权重参数逐个复制进模型的对应层中
model.load_state_dict(torch.load(model_path, map_location=device))

model.to(device)  # 把模型的所有参数和缓冲区移动到指定设备上（GPU 或 CPU），之后前向计算就在该设备上进行
model.eval()  
# 读取图片并进行预处理
image_path = "D:/code/xwdev/mlp_mse/mnist_pytorch_samples/image_0003_label_0.png"  
image = Image.open(image_path)  # 用 Pillow 打开图片文件，得到一个 PIL Image 对象

# 定义图像预处理流水线，transforms.Compose 会把多个变换按顺序串联起来依次执行
transform = transforms.Compose([
        transforms.Grayscale(num_output_channels=1),  # 转为灰度图：无论原图是 RGB 还是 RGBA，都强制变成单通道灰度图（MNIST 是灰度图，通道数必须为 1）
        transforms.Resize((28, 28)),                  # 调整大小为 28x28：MNIST 模型训练时输入就是 28x28，推理时尺寸必须一致
        transforms.ToTensor(),                        # 转为 tensor：把 PIL 图片转成 PyTorch 张量，形状为 [C, H, W]（这里是 [1, 28, 28]），
                                                      # 同时把像素值从 0~255 的整数缩放到 0.0~1.0 的浮点数
    ])

input_tensor = transform(image)  # 对图片执行上面定义的预处理流水线，得到形状为 [1, 28, 28] 的张量（通道、高、宽）
input_tensor = input_tensor.unsqueeze(0)  # 添加 batch 维度：在第 0 维插入一个大小为 1 的维度，
                                          # 形状从 [1, 28, 28] 变为 [1, 1, 28, 28]，即 NCHW 格式（批量数、通道、高、宽），
                                          # 因为模型要求输入是一“批”图片，即使只有一张也要包成 batch
input_tensor = input_tensor.to(device)  # 把输入张量移动到与模型相同的设备上，设备和模型必须一致，否则计算时会报错

# 使用模型进行预测
with torch.no_grad():  # 上下文管理器：在其内部所有操作都不会记录梯度信息。
                       # 推理阶段不需要反向传播，关闭梯度可以节省内存、加快计算
    output = model(input_tensor)  # 前向传播：把图片张量送入模型，得到输出。
                                  # 由于网络最后一层是 Softmax(dim=1)，output 是形状 [1, 10] 的概率分布，
                                  # 每一维表示该图片属于数字 0~9 的概率，10 个概率之和为 1

    predicted_class = torch.argmax(output, dim=1).item()  # 取概率最大的那一维的索引作为预测类别：
                                                          # argmax(output, dim=1) 在第 1 维（类别维）上找最大值的下标，返回形状 [1] 的张量
                                                          # .item() 把只含一个元素的张量转成普通的 Python 整数
    print(f"Predicted class: {predicted_class}")  # 打印预测结果，例如 "Predicted class: 0"

    # 下面三行是把送入模型的图片显示出来，方便人工核对预处理结果
    img_display = input_tensor.squeeze().cpu().numpy()  # 把张量变回 NumPy 数组以便转成图片：
                                                        # squeeze() 去掉所有大小为 1 的维度，[1,1,28,28] -> [28,28]
                                                        # .cpu() 把张量从 GPU 移回 CPU（NumPy 只能处理 CPU 上的数据）
                                                        # .numpy() 把 PyTorch 张量转成 NumPy 数组，此时像素值仍是 0.0~1.0 的浮点数
    image = Image.fromarray((img_display * 255).astype(np.uint8))  # 把 NumPy 数组转回 PIL 图片：
                                                                   # *255 把像素值从 0~1 还原回 0~255
                                                                   # astype(np.uint8) 转成 8 位无符号整数（图片标准格式）
                                                                   # Image.fromarray 根据数组创建灰度图对象
    image.show()  # 调用系统默认的图片查看器显示这张图片