diff --git a/net/food_net.py b/net/food_net.py index 37fd284..d5ebb46 100644 --- a/net/food_net.py +++ b/net/food_net.py @@ -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}") \ No newline at end of file + 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}") \ No newline at end of file diff --git a/settings/settings.py b/settings/settings.py index c92e99a..ea56fa1 100644 --- a/settings/settings.py +++ b/settings/settings.py @@ -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')