基于UPerNet模型的视网膜血管语义分割 深度学习医学图像处理Pytorch运行环境pytorch1.10.0 cpu or gpu都可以python3.8代码中有两个主要程序一个是run.py直接运行开始训练一个是gui界面可以加载训练好的模型然后选择图片进行语义分割基于UPerNet模型的视网膜血管语义分割系统。我们将使用PyTorch 1.10.0并提供两个主要程序一个是run.py用于训练模型另一个是gui.py用于加载训练好的模型并通过图形用户界面进行语义分割。1. 环境准备首先确保你已经安装了所需的依赖项。你可以使用以下命令安装这些依赖项pipinstalltorch1.10.0 torchvision matplotlib opencv-python2. 数据集准备假设你的数据集已经准备好并且分为训练集和验证集。数据集目录结构如下retinal_vessel_dataset/ ├── images/ │ ├── train/ │ └── val/ ├── masks/ │ ├── train/ │ └── val/3. UPerNet模型定义我们将使用UPerNet模型进行语义分割。这里是一个简化的UPerNet模型定义importtorchimporttorch.nnasnnimporttorch.nn.functionalasFfromtorchvision.modelsimportresnet50classUPerNet(nn.Module):def__init__(self,num_classes2):super(UPerNet,self).__init__()# Backbone: ResNet50self.backboneresnet50(pretrainedTrue)self.backbone.fcnn.Identity()# Remove the fully connected layer# PPM (Pyramid Pooling Module)self.ppmnn.ModuleList([nn.Sequential(nn.Conv2d(2048,512,kernel_size1,biasFalse),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue)),nn.Sequential(nn.Conv2d(2048,512,kernel_size1,biasFalse),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue),nn.Upsample(scale_factor2,modebilinear,align_cornersFalse)),nn.Sequential(nn.Conv2d(2048,512,kernel_size1,biasFalse),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue),nn.Upsample(scale_factor4,modebilinear,align_cornersFalse)),nn.Sequential(nn.Conv2d(2048,512,kernel_size1,biasFalse),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue),nn.Upsample(scale_factor8,modebilinear,align_cornersFalse))])# Fusion Layerself.fusionnn.Sequential(nn.Conv2d(2048512*4,512,kernel_size3,padding1,biasFalse),nn.BatchNorm2d(512),nn.ReLU(inplaceTrue))# Final Convolutionself.final_convnn.Conv2d(512,num_classes,kernel_size1)defforward(self,x):# Backbonexself.backbone.conv1(x)xself.backbone.bn1(x)xself.backbone.relu(x)xself.backbone.maxpool(x)c1self.backbone.layer1(x)c2self.backbone.layer2(c1)c3self.backbone.layer3(c2)c4self.backbone.layer4(c3)# PPMppm_out[c4]forpoolinself.ppm:ppm_out.append(pool(c4))# Fusionfusion_outself.fusion(torch.cat(ppm_out,dim1))# Final Convolutionoutself.final_conv(fusion_out)outF.interpolate(out,sizex.size()[2:],modebilinear,align_cornersFalse)returnout4. 训练脚本 (run.py)importtorchimporttorch.optimasoptimimporttorch.nn.functionalasFfromtorch.utils.dataimportDataLoaderfromtorchvision.transformsimportCompose,ToTensor,NormalizefromdatasetimportRetinalVesselDatasetfrommodelimportUPerNet# Hyperparametersbatch_size4learning_rate1e-4num_epochs100devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)# Data transformstransformCompose([ToTensor(),Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# Datasetstrain_datasetRetinalVesselDataset(root_dirretinal_vessel_dataset,splittrain,transformtransform)val_datasetRetinalVesselDataset(root_dirretinal_vessel_dataset,splitval,transformtransform)# DataLoaderstrain_loaderDataLoader(train_dataset,batch_sizebatch_size,shuffleTrue,num_workers4)val_loaderDataLoader(val_dataset,batch_sizebatch_size,shuffleFalse,num_workers4)# ModelmodelUPerNet(num_classes2).to(device)# Loss and optimizercriterionnn.CrossEntropyLoss()optimizeroptim.Adam(model.parameters(),lrlearning_rate)# Training loopforepochinrange(num_epochs):model.train()running_loss0.0forimages,masksintrain_loader:images,masksimages.to(device),masks.to(device)optimizer.zero_grad()outputsmodel(images)losscriterion(outputs,masks)loss.backward()optimizer.step()running_lossloss.item()print(fEpoch [{epoch1}/{num_epochs}], Loss:{running_loss/len(train_loader):.4f})# Validationmodel.eval()withtorch.no_grad():val_loss0.0forimages,masksinval_loader:images,masksimages.to(device),masks.to(device)outputsmodel(images)losscriterion(outputs,masks)val_lossloss.item()print(fValidation 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,splittrain,transformNone):self.root_dirroot_dir self.splitsplit self.transformtransform self.image_pathssorted(os.listdir(os.path.join(root_dir,images,split)))self.mask_pathssorted(os.listdir(os.path.join(root_dir,masks,split)))def__len__(self):returnlen(self.image_paths)def__getitem__(self,idx):image_pathos.path.join(self.root_dir,images,self.split,self.image_paths[idx])mask_pathos.path.join(self.root_dir,masks,self.split,self.mask_paths[idx])imagecv2.imread(image_path)maskcv2.imread(mask_path,cv2.IMREAD_GRAYSCALE)ifself.transform:imageself.transform(image)masktorch.tensor(mask,dtypetorch.long)returnimage,mask6. GUI界面 (gui.py)importtkinterastkfromtkinterimportfiledialogimportcv2importtorchimportnumpyasnpfromtorchvision.transformsimportCompose,ToTensor,NormalizefrommodelimportUPerNet# Load the trained modelmodelUPerNet(num_classes2)model.load_state_dict(torch.load(upernet_retinal_vessel.pth,map_locationtorch.device(cpu)))model.eval()# Data transformstransformCompose([ToTensor(),Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])defload_image():file_pathfiledialog.askopenfilename()iffile_path:imagecv2.imread(file_path)imagecv2.cvtColor(image,cv2.COLOR_BGR2RGB)original_imageimage.copy()# Preprocess the imageimagetransform(image).unsqueeze(0)# Perform inferencewithtorch.no_grad():outputmodel(image)outputtorch.argmax(output.squeeze(),dim0).numpy()# Overlay the segmentation mask on the original imageoverlayoriginal_image.copy()mask_color(0,255,0)# Green color for the vesselsoverlay[output0]mask_color alpha0.5segmented_imagecv2.addWeighted(original_image,1-alpha,overlay,alpha,0)# Display the resultcv2.imshow(Segmented Image,segmented_image)cv2.waitKey(0)cv2.destroyAllWindows()# Create the GUIroottk.Tk()root.title(Retinal Vessel Segmentation)load_buttontk.Button(root,textLoad Image,commandload_image)load_button.pack(pady20)root.mainloop()7. 运行脚本训练模型python run.py启动GUI界面python gui.py总结通过以上步骤你可以构建一个基于UPerNet模型的视网膜血管语义分割系统。run.py用于训练模型gui.py用于加载训练好的模型并通过图形用户界面进行语义分割。