基于UPerNet模型的视网膜血管语义分割 深度学习医学图像处理Pytorch运行环境

pytorch1.10.0 cpu or gpu都可以
python3.8
代码中有两个主要程序，一个是run.py，直接运行开始训练；一个是gui界面，可以加载训练好的模型，然后选择图片进行语义分割

基于UPerNet模型的视网膜血管语义分割系统。我们将使用PyTorch 1.10.0，并提供两个主要程序：一个是run.py，用于训练模型；另一个是gui.py，用于加载训练好的模型并通过图形用户界面进行语义分割。

1. 环境准备

首先，确保你已经安装了所需的依赖项。你可以使用以下命令安装这些依赖项：

pip install torch==1.10.0 torchvision matplotlib opencv-python

2. 数据集准备

假设你的数据集已经准备好，并且分为训练集和验证集。数据集目录结构如下：

retinal_vessel_dataset/
├── images/
│   ├── train/
│   └── val/
├── masks/
│   ├── train/
│   └── val/

3. UPerNet模型定义

我们将使用UPerNet模型进行语义分割。这里是一个简化的UPerNet模型定义：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50

class UPerNet(nn.Module):
    def __init__(self, num_classes=2):
        super(UPerNet, self).__init__()
        
        # Backbone: ResNet50
        self.backbone = resnet50(pretrained=True)
        self.backbone.fc = nn.Identity()  # Remove the fully connected layer
        
        # PPM (Pyramid Pooling Module)
        self.ppm = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(2048, 512, kernel_size=1, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True)
            ),
            nn.Sequential(
                nn.Conv2d(2048, 512, kernel_size=1, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True),
                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
            ),
            nn.Sequential(
                nn.Conv2d(2048, 512, kernel_size=1, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True),
                nn.Upsample(scale_factor=4, mode='bilinear', align_corners=False)
            ),
            nn.Sequential(
                nn.Conv2d(2048, 512, kernel_size=1, bias=False),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True),
                nn.Upsample(scale_factor=8, mode='bilinear', align_corners=False)
            )
        ])
        
        # Fusion Layer
        self.fusion = nn.Sequential(
            nn.Conv2d(2048 + 512 * 4, 512, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )
        
        # Final Convolution
        self.final_conv = nn.Conv2d(512, num_classes, kernel_size=1)
    
    def forward(self, x):
        # Backbone
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)
        
        c1 = self.backbone.layer1(x)
        c2 = self.backbone.layer2(c1)
        c3 = self.backbone.layer3(c2)
        c4 = self.backbone.layer4(c3)
        
        # PPM
        ppm_out = [c4]
        for pool in self.ppm:
            ppm_out.append(pool(c4))
        
        # Fusion
        fusion_out = self.fusion(torch.cat(ppm_out, dim=1))
        
        # Final Convolution
        out = self.final_conv(fusion_out)
        out = F.interpolate(out, size=x.size()[2:], mode='bilinear', align_corners=False)
        
        return out

4. 训练脚本 (run.py)

import torch
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.transforms import Compose, ToTensor, Normalize
from dataset import RetinalVesselDataset
from model import UPerNet

# Hyperparameters
batch_size = 4
learning_rate = 1e-4
num_epochs = 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Data transforms
transform = Compose([ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

# Datasets
train_dataset = RetinalVesselDataset(root_dir='retinal_vessel_dataset', split='train', transform=transform)
val_dataset = RetinalVesselDataset(root_dir='retinal_vessel_dataset', split='val', transform=transform)

# DataLoaders
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

# Model
model = UPerNet(num_classes=2).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Training loop
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, masks in train_loader:
        images, masks = images.to(device), masks.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")
    
    # Validation
    model.eval()
    with torch.no_grad():
        val_loss = 0.0
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            loss = criterion(outputs, masks)
            val_loss += loss.item()
        
        print(f"Validation Loss: {val_loss/len(val_loader):.4f}")

# Save the model
torch.save(model.state_dict(), 'upernet_retinal_vessel.pth')

5. 数据集类 (dataset.py)

import os
import cv2
import numpy as np
from torch.utils.data import Dataset

class RetinalVesselDataset(Dataset):
    def __init__(self, root_dir, split='train', transform=None):
        self.root_dir = root_dir
        self.split = split
        self.transform = transform
        self.image_paths = sorted(os.listdir(os.path.join(root_dir, 'images', split)))
        self.mask_paths = sorted(os.listdir(os.path.join(root_dir, 'masks', split)))

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

    def __getitem__(self, idx):
        image_path = os.path.join(self.root_dir, 'images', self.split, self.image_paths[idx])
        mask_path = os.path.join(self.root_dir, 'masks', self.split, self.mask_paths[idx])
        
        image = cv2.imread(image_path)
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        
        if self.transform:
            image = self.transform(image)
            mask = torch.tensor(mask, dtype=torch.long)
        
        return image, mask

6. GUI界面 (gui.py)

import tkinter as tk
from tkinter import filedialog
import cv2
import torch
import numpy as np
from torchvision.transforms import Compose, ToTensor, Normalize
from model import UPerNet

# Load the trained model
model = UPerNet(num_classes=2)
model.load_state_dict(torch.load('upernet_retinal_vessel.pth', map_location=torch.device('cpu')))
model.eval()

# Data transforms
transform = Compose([ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

def load_image():
    file_path = filedialog.askopenfilename()
    if file_path:
        image = cv2.imread(file_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        original_image = image.copy()
        
        # Preprocess the image
        image = transform(image).unsqueeze(0)
        
        # Perform inference
        with torch.no_grad():
            output = model(image)
            output = torch.argmax(output.squeeze(), dim=0).numpy()
        
        # Overlay the segmentation mask on the original image
        overlay = original_image.copy()
        mask_color = (0, 255, 0)  # Green color for the vessels
        overlay[output > 0] = mask_color
        alpha = 0.5
        segmented_image = cv2.addWeighted(original_image, 1 - alpha, overlay, alpha, 0)
        
        # Display the result
        cv2.imshow('Segmented Image', segmented_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

# Create the GUI
root = tk.Tk()
root.title("Retinal Vessel Segmentation")
load_button = tk.Button(root, text="Load Image", command=load_image)
load_button.pack(pady=20)
root.mainloop()

7. 运行脚本

1. 训练模型：

   python run.py
   

2. 启动GUI界面：

   python gui.py
   

总结

通过以上步骤，你可以构建一个基于UPerNet模型的视网膜血管语义分割系统。run.py用于训练模型，gui.py用于加载训练好的模型并通过图形用户界面进行语义分割。