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

医学语义分割类-基于UPerNet模型的视网膜血管语义分割深度学习医学图像处理视觉眼睛视网膜血管语义分割

news 2026/5/11 11:55:55

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

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

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

1. 环境准备

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

pipinstalltorch==1.10.0 torchvision matplotlib opencv-python

2. 数据集准备

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

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

3. UPerNet模型定义

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

importtorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorchvision.modelsimportresnet50classUPerNet(nn.Module):def__init__(self,num_classes=2):super(UPerNet,self).__init__()# Backbone: ResNet50self.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 Layerself.fusion=nn.Sequential(nn.Conv2d(2048+512*4,512,kernel_size=3,padding=1,bias=False),nn.BatchNorm2d(512),nn.ReLU(inplace=True))# Final Convolutionself.final_conv=nn.Conv2d(512,num_classes,kernel_size=1)defforward(self,x):# Backbonex=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)# PPMppm_out=[c4]forpoolinself.ppm:ppm_out.append(pool(c4))# Fusionfusion_out=self.fusion(torch.cat(ppm_out,dim=1))# Final Convolutionout=self.final_conv(fusion_out)out=F.interpolate(out,size=x.size()[2:],mode='bilinear',align_corners=False)returnout

4. 训练脚本 (`run.py`)

importtorchimporttorch.optimasoptimimporttorch.nn.functionalasFfromtorch.utils.dataimportDataLoaderfromtorchvision.transformsimportCompose,ToTensor,NormalizefromdatasetimportRetinalVesselDatasetfrommodelimportUPerNet# Hyperparametersbatch_size=4learning_rate=1e-4num_epochs=100device=torch.device("cuda"iftorch.cuda.is_available()else"cpu")# Data transformstransform=Compose([ToTensor(),Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])])# Datasetstrain_dataset=RetinalVesselDataset(root_dir='retinal_vessel_dataset',split='train',transform=transform)val_dataset=RetinalVesselDataset(root_dir='retinal_vessel_dataset',split='val',transform=transform)# DataLoaderstrain_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)# Modelmodel=UPerNet(num_classes=2).to(device)# Loss and optimizercriterion=nn.CrossEntropyLoss()optimizer=optim.Adam(model.parameters(),lr=learning_rate)# Training loopforepochinrange(num_epochs):model.train()running_loss=0.0forimages,masksintrain_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}")# Validationmodel.eval()withtorch.no_grad():val_loss=0.0forimages,masksinval_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 modeltorch.save(model.state_dict(),'upernet_retinal_vessel.pth')

5. 数据集类 (`dataset.py`)

importosimportcv2importnumpyasnpfromtorch.utils.dataimportDatasetclassRetinalVesselDataset(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):returnlen(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)ifself.transform:image=self.transform(image)mask=torch.tensor(mask,dtype=torch.long)returnimage,mask

6. GUI界面 (`gui.py`)

importtkinterastkfromtkinterimportfiledialogimportcv2importtorchimportnumpyasnpfromtorchvision.transformsimportCompose,ToTensor,NormalizefrommodelimportUPerNet# Load the trained modelmodel=UPerNet(num_classes=2)model.load_state_dict(torch.load('upernet_retinal_vessel.pth',map_location=torch.device('cpu')))model.eval()# Data transformstransform=Compose([ToTensor(),Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])])defload_image():file_path=filedialog.askopenfilename()iffile_path:image=cv2.imread(file_path)image=cv2.cvtColor(image,cv2.COLOR_BGR2RGB)original_image=image.copy()# Preprocess the imageimage=transform(image).unsqueeze(0)# Perform inferencewithtorch.no_grad():output=model(image)output=torch.argmax(output.squeeze(),dim=0).numpy()# Overlay the segmentation mask on the original imageoverlay=original_image.copy()mask_color=(0,255,0)# Green color for the vesselsoverlay[output>0]=mask_color alpha=0.5segmented_image=cv2.addWeighted(original_image,1-alpha,overlay,alpha,0)# Display the resultcv2.imshow('Segmented Image',segmented_image)cv2.waitKey(0)cv2.destroyAllWindows()# Create the GUIroot=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()