图像分割：语义分割与实例分割

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms

# 分割任务的输入输出尺寸对照
# 输入: (B, 3, H, W) → 输出: (B, C, H, W)
# 每个像素输出 C 个类别的 logits
batch_size, num_classes = 4, 21
H, W = 512, 512

dummy_input = torch.randn(batch_size, 3, H, W)
dummy_output = torch.randn(batch_size, num_classes, H, W)

# 交叉熵损失: 自动对每个像素计算
criterion = nn.CrossEntropyLoss()
loss = criterion(dummy_output, torch.randint(0, num_classes, (batch_size, H, W)))
print(f"Pixel-wise CE Loss: {loss.item():.4f}")

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleFCN(nn.Module):
    """简化版 FCN：VGG-style 编码器 + 转置卷积解码器"""
    def __init__(self, num_classes=21):
        super().__init__()
        # 编码器 (模拟 VGG16 的 pool5, 输出 stride=32)
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),  # /2
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),  # /4
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),  # /8
            nn.Conv2d(256, 512, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),  # /16
            nn.Conv2d(512, 512, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),  # /32
        )
        # 1x1 卷积替代全连接层
        self.classifier = nn.Conv2d(512, num_classes, 1)
        # 转置卷积上采样 (32x → 1x)
        self.upsample = nn.ConvTranspose2d(
            num_classes, num_classes, kernel_size=64,
            stride=32, padding=16
        )

    def forward(self, x):
        h, w = x.shape[2], x.shape[3]
        x = self.encoder(x)          # (B, 512, H/32, W/32)
        x = self.classifier(x)       # (B, C, H/32, W/32)
        x = self.upsample(x)         # (B, C, H, W)
        return x[:, :, :h, :w]       # 裁剪到原图尺寸

import torch
import torch.nn as nn

class UNetBlock(nn.Module):
    """U-Net 的基本卷积块: Conv → BN → ReLU → Conv → BN → ReLU"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.block = 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.block(x)

class UNet(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        # 编码器
        self.enc1 = UNetBlock(3, 64)
        self.enc2 = UNetBlock(64, 128)
        self.enc3 = UNetBlock(128, 256)
        self.pool = nn.MaxPool2d(2)
        # 瓶颈层
        self.bottleneck = UNetBlock(256, 512)
        # 解码器
        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = UNetBlock(512, 256)  # 256+256=512
        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = UNetBlock(256, 128)  # 128+128=256
        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = UNetBlock(128, 64)   # 64+64=128
        # 输出
        self.out_conv = nn.Conv2d(64, num_classes, 1)

    def forward(self, x):
        # 编码
        e1 = self.enc1(x)       # (B, 64, H, W)
        e2 = self.enc2(self.pool(e1))   # (B, 128, H/2, W/2)
        e3 = self.enc3(self.pool(e2))   # (B, 256, H/4, W/4)
        # 瓶颈
        b = self.bottleneck(self.pool(e3))  # (B, 512, H/8, W/8)
        # 解码 + 跳跃连接
        d3 = self.up3(b)        # (B, 256, H/4, W/4)
        d3 = self.dec3(torch.cat([d3, e3], dim=1))
        d2 = self.up2(d3)
        d2 = self.dec2(torch.cat([d2, e2], dim=1))
        d1 = self.up1(d2)
        d1 = self.dec1(torch.cat([d1, e1], dim=1))
        return self.out_conv(d1)

import torch
import torch.nn as nn
import torch.nn.functional as F

class ASPP(nn.Module):
    """Atrous Spatial Pyramid Pooling - DeepLab v3 核心"""
    def __init__(self, in_channels=2048, out_channels=256, rates=[6, 12, 18]):
        super().__init__()
        # 1x1 卷积分支
        self.branch1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )
        # 多尺度空洞卷积分支
        self.branches = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 3, padding=r, dilation=r),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
            ) for r in rates
        ])
        # 全局平均池化分支
        self.global_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )
        self.fuse = nn.Sequential(
            nn.Conv2d(out_channels * (len(rates) + 2), out_channels, 1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        h, w = x.shape[2], x.shape[3]
        feats = [self.branch1(x)]
        feats.extend([branch(x) for branch in self.branches])
        # 全局池化后上采样回原尺寸
        g = self.global_pool(x)
        g = F.interpolate(g, size=(h, w), mode='bilinear', align_corners=True)
        feats.append(g)
        return self.fuse(torch.cat(feats, dim=1))

import torch
import torch.nn as nn
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.rpn import AnchorGenerator
import torchvision.transforms as T

# 使用 torchvision 预训练 Mask R-CNN
model = maskrcnn_resnet50_fpn(weights="DEFAULT")

# 修改分类头以适应自定义类别数
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = nn.Sequential(
    nn.Linear(in_features, 1024),
    nn.ReLU(),
    nn.Linear(1024, 5)  # 4 类 + 背景
)
# 修改掩码头
model.roi_heads.mask_predictor = nn.Sequential(
    nn.ConvTranspose2d(256, 256, 2, 2),
    nn.ReLU(),
    nn.Conv2d(256, 5, 1)  # 4 类 + 背景
)

# 推理示例
model.eval()
transform = T.Compose([T.ToTensor()])
img = torch.randn(3, 800, 1200)  # 模拟输入图像
with torch.no_grad():
    predictions = model([img])

# 解析输出
pred = predictions[0]
boxes = pred["boxes"]          # [N, 4] 边界框
labels = pred["labels"]        # [N] 类别
scores = pred["scores"]        # [N] 置信度
masks = pred["masks"]          # [N, 1, H, W] 掩码概率图
print(f"检测到 {len(boxes)} 个实例")
# 阈值过滤: scores > 0.5
keep = scores > 0.5
print(f"高置信度实例: {keep.sum().item()} 个")

import torch
import numpy as np
from typing import Dict

def compute_iou(pred: torch.Tensor, target: torch.Tensor, num_classes: int) -> torch.Tensor:
    """计算每个类别的 IoU
    pred, target: (H, W) 长整型类别图
    返回: (num_classes,) 每个类别的 IoU
    """
    ious = []
    for c in range(num_classes):
        pred_c = (pred == c)
        target_c = (target == c)
        intersection = (pred_c & target_c).sum().float()
        union = (pred_c | target_c).sum().float()
        if union == 0:
            ious.append(torch.tensor(1.0))  # 该类别不存在，视为完美
        else:
            ious.append(intersection / union)
    return torch.tensor(ious)

def compute_dice(pred: torch.Tensor, target: torch.Tensor, num_classes: int) -> torch.Tensor:
    """计算每个类别的 Dice 系数"""
    dices = []
    for c in range(num_classes):
        pred_c = (pred == c).float()
        target_c = (target == c).float()
        intersection = (pred_c * target_c).sum()
        dice = (2 * intersection) / (pred_c.sum() + target_c.sum() + 1e-8)
        dices.append(dice)
    return torch.tensor(dices)

# 模拟评估
pred_map = torch.randint(0, 5, (512, 512))
target_map = torch.randint(0, 5, (512, 512))

ious = compute_iou(pred_map, target_map, 5)
mIoU = ious.mean()
dices = compute_dice(pred_map, target_map, 5)

print(f"各类别 IoU: {ious.tolist()}")
print(f"mIoU: {mIoU.item():.4f}")
print(f"各类别 Dice: {dices.tolist()}")
print(f"Mean Dice: {dices.mean().item():.4f}")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
import numpy as np
from pathlib import Path

# ========== 1. 自定义数据集 ==========
class SegmentationDataset(Dataset):
    """自定义分割数据集：图像 + 掩码对"""
    def __init__(self, img_dir, mask_dir, transform=None):
        self.img_paths = sorted(Path(img_dir).glob("*.jpg"))
        self.mask_paths = sorted(Path(mask_dir).glob("*.png"))
        self.transform = transform

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

    def __getitem__(self, idx):
        from PIL import Image
        img = Image.open(self.img_paths[idx]).convert("RGB")
        mask = Image.open(self.mask_paths[idx]).convert("L")
        if self.transform:
            seed = np.random.randint(2147483647)
            torch.manual_seed(seed)
            img = self.transform(img)
            torch.manual_seed(seed)
            mask = self.transform(mask)
        return img, torch.tensor(np.array(mask), dtype=torch.long)

# 数据增强（图像和掩码同步）
train_transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.RandomHorizontalFlip(0.5),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
mask_transform = transforms.Compose([
    transforms.Resize((512, 512), interpolation=Image.NEAREST),
    transforms.ToTensor(),
])

# ========== 2. 模型 + 损失 ==========
from torchvision.models.segmentation import deeplabv3_resnet50

class DiceBCELoss(nn.Module):
    """混合损失: CE + Dice"""
    def __init__(self, alpha=0.5, beta=0.5, num_classes=21):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.ce = nn.CrossEntropyLoss()
        self.num_classes = num_classes

    def forward(self, logits, targets):
        ce_loss = self.ce(logits, targets)
        # Dice Loss
        probs = torch.softmax(logits, dim=1)
        dice = 0.0
        for c in range(self.num_classes):
            p_c = probs[:, c]
            t_c = (targets == c).float()
            intersection = (p_c * t_c).sum()
            dice += 1 - (2 * intersection + 1e-8) / (p_c.sum() + t_c.sum() + 1e-8)
        dice_loss = dice / self.num_classes
        return self.alpha * ce_loss + self.beta * dice_loss

# ========== 3. 训练循环 ==========
def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    for images, masks in loader:
        images = images.to(device)
        masks = masks.to(device).squeeze(1)
        outputs = model(images)["out"]
        loss = criterion(outputs, masks)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)

@torch.no_grad()
def evaluate(model, loader, device, num_classes=21):
    model.eval()
    ious = torch.zeros(num_classes, device=device)
    counts = torch.zeros(num_classes, device=device)
    for images, masks in loader:
        outputs = model(images.to(device))["out"]
        preds = outputs.argmax(dim=1)
        masks = masks.squeeze(1).to(device)
        for c in range(num_classes):
            pred_c = (preds == c)
            mask_c = (masks == c)
            inter = (pred_c & mask_c).sum().float()
            union = (pred_c | mask_c).sum().float()
            if union > 0:
                ious[c] += inter / union
                counts[c] += 1
    valid = counts > 0
    return (ious[valid] / counts[valid]).mean().item() if valid.any() else 0.0

# ========== 4. 主流程 ==========
if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = deeplabv3_resnet50(weights=None, num_classes=21).to(device)
    criterion = DiceBCELoss(num_classes=21)
    optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)

    train_dataset = SegmentationDataset("data/train/images", "data/train/masks", train_transform)
    train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4)

    epochs = 50
    for epoch in range(epochs):
        loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
        print(f"Epoch {epoch+1}/{epochs} | Loss: {loss:.4f}")
        if (epoch + 1) % 10 == 0:
            torch.save(model.state_dict(), f"deeplab_epoch_{epoch+1}.pth")