传统乳腺结节诊断流程模拟
def traditional_breast_diagnosis(mammogram):
     Step 1: 放射科医生A阅片（主观性强，易受疲劳影响）
    doctor_a_birads = read_mammogram_manually(mammogram)   耗时3-5分钟/例
     Step 2: 放射科医生B独立阅片
    doctor_b_birads = read_mammogram_manually(mammogram)
     Step 3: 结果比对与仲裁
    if doctor_a_birads == doctor_b_birads:
        return doctor_a_birads
    else:
        return arbitration(doctor_a_birads, doctor_b_birads)   引入第三位专家

BI-RADS分级体系（临床标准）
       ↓ 提供分级标准
  乳腺结节AI助手
       ↓ 由两类任务支撑
   ┌───┴───┐
   ↓       ↓
 检测     分割
（在哪里） （长什么样）
       ↓ 共同输入
  可解释AI技术
（为什么这么判断）

 -- coding: utf-8 --
"""
乳腺结节AI助手核心实现（基于PyTorch）
功能模块：数据预处理 | UNet分割 | 分类推理 | 热力图可视化
版本：2026年4月
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

 ==================== 1. 数据预处理模块 ====================
class BreastUltrasoundDataset(Dataset):
    """乳腺超声数据集加载器（支持BUSI数据集格式）"""
    def __init__(self, image_paths, mask_paths=None, transform=None):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform
        
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
         加载原始超声图像，归一化到[0,1]
        image = np.array(Image.open(self.image_paths[idx])) / 255.0
        if len(image.shape) == 2:
            image = np.stack([image]3, axis=2)   转为3通道
        
        if self.mask_paths is not None:
            mask = np.array(Image.open(self.mask_paths[idx])) / 255.0
            mask = (mask > 0.5).astype(np.float32)
            return torch.FloatTensor(image).permute(2,0,1), torch.FloatTensor(mask).unsqueeze(0)
        return torch.FloatTensor(image).permute(2,0,1)

 ==================== 2. UNet分割模型架构 ====================
class DoubleConv(nn.Module):
    """双卷积块：Conv2d -> BN -> ReLU -> Conv2d -> BN -> ReLU"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True)
        )
    def forward(self, x): return self.conv(x)

class UNet(nn.Module):
    """U-Net：医学图像分割领域最经典的编码器-解码器架构"""
    def __init__(self, in_channels=3, out_channels=1, features=[64,128,256,512]):
        super().__init__()
         编码器（下采样路径）—— 提取多尺度特征
        self.encoders = nn.ModuleList()
        self.pools = nn.ModuleList()
        for idx, f in enumerate(features):
            self.encoders.append(DoubleConv(in_channels if idx==0 else features[idx-1], f))
            self.pools.append(nn.MaxPool2d(kernel_size=2, stride=2))
        
         瓶颈层 —— 最深层的语义特征
        self.bottleneck = DoubleConv(features[-1], features[-1]2)
        
         解码器（上采样路径）—— 恢复空间分辨率 + 跳跃连接融合
        self.upconvs = nn.ModuleList()
        self.decoders = nn.ModuleList()
        for idx, f in enumerate(reversed(features)):
            self.upconvs.append(nn.ConvTranspose2d(f2 if idx==0 else f, f, kernel_size=2, stride=2))
            self.decoders.append(DoubleConv(f2 if idx==0 else f3, f))
        
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
    
    def forward(self, x):
        skip_connections = []
         编码阶段：保存每层输出用于跳跃连接
        for encoder, pool in zip(self.encoders, self.pools):
            x = encoder(x)
            skip_connections.append(x)
            x = pool(x)
        x = self.bottleneck(x)
        
         解码阶段：上采样后与对应编码层特征拼接
        skip_connections = skip_connections[::-1]
        for idx, (upconv, decoder) in enumerate(zip(self.upconvs, self.decoders)):
            x = upconv(x)
            if x.shape != skip_connections[idx].shape:
                x = F.interpolate(x, size=skip_connections[idx].shape[2:], mode='bilinear')
            x = torch.cat([skip_connections[idx], x], dim=1)
            x = decoder(x)
        return torch.sigmoid(self.final_conv(x))

 ==================== 3. 结节分类模型 ====================
class NoduleClassifier(nn.Module):
    """
    基于ResNet风格的结节良恶性分类器
    输入：分割掩码框定后的结节ROI区域
    输出：良性概率 vs 恶性概率
    """
    def __init__(self, num_classes=2):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 7, stride=2, padding=3)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
         简化的残差块
        self.layer1 = self._make_layer(64, 64, 2)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Linear(256, num_classes)
    
    def _make_layer(self, in_ch, out_ch, blocks, stride=1):
        layers = []
        layers.append(nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1))
        layers.append(nn.BatchNorm2d(out_ch))
        layers.append(nn.ReLU(inplace=True))
        for _ in range(1, blocks):
            layers.append(nn.Conv2d(out_ch, out_ch, 3, padding=1))
            layers.append(nn.BatchNorm2d(out_ch))
            layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(layers)
    
    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.fc(x)

 ==================== 4. 训练与推理 ====================
def train_segmentation_model(model, train_loader, epochs=50, lr=1e-4):
    """
    训练分割模型
    损失函数：组合Dice Loss + BCE Loss
    """
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
    
    def dice_loss(pred, target, smooth=1e-6):
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        intersection = (pred_flat  target_flat).sum()
        return 1 - (2.  intersection + smooth) / (pred_flat.sum() + target_flat.sum() + smooth)
    
    for epoch in range(epochs):
        model.train()
        epoch_loss = 0
        for images, masks in train_loader:
            optimizer.zero_grad()
            outputs = model(images)
            bce = F.binary_cross_entropy(outputs, masks)
            dice = dice_loss(outputs, masks)
            loss = bce + dice
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        if epoch % 10 == 0:
            print(f"Epoch {epoch}, Loss: {epoch_loss/len(train_loader):.4f}")

def infer_with_heatmap(seg_model, cls_model, image, use_gradcam=True):
    """
    完整推理流程：分割 + 分类 + Grad-CAM热力图
    """
    seg_model.eval(); cls_model.eval()
    with torch.no_grad():
        seg_mask = seg_model(image.unsqueeze(0)).squeeze()
        seg_mask_bin = (seg_mask > 0.5).float()
    
     提取ROI区域用于分类
    roi = image  seg_mask_bin
    cls_output = cls_model(roi.unsqueeze(0))
    probability = torch.softmax(cls_output, dim=1)
    
     Grad-CAM热力图生成（简化版）
    heatmap = None
    if use_gradcam:
         获取分类模型最后一个卷积层的特征图
        features = cls_model.layer3[:2](roi.unsqueeze(0))
        cls_model.zero_grad()
        cls_output[:, 1].backward()   针对恶性类别反向传播
        gradients = cls_model.layer3[:2][-1][-1].weight.grad   简化示意
         生成热力图并叠加到原图上
        heatmap = np.mean(gradients.cpu().detach().numpy(), axis=(2,3))
        heatmap = np.maximum(heatmap, 0)
        heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
    
    return seg_mask, probability, heatmap

 ==================== 5. 主程序示例 ====================
if __name__ == "__main__":
     初始化模型
    seg_model = UNet(in_channels=3, out_channels=1)
    cls_model = NoduleClassifier(num_classes=2)
    
     模拟数据加载（实际使用时替换为真实图像路径）
     train_dataset = BreastUltrasoundDataset(image_paths, mask_paths)
     train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
     train_segmentation_model(seg_model, train_loader)
    
    print("模型已初始化，参数量统计：")
    print(f"UNet分割模型参数量: {sum(p.numel() for p in seg_model.parameters()):,}")
    print(f"分类模型参数量: {sum(p.numel() for p in cls_model.parameters()):,}")
    
     推理示例（使用随机输入演示）
    dummy_image = torch.rand(3, 256, 256)
    mask, prob, hm = infer_with_heatmap(seg_model, cls_model, dummy_image)
    print(f"推理完成：恶性概率 = {prob[0][1].item():.3f}")