在模型中增加预处理步骤!
This commit is contained in:
+43
-1
@@ -1,6 +1,8 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision import transforms
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class FoodCNN(nn.Module):
|
||||
@@ -10,6 +12,13 @@ class FoodCNN(nn.Module):
|
||||
"""
|
||||
def __init__(self):
|
||||
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.conv2 = nn.Conv2d(32, 32, 3, padding=1)
|
||||
@@ -33,6 +42,28 @@ class FoodCNN(nn.Module):
|
||||
self.dropout4 = nn.Dropout(0.5)
|
||||
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):
|
||||
# 第一个卷积块
|
||||
x = F.relu(self.conv1(x))
|
||||
@@ -81,4 +112,15 @@ if __name__ == "__main__":
|
||||
# 测试前向传播
|
||||
dummy_input = torch.randn(1, 3, 32, 32)
|
||||
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}")
|
||||
@@ -15,7 +15,7 @@ VAL_DATA_DIR = os.path.join(DATASET_DIR, 'val')
|
||||
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')
|
||||
TRAINING_CURVES_PATH = os.path.join(MODEL_DIR, 'training_curves.png')
|
||||
TRAINING_RESULTS_PATH = os.path.join(MODEL_DIR, 'training_results.txt')
|
||||
|
||||
Reference in New Issue
Block a user