TransUNet遥感河流分割项目 pytorch模型文章目录TransUNet遥感河流分割项目 pytorch模型一、环境搭建1. 安装依赖项2. 确认PyTorch版本二、数据准备数据集结构数据预处理三、模型训练TransUNet模型定义四、构建GUI界面使用PyQt5创建GUI界面运行项目代码示例仅供参考实现原理*项目使用数据集有训练样本共82个测试样本8个原图是三通道彩色图像label是二值化(0与255)掩模图像代码运行环境pytorch1.10.0 cpu or gpu都可python3.8可用pip install -r命令安装基于TransUNet的遥感河流分割项目从环境搭建、数据准备、模型训练到GUI界面开发逐步进行。代码示例一、环境搭建1. 安装依赖项确保你的系统已安装Python 3.8并使用requirements.txt文件安装必要的依赖项。pipinstall-rrequirements.txt2. 确认PyTorch版本确保你安装了正确的PyTorch版本1.10.0。python-cimport torch; print(torch.__version__)二、数据准备假设你已经有82个训练样本和8个测试样本的数据集并且这些图像已经按照标准格式组织。数据集结构确保数据集按如下结构组织river_segmentation_dataset/ ├── images/ │ ├── train/ │ └── test/ └── labels/ ├── train/ └── test/数据预处理编写一个简单的脚本来加载和预处理数据。importosimportcv2importnumpyasnpfromtorch.utils.dataimportDataset,DataLoaderfromtorchvisionimporttransformsclassRiverDataset(Dataset):def__init__(self,image_dir,label_dir,transformNone):self.image_dirimage_dir self.label_dirlabel_dir self.transformtransform self.imagesos.listdir(image_dir)def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_pathos.path.join(self.image_dir,self.images[idx])label_pathos.path.join(self.label_dir,self.images[idx].replace(.jpg,.png))imagecv2.imread(img_path)labelcv2.imread(label_path,cv2.IMREAD_GRAYSCALE)label(label0).astype(np.float32)# Convert to binary maskifself.transform:imageself.transform(image)labelself.transform(label)returnimage,label# Define transformationstransformtransforms.Compose([transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225])])# Create datasets and dataloaderstrain_datasetRiverDataset(river_segmentation_dataset/images/train,river_segmentation_dataset/labels/train,transformtransform)test_datasetRiverDataset(river_segmentation_dataset/images/test,river_segmentation_dataset/labels/test,transformtransform)train_loaderDataLoader(train_dataset,batch_size4,shuffleTrue)test_loaderDataLoader(test_dataset,batch_size4,shuffleFalse)三、模型训练TransUNet模型定义importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassTransUNet(nn.Module):def__init__(self):super(TransUNet,self).__init__()# Define your TransUNet architecture herepassdefforward(self,x):# Implement the forward passpass# Initialize the modelmodelTransUNet()iftorch.cuda.is_available():modelmodel.cuda()# Loss function and optimizercriterionnn.BCEWithLogitsLoss()optimizertorch.optim.Adam(model.parameters(),lr0.001)# Training loopnum_epochs100forepochinrange(num_epochs):fori,(images,labels)inenumerate(train_loader):iftorch.cuda.is_available():imagesimages.cuda()labelslabels.cuda()outputsmodel(images)losscriterion(outputs,labels.unsqueeze(1))optimizer.zero_grad()loss.backward()optimizer.step()if(i1)%100:print(fEpoch [{epoch1}/{num_epochs}], Step [{i1}/{len(train_loader)}], Loss:{loss.item():.4f})四、构建GUI界面使用PyQt5创建GUI界面importsysfromPyQt5.QtWidgetsimportQApplication,QMainWindow,QLabel,QPushButton,QFileDialog,QVBoxLayout,QWidgetfromPyQt5.QtGuiimportQPixmapimportcv2importnumpyasnpimporttorchclassMainWindow(QMainWindow):def__init__(self):super().__init__()self.setWindowTitle(Remote Sensing Image River Segmentation)self.setGeometry(100,100,800,600)self.central_widgetQWidget()self.setCentralWidget(self.central_widget)self.layoutQVBoxLayout(self.central_widget)self.upload_buttonQPushButton(Upload Image,self)self.upload_button.clicked.connect(self.upload_image)self.layout.addWidget(self.upload_button)self.image_labelQLabel(self)self.layout.addWidget(self.image_label)self.segment_buttonQPushButton(Perform Segmentation,self)self.segment_button.clicked.connect(self.perform_segmentation)self.layout.addWidget(self.segment_button)self.result_labelQLabel(self)self.layout.addWidget(self.result_label)self.modeltorch.load(path_to_your_trained_model.pth)# Load your trained modeldefupload_image(self):fname,_QFileDialog.getOpenFileName(self,Open file,,Image files (*.jpg *.png))iffname:self.image_pathfname pixmapQPixmap(fname)self.image_label.setPixmap(pixmap)defperform_segmentation(self):imagecv2.imread(self.image_path)imagecv2.resize(image,(256,256))# Resize to match input size of your modelimageimage/255.0# Normalizeimagenp.transpose(image,(2,0,1))# Change order from HWC to CHWimagetorch.tensor(image,dtypetorch.float32).unsqueeze(0)iftorch.cuda.is_available():imageimage.cuda()self.modelself.model.cuda()withtorch.no_grad():outputself.model(image)outputtorch.sigmoid(output)outputoutput.cpu().numpy()[0][0]# Get the first channel of the first batch itemoutput(output0.5).astype(np.uint8)*255# Threshold and convert to binary maskresultcv2.cvtColor(output,cv2.COLOR_GRAY2BGR)resultcv2.addWeighted(cv2.imread(self.image_path),0.7,result,0.3,0)# Overlay on original imageheight,width,channelresult.shape bytes_per_line3*width q_imgQImage(result.data,width,height,bytes_per_line,QImage.Format_RGB888).rgbSwapped()self.result_label.setPixmap(QPixmap.fromImage(q_img))if__name____main__:appQApplication(sys.argv)windowMainWindow()window.show()sys.exit(app.exec_())运行项目保存上述代码为gui.py然后运行它python gui.py启动一个窗口应用程序允许用户上传遥感图像并进行河流分割。记得在实际部署前对模型进行充分的测试与调优。