基于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-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_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)returnout4. 训练脚本 (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,mask6. 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()7. 运行脚本
训练模型：
python run.py启动GUI界面：
python gui.py
总结
通过以上步骤，你可以构建一个基于UPerNet模型的视网膜血管语义分割系统。run.py用于训练模型，gui.py用于加载训练好的模型并通过图形用户界面进行语义分割。
