#!/usr/bin/env python3
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms, models
from PIL import Image
import os
from datetime import datetime

class CustomDataset(Dataset):
    def __init__(self, img_dir, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.img_labels = self._load_labels()

    def _load_labels(self):
        # 假设标签存储在labels.txt文件中，每行格式为：图片名 标签
        with open(os.path.join(self.img_dir, 'labels.txt'), 'r') as f:
            lines = f.readlines()
        return [line.strip().split() for line in lines if line.strip()]

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

    def __getitem__(self, idx):
        img_name, label = self.img_labels[idx]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        label = int(label)
        return image, label

# 数据预处理
transform_train = transforms.Compose([
    transforms.RandomResizedCrop((128, 64)),  # 随机裁剪
    transforms.RandomHorizontalFlip(),       # 随机水平翻转
    transforms.ToTensor(),                   # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

transform_val = transforms.Compose([
    transforms.Resize((128, 64)),            # 调整大小
    transforms.ToTensor(),                   # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 加载数据集
img_dir = os.path.abspath("data/roi/train")
train_dataset = CustomDataset(img_dir=img_dir, transform=transform_train)
val_dataset = CustomDataset(img_dir=img_dir, transform=transform_val)

# 创建DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

# 加载预训练的ResNet18模型
model = models.resnet18(pretrained=True)

# 修改最后一层全连接层，使其输出为2（二分类）
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 2)

# 将模型移动到GPU（如果可用）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    # 训练阶段
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

    # 验证阶段
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Validation Accuracy: {100 * correct / total:.2f}%')

dt = datetime.now().strftime("%Y%m%d_%H%M%S")
torch.save(model.state_dict(), f'data/roi/resnet18_{dt}.pth')