使用深度学习框架如PyTorch/YOLO训练指针式仪表的目标检测和语义分割模型并建立推理代码及可视化界面训练指针式仪表数据集 检测和分割数据集文章目录使用深度学习框架如PyTorch/YOLO训练指针式仪表的目标检测和语义分割模型并建立推理代码及可视化界面训练指针式仪表数据集 检测和分割数据集**任务概述****步骤 1环境准备****步骤 2数据集处理****目标检测数据集****语义分割数据集****步骤 3目标检测训练代码****YOLOv8 训练代码****自定义目标检测模型PyTorch Faster R-CNN****步骤 4语义分割训练代码****DeepLabV3 训练代码****步骤 5推理代码与可视化界面****推理代码****可视化界面****总结**3000多张指针式仪表目标检测语义分割数据集和标注文件11使用深度学习框架如PyTorch或YOLO训练指针式仪表的目标检测和语义分割模型并建立推理代码及可视化界面。任务概述目标目标检测识别仪表盘上的指针位置。语义分割对仪表盘的各个部件进行像素级分类。数据集3000多张图像已标注。标注文件可能是边界框目标检测或像素掩码语义分割。工具目标检测YOLOv5/v8 或 PyTorch Faster R-CNN。语义分割PyTorch DeepLabV3 或 MMSegmentation。可视化界面使用Tkinter或Streamlit构建图形化界面。步骤 1环境准备确保安装了必要的库。# 安装依赖pipinstalltorch torchvision matplotlib opencv-python ultralytics streamlit步骤 2数据集处理目标检测数据集格式YOLO需要每个图像对应一个txt文件包含目标的类别和归一化的边界框坐标。目录结构dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ └── labels/ ├── train/ ├── val/ └── test/语义分割数据集格式每个图像对应一个掩码图像掩码图像的像素值表示类别索引。目录结构dataset/ ├── images/ │ ├── train/ │ ├── val/ │ └── test/ └── masks/ ├── train/ ├── val/ └── test/步骤 3目标检测训练代码YOLOv8 训练代码fromultralyticsimportYOLO# 加载预训练模型modelYOLO(yolov8n.pt)# 使用YOLOv8 Nano模型# 配置数据集data_yamldataset.yaml# 开始训练resultsmodel.train(datadata_yaml,epochs50,# 训练轮数imgsz640,# 输入图像尺寸batch16,# 批量大小devicecuda# 使用GPU)# 验证模型metricsmodel.val()# 保存模型model.export(formatonnx)# 导出为ONNX格式自定义目标检测模型PyTorch Faster R-CNNimporttorchimporttorchvisionfromtorchvision.models.detection.faster_rcnnimportFastRCNNPredictorfromtorch.utils.dataimportDataLoaderfromtorchvision.transformsimportToTensorfromPILimportImageimportos# 自定义数据集类classPointerDataset(torch.utils.data.Dataset):def__init__(self,root,transformsNone):self.rootroot self.transformstransforms self.imageslist(sorted(os.listdir(os.path.join(root,images))))self.labelslist(sorted(os.listdir(os.path.join(root,labels))))def__getitem__(self,idx):img_pathos.path.join(self.root,images,self.images[idx])label_pathos.path.join(self.root,labels,self.labels[idx])imageImage.open(img_path).convert(RGB)withopen(label_path,r)asf:boxes[]forlineinf.readlines():class_id,x_center,y_center,width,heightmap(float,line.strip().split())x_min(x_center-width/2)*image.width x_max(x_centerwidth/2)*image.width y_min(y_center-height/2)*image.height y_max(y_centerheight/2)*image.height boxes.append([x_min,y_min,x_max,y_max])boxestorch.as_tensor(boxes,dtypetorch.float32)labelstorch.ones((len(boxes),),dtypetorch.int64)# 假设所有类别为1target{boxes:boxes,labels:labels}ifself.transforms:imageself.transforms(image)returnimage,targetdef__len__(self):returnlen(self.images)# 创建模型defget_model(num_classes):modeltorchvision.models.detection.fasterrcnn_resnet50_fpn(pretrainedTrue)in_featuresmodel.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictorFastRCNNPredictor(in_features,num_classes)returnmodel# 数据加载器datasetPointerDataset(rootdataset/train,transformsToTensor())data_loaderDataLoader(dataset,batch_size2,shuffleTrue,num_workers4)# 训练模型devicetorch.device(cuda)iftorch.cuda.is_available()elsetorch.device(cpu)modelget_model(num_classes5)# 包括背景类model.to(device)optimizertorch.optim.SGD(model.parameters(),lr0.005,momentum0.9,weight_decay0.0005)forepochinrange(10):# 训练10轮model.train()forimages,targetsindata_loader:imageslist(image.to(device)forimageinimages)targets[{k:v.to(device)fork,vint.items()}fortintargets]loss_dictmodel(images,targets)lossessum(lossforlossinloss_dict.values())optimizer.zero_grad()losses.backward()optimizer.step()print(fEpoch{epoch1}, Loss:{losses.item()})步骤 4语义分割训练代码DeepLabV3 训练代码importtorchimporttorch.nnasnnimporttorchvisionfromtorchvisionimporttransformsfromtorch.utils.dataimportDataLoaderfromPILimportImageimportos# 自定义数据集类classSegmentationDataset(torch.utils.data.Dataset):def__init__(self,image_dir,mask_dir,transformNone):self.image_dirimage_dir self.mask_dirmask_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])mask_pathos.path.join(self.mask_dir,self.images[idx].replace(.jpg,.png))imageImage.open(img_path).convert(RGB)maskImage.open(mask_path)ifself.transform:imageself.transform(image)maskself.transform(mask)returnimage,mask# 数据加载器transformtransforms.Compose([transforms.Resize((256,256)),transforms.ToTensor()])datasetSegmentationDataset(image_dirdataset/images/train,mask_dirdataset/masks/train,transformtransform)data_loaderDataLoader(dataset,batch_size8,shuffleTrue)# 创建模型modeltorchvision.models.segmentation.deeplabv3_resnet50(pretrainedTrue)model.classifier[-1]nn.Conv2d(256,5,kernel_size(1,1))# 修改输出类别数model.to(torch.device(cuda))# 训练模型criterionnn.CrossEntropyLoss()optimizertorch.optim.Adam(model.parameters(),lr0.001)forepochinrange(10):model.train()forimages,masksindata_loader:images,masksimages.to(cuda),masks.to(cuda)outputsmodel(images)[out]losscriterion(outputs,masks.long())optimizer.zero_grad()loss.backward()optimizer.step()print(fEpoch{epoch1}, Loss:{loss.item()})步骤 5推理代码与可视化界面推理代码importcv2importtorchfromultralyticsimportYOLO# 加载模型modelYOLO(runs/detect/train/weights/best.pt)# 推理imagecv2.imread(test_image.jpg)resultsmodel(image)# 绘制结果forboxinresults[0].boxes:x1,y1,x2,y2map(int,box.xyxy[0])cv2.rectangle(image,(x1,y1),(x2,y2),(0,255,0),2)cv2.imshow(Result,image)cv2.waitKey(0)可视化界面使用Streamlit创建简单的GUIimportstreamlitasstfromPILimportImageimporttorchfromultralyticsimportYOLO# 加载模型modelYOLO(runs/detect/train/weights/best.pt)# Streamlit应用st.title(指针式仪表检测)uploaded_filest.file_uploader(上传一张图片,type[jpg,png])ifuploaded_fileisnotNone:imageImage.open(uploaded_file)resultsmodel(image)# 显示结果st.image(results[0].plot(),caption检测结果,use_column_widthTrue)总结以上代码展示了如何使用YOLOv8和PyTorch完成目标检测和语义分割任务并通过Streamlit构建可视化界面。您可以根据实际需求调整模型、数据集和界面设计。仅供参考
