使用U-Net模型进行训练风力发电机组件语义分割数据集 检测分割风力发电叶片的分割文章目录使用 **U-Net** 模型进行训练风力发电机组件语义分割数据集 检测分割风力发电叶片的分割一、环境搭建二、数据准备与格式转换数据划分脚本示例三、U-Net模型定义四、数据加载器定义五、训练过程六、模型推理与批量推理七、性能评估风力发电机组件语义分割系统界面设计二、界面功能模块设计以 Streamlit 为例 示例界面布局️ 三、完整代码示例Streamlit PyTorch1. 安装依赖2. 文件结构建议3. utils.py —— 模型定义和工具函数4. app.py —— Streamlit 主程序四、性能评估模块 五、运行应用以下文字及代码仅供参考。风力发电机组件语义分割数据集//6000张////**11类检测目标Hub - 轮毂Trunk - 主干leading_s1 - 导流_s1leading_s2 - 导流_s2leading_s3 - 导流_s3pressure_s1 - 压力_s1pressure_s2 - 压力_s2pressure_s3 - 压力_s3suction_s1 - 吸入_s1suction_s2 - 吸入_s2suction_s3 - 吸入_s31识别和分割风力发电机组件的数据集包含6000张图像和11类检测目标。使用U-Net模型进行训练并详细讲解从数据准备到模型推理及性能评估的整个流程。一、环境搭建首先确保你的环境中已安装必要的库pipinstalltorch torchvision numpy matplotlib scikit-image opencv-python二、数据准备与格式转换假设数据集结构如下wind_turbine_dataset/ ├── images/ # 原始图片 └── masks/ # 对应的标签图每种类别用不同的灰度值表示我们需要将图像和对应的mask划分成训练集和验证集。数据划分脚本示例importosimportrandomimportshutil dataset_dirwind_turbine_datasettrain_ratio0.8images_diros.path.join(dataset_dir,images)masks_diros.path.join(dataset_dir,masks)image_filesos.listdir(images_dir)random.shuffle(image_files)split_idxint(len(image_files)*train_ratio)train_imagesimage_files[:split_idx]val_imagesimage_files[split_idx:]defcopy_files(file_list,src_dir,dst_dir):ifnotos.path.exists(dst_dir):os.makedirs(dst_dir)forfinfile_list:shutil.copy(os.path.join(src_dir,f),os.path.join(dst_dir,f))copy_files(train_images,images_dir,os.path.join(dataset_dir,train_images))copy_files(train_images,masks_dir,os.path.join(dataset_dir,train_masks))copy_files(val_images,images_dir,os.path.join(dataset_dir,val_images))copy_files(val_images,masks_dir,os.path.join(dataset_dir,val_masks))三、U-Net模型定义以下是一个简单的U-Net模型实现importtorch.nnasnnimporttorchclassUNet(nn.Module):def__init__(self,n_channels,n_classes):super(UNet,self).__init__()# 定义网络结构...# 此处简化了U-Net的具体实现请根据需要补充完整defforward(self,x):# 前向传播逻辑...returnx modelUNet(n_channels3,n_classes11).cuda()四、数据加载器定义fromtorch.utils.dataimportDataset,DataLoaderimportcv2importnumpyasnpclassWindTurbineDataset(Dataset):def__init__(self,img_dir,mask_dir,transformNone):self.img_dirimg_dir self.mask_dirmask_dir self.transformtransform self.imagessorted(os.listdir(img_dir))def__len__(self):returnlen(self.images)def__getitem__(self,idx):img_pathos.path.join(self.img_dir,self.images[idx])mask_pathos.path.join(self.mask_dir,self.images[idx])imagecv2.imread(img_path)maskcv2.imread(mask_path,0)ifself.transform:imageself.transform(image)maskself.transform(mask)returnimage,mask train_dsWindTurbineDataset(wind_turbine_dataset/train_images,wind_turbine_dataset/train_masks)val_dsWindTurbineDataset(wind_turbine_dataset/val_images,wind_turbine_dataset/val_masks)train_loaderDataLoader(train_ds,batch_size4,shuffleTrue)val_loaderDataLoader(val_ds,batch_size4)五、训练过程optimizertorch.optim.Adam(model.parameters(),lr0.001)criterionnn.CrossEntropyLoss()forepochinrange(epochs):model.train()running_loss0.0forinputs,labelsintrain_loader:inputs,labelsinputs.cuda(),labels.cuda().long()optimizer.zero_grad()outputsmodel(inputs)losscriterion(outputs,labels)loss.backward()optimizer.step()running_lossloss.item()*inputs.size(0)print(fEpoch{epoch1}/{epochs}, Loss:{running_loss/len(train_ds)})六、模型推理与批量推理model.eval()withtorch.no_grad():forinputs,_inval_loader:inputsinputs.cuda()outputsmodel(inputs)_,predstorch.max(outputs,1)# 进行预测结果处理...七、性能评估可以使用混淆矩阵、准确率、召回率、F1分数等指标来评估模型性能。风力发电机组件语义分割系统界面设计为了实现一个用户友好的风力发电机组件语义分割系统我们将使用Python结合Tkinter或Streamlit构建图形化界面。二、界面功能模块设计以 Streamlit 为例 示例界面布局[上传图像] [开始分割] [查看结果] [批量处理] [模型信息]️ 三、完整代码示例Streamlit PyTorch1. 安装依赖pipinstallstreamlit torch torchvision opencv-python matplotlib pillow2. 文件结构建议wind_turbine_segmentation_app/ ├── app.py # 主程序 ├── unet_model.pth # 训练好的U-Net模型权重 ├── utils.py # 工具函数预处理/后处理 └── sample_images/ # 示例图片3.utils.py—— 模型定义和工具函数importtorchfromtorchimportnnimportcv2importnumpyasnpfromtorchvisionimporttransformsclassUNet(nn.Module):def__init__(self,n_channels3,n_classes11):super(UNet,self).__init__()# 这里替换为你训练的 U-Net 模型结构self.modelnn.Sequential(nn.Conv2d(n_channels,64,kernel_size3,padding1),nn.ReLU(),nn.Conv2d(64,n_classes,kernel_size1))defforward(self,x):returnself.model(x)defload_model(model_pathunet_model.pth):modelUNet()model.load_state_dict(torch.load(model_path))model.eval()returnmodeldefpreprocess_image(image):transformtransforms.Compose([transforms.ToTensor(),transforms.Normalize(mean[0.485,0.456,0.406],std[0.229,0.224,0.225]),])imagecv2.resize(image,(256,256))# 假设输入大小为256x256returntransform(image).unsqueeze(0)defpostprocess_mask(mask_tensor):maskmask_tensor.argmax(dim1).squeeze().cpu().numpy()returnmask4.app.py—— Streamlit 主程序importstreamlitasstimportcv2importnumpyasnpimporttorchfromutilsimportload_model,preprocess_image,postprocess_maskimportmatplotlib.pyplotasplt st.set_page_config(page_title风电组件语义分割系统,layoutwide)st.title(️ 风力发电机组件语义分割系统)# 类别标签CLASS_NAMES{0:Background,1:Hub,2:Trunk,3:leading_s1,4:leading_s2,5:leading_s3,6:pressure_s1,7:pressure_s2,8:pressure_s3,9:suction_s1,10:suction_s2,11:suction_s3}# 加载模型st.cache_resourcedefget_model():returnload_model()modelget_model()# 图像上传uploaded_filest.file_uploader( 上传一张图像,type[jpg,png,jpeg])ifuploaded_fileisnotNone:file_bytesnp.asarray(bytearray(uploaded_file.read()),dtypenp.uint8)imagecv2.imdecode(file_bytes,1)image_rgbcv2.cvtColor(image,cv2.COLOR_BGR2RGB)st.image(image_rgb,caption原始图像,use_column_widthTrue)ifst.button( 开始分割):withst.spinner(正在推理...):input_tensorpreprocess_image(image)withtorch.no_grad():outputmodel(input_tensor)maskpostprocess_mask(output)# 显示分割结果fig,axplt.subplots()ax.imshow(mask,cmaptab20)ax.set_title(分割结果)ax.axis(off)st.pyplot(fig)# 显示类别图例st.markdown(### 类别说明)foridx,nameinCLASS_NAMES.items():st.markdown(f- **{name}**: {idx})# 批量处理st.sidebar.header(⚙️ 功能选项)batch_modest.sidebar.checkbox( 批量处理模式)ifbatch_mode:st.sidebar.info(请将所有图像放入指定文件夹中点击下方按钮进行批量推理。)ifst.sidebar.button( 执行批量推理):# 示例逻辑需根据实际路径调整st.sidebar.success(✅ 批量推理完成可在指定路径获取结果。)# 模型信息withst.sidebar.expander(ℹ️ 模型信息):st.write(模型类型U-Net)st.write(类别数量11类)st.write(输入尺寸256x256)st.write(训练框架PyTorch)四、性能评估模块你可以添加一个页面来展示模型在验证集上的表现如混淆矩阵、IoU、Dice Score 等。fromsklearn.metricsimportconfusion_matrix,ConfusionMatrixDisplay# 示例混淆矩阵展示cmnp.random.randint(0,100,size(11,11))fig,axplt.subplots()ConfusionMatrixDisplay(confusion_matrixcm,display_labelsCLASS_NAMES.values()).plot(axax)plt.xticks(rotation45)st.pyplot(fig) 五、运行应用streamlit run app.py