在模型中增加预处理步骤!

This commit is contained in:
zhanghuan
2025-09-11 13:52:36 +08:00
parent adefc3ecb2
commit ee829dc8c9
2 changed files with 44 additions and 2 deletions
+43 -1
View File
@@ -1,6 +1,8 @@
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
class FoodCNN(nn.Module): class FoodCNN(nn.Module):
@@ -10,6 +12,13 @@ class FoodCNN(nn.Module):
""" """
def __init__(self): def __init__(self):
super(FoodCNN, self).__init__() super(FoodCNN, self).__init__()
# 图片预处理变换
self.preprocess = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
# 第一个卷积块 # 第一个卷积块
self.conv1 = nn.Conv2d(3, 32, 3, padding=1) self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
@@ -33,6 +42,28 @@ class FoodCNN(nn.Module):
self.dropout4 = nn.Dropout(0.5) self.dropout4 = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 3) # 3分类 self.fc2 = nn.Linear(512, 3) # 3分类
def preprocess_image(self, image):
"""
预处理单张图片
Args:
image: PIL Image 或 numpy array
Returns:
torch.Tensor: 预处理后的张量,形状为 (1, 3, 32, 32)
"""
if not isinstance(image, Image.Image):
# 如果是numpy array,转换为PIL Image
if hasattr(image, 'shape'):
image = Image.fromarray(image)
else:
raise ValueError("输入必须是PIL Image或numpy array")
# 应用预处理变换
processed = self.preprocess(image)
# 添加batch维度
return processed.unsqueeze(0)
def forward(self, x): def forward(self, x):
# 第一个卷积块 # 第一个卷积块
x = F.relu(self.conv1(x)) x = F.relu(self.conv1(x))
@@ -81,4 +112,15 @@ if __name__ == "__main__":
# 测试前向传播 # 测试前向传播
dummy_input = torch.randn(1, 3, 32, 32) dummy_input = torch.randn(1, 3, 32, 32)
output = model(dummy_input) output = model(dummy_input)
print(f"输出形状: {output.shape}") print(f"输出形状: {output.shape}")
# 测试图片预处理
try:
import numpy as np
# 创建一个测试图片 (RGB格式)
test_image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
processed = model.preprocess_image(test_image)
print(f"预处理后图片形状: {processed.shape}")
print(f"预处理后数值范围: [{processed.min():.3f}, {processed.max():.3f}]")
except Exception as e:
print(f"预处理测试失败: {e}")
+1 -1
View File
@@ -15,7 +15,7 @@ VAL_DATA_DIR = os.path.join(DATASET_DIR, 'val')
TEST_DATA_DIR = os.path.join(DATASET_DIR, 'test') TEST_DATA_DIR = os.path.join(DATASET_DIR, 'test')
# 模型保存路径 # 模型保存路径
MODEL_DIR = os.path.join(BASE_DIR, 'model', '06') MODEL_DIR = os.path.join(BASE_DIR, 'model', '07')
BEST_MODEL_PATH = os.path.join(MODEL_DIR, 'best_food_model.pth') BEST_MODEL_PATH = os.path.join(MODEL_DIR, 'best_food_model.pth')
TRAINING_CURVES_PATH = os.path.join(MODEL_DIR, 'training_curves.png') TRAINING_CURVES_PATH = os.path.join(MODEL_DIR, 'training_curves.png')
TRAINING_RESULTS_PATH = os.path.join(MODEL_DIR, 'training_results.txt') TRAINING_RESULTS_PATH = os.path.join(MODEL_DIR, 'training_results.txt')