# 强烈建议使用虚拟环境
python -m venv venv
source venv/bin/activate

# 安装 PyTorch 和 torchvision
pip install torch torchvision

# src/dataset.py
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
import pandas as pd

class CalorieDataset(Dataset):
    def __init__(self, csv_file, transform=None):
        """
        参数：
            csv_file (string): 包含标注的 csv 文件路径。
            transform (callable, optional): 可选的变换应用于样本。
        """
        self.food_frame = pd.read_csv(csv_file)
        self.transform = transform

    def __len__(self):
        return len(self.food_frame)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        img_path = self.food_frame.iloc[idx, 0]
        try:
            image = Image.open(img_path).convert('RGB')
        except FileNotFoundError:
            print(f"警告：在 {img_path} 找不到图像。跳过。")
            return None, None # 处理缺失图像

        calories = self.food_frame.iloc[idx, 1]
        calories = torch.tensor([calories], dtype=torch.float32)

        if self.transform:
            image = self.transform(image)

        return image, calories

# 定义变换
# 这些应该根据你的特定数据集调整
transform = transforms.Compose([
    transforms.Resize((224, 224)), # 将图像调整到固定大小
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 示例用法：
# calorie_dataset = CalorieDataset(csv_file='data/food_data.csv', transform=transform)
# dataloader = torch.utils.data.DataLoader(calorie_dataset, batch_size=32, shuffle=True)

# src/model.py
import torch
import torch.nn as nn
import torchvision.models as models

def get_calorie_estimation_model(pretrained=True):
    # 加载预训练模型
    model = models.resnet50(pretrained=pretrained)

    # 冻结预训练模型中的所有参数
    for param in model.parameters():
        param.requires_grad = False

    # 获取分类器输入特征数量
    num_ftrs = model.fc.in_features

    # 用我们的回归头替换最后的全连接层
    # 我们想要一个输出神经元用于卡路里值。
    model.fc = nn.Sequential(
        nn.Linear(num_ftrs, 512),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(512, 1) # 输出是单个连续值
    )
    
    return model

# 示例用法：
# model = get_calorie_estimation_model()
# print(model)

# src/train.py
import torch
import torch.optim as optim
from model import get_calorie_estimation_model
# 假设 dataloader 如步骤一所示已创建

def train_model(model, dataloader, num_epochs=10):
    # 定义损失函数和优化器
    criterion = torch.nn.MSELoss()
    # 我们只想优化新回归头的参数
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)

    for epoch in range(num_epochs):
        model.train() # 将模型设置为训练模式
        running_loss = 0.0

        for inputs, labels in dataloader:
            # 处理找不到图像的情况
            if inputs is None:
                continue
            
            inputs = inputs.to(device)
            labels = labels.to(device)

            # 清零参数梯度
            optimizer.zero_grad()

            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)

        epoch_loss = running_loss / len(dataloader.dataset)
        print(f"Epoch {epoch}/{num_epochs - 1}, 损失: {epoch_loss:.4f}")

    print("训练完成")
    return model

# 示例用法：
# model = get_calorie_estimation_model()
# trained_model = train_model(model, dataloader)

# src/inference.py
def predict_calories(model, image_path, transform):
    model.eval() # 将模型设置为评估模式
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0) # 添加批次维度

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    image = image.to(device)

    with torch.no_grad():
        prediction = model(image)
        
    return prediction.item()

# 示例用法：
# estimated_calories = predict_calories(trained_model, 'path/to/my_pizza.jpg', transform)
# print(f"估算卡路里: {estimated_calories:.0f}")