# 1. 导入工具库
import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader

# 2. 定义数据路径与标签映射（需根据实际文件结构调整）
train_val_dir = "path/to/train_validation_images"  # 5000张已标注图像路径
test_dir = "path/to/test_images"  # 1800张未标注图像路径

# 标签映射：密度分级（D1-D4）+良恶性（B=Malignant，M=Benign）→ 统一标签（如D1B=0，D1M=1...D4M=7）
label_map = {
    "D1B": 0, "D1M": 1,
    "D2B": 2, "D2M": 3,
    "D3B": 4, "D3M": 5,
    "D4B": 6, "D4M": 7
}

# 3. 自定义数据集类（读取图像+标注）
class BreastXRayDataset(Dataset):
    def __init__(self, image_dir, label_map=None, is_test=False, transform=None):
        self.image_dir = image_dir
        self.image_names = os.listdir(image_dir)
        self.label_map = label_map
        self.is_test = is_test  # 测试集无标注，返回图像名
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = self.image_names[idx]
        img_path = os.path.join(self.image_dir, img_name)
        
        # 读取图像（224×224×3），OpenCV默认BGR，转为RGB
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 适配多数模型的RGB输入要求
        img = np.array(img, dtype=np.float32)
        
        # 标准化：像素值从[0,255]转为[0,1]（或用ImageNet均值方差标准化，视模型而定）
        img = img / 255.0
        
        # 处理标签（训练验证集）或返回图像名（测试集）
        if not self.is_test:
            # 假设图像名包含标签（如"D1B_001.jpg"），提取标签部分
            label_str = img_name.split("_")[0]  # 需根据实际命名规则调整
            label = self.label_map[label_str]
            return img.transpose(2, 0, 1), label  # 转为[C, H, W]格式（适配PyTorch）
        else:
            return img.transpose(2, 0, 1), img_name

# 4. 划分训练集与验证集（7:3）
# 先获取训练验证集所有图像名与标签
all_labels = [label_map[img_name.split("_")[0]] for img_name in os.listdir(train_val_dir)]
train_names, val_names, train_labels, val_labels = train_test_split(
    os.listdir(train_val_dir), all_labels, test_size=0.3, random_state=42, stratify=all_labels
)

# 5. 创建DataLoader（批量加载数据）
train_dataset = BreastXRayDataset(
    image_dir=train_val_dir, label_map=label_map, is_test=False, transform=None
)
val_dataset = BreastXRayDataset(
    image_dir=train_val_dir, label_map=label_map, is_test=False, transform=None
)
test_dataset = BreastXRayDataset(
    image_dir=test_dir, label_map=None, is_test=True, transform=None
)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

import torch
import torch.nn as nn
from torch.optim import Adam
from torchvision.models import efficientnet_b0

# 1. 加载预训练模型并修改输出层（原输出1000类，改为8类）
model = efficientnet_b0(pretrained=True)  # 加载ImageNet预训练权重
in_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(in_features, 8)  # 8类：D1B-D4M
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 2. 定义损失函数与优化器
criterion = nn.CrossEntropyLoss()  # 多分类交叉熵损失
optimizer = Adam(model.parameters(), lr=1e-4)  # 学习率1e-4，平衡收敛与震荡

# 3. 模型训练（简化版，实际需加早停、学习率衰减）
num_epochs = 20
for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    train_acc = 0.0
    
    # 训练循环
    for imgs, labels in train_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        
        # 前向传播
        outputs = model(imgs)
        loss = criterion(outputs, labels)
        
        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 计算训练损失与准确率
        train_loss += loss.item() * imgs.size(0)
        _, preds = torch.max(outputs, 1)
        train_acc += torch.sum(preds == labels.data)
    
    # 计算epoch平均损失与准确率
    train_loss = train_loss / len(train_loader.dataset)
    train_acc = train_acc.double() / len(train_loader.dataset)
    
    # 验证阶段
    model.eval()
    val_loss = 0.0
    val_acc = 0.0
    with torch.no_grad():
        for imgs, labels in val_loader:
            imgs, labels = imgs.to(device), labels.to(device)
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            
            val_loss += loss.item() * imgs.size(0)
            _, preds = torch.max(outputs, 1)
            val_acc += torch.sum(preds == labels.data)
    
    val_loss = val_loss / len(val_loader.dataset)
    val_acc = val_acc.double() / len(val_loader.dataset)
    
    # 打印日志
    print(f"Epoch {epoch+1}/{num_epochs}")
    print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
    print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}\n")

# 4. 测试集推理（无标注，保存预测结果）
model.eval()
test_preds = []
test_img_names = []
with torch.no_grad():
    for imgs, img_names in test_loader:
        imgs = imgs.to(device)
        outputs = model(imgs)
        _, preds = torch.max(outputs, 1)
        
        test_preds.extend(preds.cpu().numpy())
        test_img_names.extend(img_names)

# 保存测试集预测结果（图像名+预测标签）
result_df = pd.DataFrame({
    "image_name": test_img_names,
    "predicted_label": [list(label_map.keys())[list(label_map.values()).index(p)] for p in test_preds]
})
result_df.to_csv("breast_xray_test_predictions.csv", index=False)