当前位置：首页 > news >正文

自监督骨干（DINOv2）用于内镜分割与跟踪的全面实现分析

news 2025/11/7 7:10:07

在这里插入图片描述

1. 引言：内镜图像处理的挑战与机遇

内镜检查是现代医学诊断的关键工具，但内镜图像分析面临诸多独特挑战：

视觉干扰：镜面反光、烟雾/血污、液体遮挡
动态变化：快速相机移动、视野缩放、器械运动
域偏移：不同设备、光照条件导致的颜色变化
标注稀缺：医疗数据标注成本高，专家依赖性强

自监督学习（SSL）为解决这些问题提供了新途径。DINOv2等模型在无标签数据上学习到的视觉特征具有强鲁棒性，特别适合医学图像处理任务。

为什么选择DINOv2？

在ImageNet上预训练的视觉特征可直接迁移
对光照、纹理变化不敏感
适合小样本医疗场景
提供多尺度特征表示

2. 自监督学习与DINOv2基础

2.1 DINOv2核心机制

DINOv2通过自蒸馏学习视觉表示：

使用Vision Transformer架构
通过对比学习优化特征空间
不依赖人工标注，直接从图像中学习

2.2 关键优势

特征一致性：相似对象特征空间接近
尺度不变性：对图像缩放鲁棒
干扰抵抗：对遮挡、光照变化稳定

# DINOv2基础加载示例
import torch
import torchvision.transforms as T
from PIL import Imagedef load_dinov2(model_name="dinov2_vitl14"):"""加载预训练DINOv2模型"""device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = torch.hub.load('facebookresearch/dinov2', model_name)model.eval().to(device)# 标准预处理transform = T.Compose([T.Resize(256),T.CenterCrop(224),T.ToTensor(),T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])return model, transform, device

3. 系统架构设计

3.1 整体流程图

[输入视频] → [帧预处理] → [DINOv2特征提取] → ├─[分割网络] → [初始掩膜]└─[光流估计] → [掩膜传播] → [时序融合] → [最终输出]

3.2 核心组件

特征提取器：冻结的DINOv2骨干
分割头：轻量级上采样解码器
运动估计：RAFT光流网络
时序融合：特征匹配与掩膜优化

3.3 创新点

双分支架构（静态分割+动态传播）
领域特定数据增强
自适应关键帧选择机制

4. 分割模块详细实现

4.1 数据加载与预处理

import torch
from torch.utils.data import Dataset, DataLoader
import albumentations as A
import cv2
import numpy as np
import glob
import osclass EndoscopicDataset(Dataset):"""内镜图像数据集，支持增强和预处理"""def __init__(self, img_dir, mask_dir=None, transform=None, size=512):self.img_paths = sorted(glob.glob(os.path.join(img_dir, '*')))self.mask_paths = sorted(glob.glob(os.path.join(mask_dir, '*'))) if mask_dir else Noneself.transform = transformself.size = size# 内镜特定增强if transform is None:self.transform = A.Compose([A.LongestMaxSize(max_size=size),A.PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT, value=0),A.HorizontalFlip(p=0.5),A.OneOf([A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2),], p=0.8),A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),A.OneOf([A.GaussianBlur(blur_limit=(3, 5)),A.MedianBlur(blur_limit=5),], p=0.2),# 内镜特定增强A.RandomSunFlare(num_flare_circles_lower=0, num_flare_circles_upper=2, src_radius=100, p=0.1),A.GridDistortion(num_steps=5, distort_limit=0.1, p=0.2),])def __len__(self):return len(self.img_paths)def __getitem__(self, idx):img = cv2.imread(self.img_paths[idx])img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)if self.mask_paths:mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE)mask = (mask > 127).astype(np.float32)if self.transform:augmented = self.transform(image=img, mask=mask)img, mask = augmented['image'], augmented['mask']# 标准化img = img.astype(np.float32) / 255.0mean = np.array([0.485, 0.456, 0.406])std = np.array([0.229, 0.224, 0.225])img = (img - mean) / stdimg = np.transpose(img, (2, 0, 1))return torch.tensor(img, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32).unsqueeze(0)else:if self.transform:img = self.transform(image=img)['image']img = img.astype(np.float32) / 255.0mean = np.array([0.485, 0.456, 0.406])std = np.array([0.229, 0.224, 0.225])img = (img - mean) / stdimg = np.transpose(img, (2, 0, 1))return torch.tensor(img, dtype=torch.float32)

4.2 分割网络架构

import torch.nn as nn
import torch.nn.functional as F
import timmclass DINOv2SegNet(nn.Module):"""基于DINOv2的分割网络"""def __init__(self, backbone_name='vit_large_patch14_dinov2', num_classes=1):super().__init__()# 加载DINOv2骨干self.backbone = timm.create_model(backbone_name, pretrained=True, features_only=True,out_indices=[3, 6, 9, 12]  # 提取不同层级特征)# 获取特征通道数feature_channels = self.backbone.feature_info.channels()print(f"Feature channels: {feature_channels}")# 特征融合模块self.fusion_conv = nn.ModuleList([nn.Conv2d(ch, 256, kernel_size=1) for ch in feature_channels])# 解码器路径self.decoder = nn.ModuleList([nn.ConvTranspose2d(256, 256, kernel_size=4, stride=2, padding=1),nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)])# 最终分类头self.head = nn.Sequential(nn.Conv2d(64, 32, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(32, num_classes, kernel_size=1))# 特征注意力self.attention = nn.ModuleList([SEBlock(256) for _ in range(len(feature_channels))])# 初始化self._init_weights()def _init_weights(self):"""权重初始化"""for m in self.modules():if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')if m.bias is not None:nn.init.constant_(m.bias, 0)def forward(self, x):# 获取多尺度特征features = self.backbone(x)# 特征处理processed_feats = []for i, feat in enumerate(features):feat = self.fusion_conv[i](feat)feat = self.attention[i](feat)processed_feats.append(feat)# 特征融合 - 从深到浅x = processed_feats[-1]  # 最深层特征# 解码路径for i, deconv in enumerate(self.decoder):x = deconv(x)if i < len(processed_feats) - 1:# 跳跃连接skip_feat = processed_feats[-(i+2)]if x.shape[2:] != skip_feat.shape[2:]:x = F.interpolate(x, size=skip_feat.shape[2:], mode='bilinear', align_corners=False)x = x + skip_featx = F.relu(x, inplace=True)# 最终输出out = self.head(x)# 上采样到输入尺寸out = F.interpolate(out, size=x.shape[-2:], mode='bilinear', align_corners=False)return outclass SEBlock(nn.Module):"""Squeeze-and-Excitation注意力模块"""def __init__(self, channel, reduction=16):super().__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.fc = nn.Sequential(nn.Linear(channel, channel // reduction, bias=False),nn.ReLU(inplace=True),nn.Linear(channel // reduction, channel, bias=False),nn.Sigmoid())def forward(self, x):b, c, _, _ = x.size()y = self.avg_pool(x).view(b, c)y = self.fc(y).view(b, c, 1, 1)return x * y.expand_as(x)

4.3 损失函数设计

class CombinedLoss(nn.Module):"""组合损失函数：Dice + BCE + 边界损失"""def __init__(self, weight_dice=1.0, weight_bce=1.0, weight_boundary=0.5):super().__init__()self.weight_dice = weight_diceself.weight_bce = weight_bceself.weight_boundary = weight_boundarydef dice_loss(self, pred, target, eps=1e-6):"""Dice损失"""pred = torch.sigmoid(pred)intersection = (pred * target).sum(dim=(2,3))union = pred.sum(dim=(2,3)) + target.sum(dim=(2,3))dice = (2. * intersection + eps) / (union + eps)return 1. - dice.mean()def boundary_loss(self, pred, target):"""边界损失"""pred = torch.sigmoid(pred)# 计算边界kernel = torch.ones(1, 1, 3, 3, device=pred.device)target_dilated = F.conv2d(target, kernel, padding=1) > 0target_eroded = F.conv2d(target, kernel, padding=1) == 9boundary = target_dilated & ~target_erodedif boundary.sum() == 0:return torch.tensor(0., device=pred.device)return F.mse_loss(pred[boundary], target[boundary].float())def forward(self, pred, target):dice = self.dice_loss(pred, target)bce = F.binary_cross_entropy_with_logits(pred, target)boundary = self.boundary_loss(pred, target)return (self.weight_dice * dice + self.weight_bce * bce + self.weight_boundary * boundary)

4.4 训练循环实现

def train_segmentation(model, dataloader, val_loader, device, epochs=50, lr=1e-4):"""分割模型训练函数"""# 优化器 - 分层学习率backbone_params = []decoder_params = []for name, param in model.named_parameters():if 'backbone' in name:backbone_params.append(param)else:decoder_params.append(param)optimizer = torch.optim.AdamW([{'params': backbone_params, 'lr': lr * 0.1},  # 骨干学习率更低{'params': decoder_params, 'lr': lr}], weight_decay=1e-4)# 学习率调度器scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=lr*0.01)# 损失函数criterion = CombinedLoss()# 训练历史train_history = {'loss': [], 'dice': []}val_history = {'loss': [], 'dice': []}best_val_dice = 0.0for epoch in range(epochs):model.train()epoch_loss = 0.0epoch_dice = 0.0# 训练阶段for imgs, masks in dataloader:imgs, masks = imgs.to(device), masks.to(device)optimizer.zero_grad()outputs = model(imgs)loss = criterion(outputs, masks)loss.backward()# 梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)optimizer.step()epoch_loss += loss.item()# 计算Dice分数with torch.no_grad():preds = torch.sigmoid(outputs) > 0.5dice = (2 * (preds * masks).sum() + 1e-6) / (preds.sum() + masks.sum() + 1e-6)epoch_dice += dice.item()# 验证阶段val_loss, val_dice = evaluate_segmentation(model, val_loader, criterion, device)# 更新学习率scheduler.step()# 记录历史train_history['loss'].append(epoch_loss / len(dataloader))train_history['dice'].append(epoch_dice / len(dataloader))val_history['loss'].append(val_loss)val_history['dice'].append(val_dice)print(f"Epoch {epoch+1}/{epochs} - "f"Train Loss: {train_history['loss'][-1]:.4f}, "f"Train Dice: {train_history['dice'][-1]:.4f}, "f"Val Loss: {val_loss:.4f}, "f"Val Dice: {val_dice:.4f}")# 保存最佳模型if val_dice > best_val_dice:best_val_dice = val_dicetorch.save({'epoch': epoch,'model_state_dict': model.state_dict(),'optimizer_state_dict': optimizer.state_dict(),'dice': val_dice,}, 'best_seg_model.pth')return train_history, val_historydef evaluate_segmentation(model, dataloader, criterion, device):"""分割模型评估"""model.eval()total_loss = 0.0total_dice = 0.0with torch.no_grad():for imgs, masks in dataloader:imgs, masks = imgs.to(device), masks.to(device)outputs = model(imgs)loss = criterion(outputs, masks)total_loss += loss.item()# 计算Dicepreds = torch.sigmoid(outputs) > 0.5dice = (2 * (preds * masks).sum() + 1e-6) / (preds.sum() + masks.sum() + 1e-6)total_dice += dice.item()return total_loss / len(dataloader), total_dice / len(dataloader)

5. 跟踪与掩膜传播模块

5.1 光流估计集成

# 首先安装RAFT: pip install raft-pytorch
from raft import RAFT
from utils.utils import InputPadderclass OpticalFlowEstimator:"""RAFT光流估计器封装"""def __init__(self, model_path='raft/models/raft-things.pth', device='cuda'):self.device = deviceself.