道路积水数据集 路面积水识别数据集 图片数量4524，xml和txt标签都有；公路积水数据集 ✓类别：puddle；
道路积水数据集 路面积水识别数据集 图片数量4524，xml和txt标签都有；公路积水数据集
✓类别：puddle；
✓图片数量4524，xml和txt标签都有；
✓类别：puddle；
好的，下面是一个完整的项目代码和详细配置说明，用于基于YOLOv8进行路面积水识别。该项目包括数据集准备、模型训练、评估和推理，并附带一个简单的Streamlit界面。
完整项目代码
1. 数据集准备与转换
首先，我们需要将数据集从XML格式转换为YOLO格式（如果数据集已经是YOLO格式，则跳过此步骤）。
importxml.etree.ElementTreeasETimportosfrompathlibimportPathdefconvert_xml_to_yolo(xml_dir,output_dir):forxml_fileinPath(xml_dir).glob('*.xml'):tree=ET.parse(xml_file)root=tree.getroot()img_name=root.find('filename').text img_width=int(root.find('size/width').text)img_height=int(root.find('size/height').text)txt_file_path=os.path.join(output_dir,Path(img_name).stem+'.txt')withopen(txt_file_path,'w')astxt_file:forobjinroot.findall('object'):cls=obj.find('name').textifcls!='puddle':continuebbox=obj.find('bndbox')xmin=float(bbox.find('xmin').text)ymin=float(bbox.find('ymin').text)xmax=float(bbox.find('xmax').text)ymax=float(bbox.find('ymax').text)x_center=(xmin+xmax)/2/img_width y_center=(ymin+ymax)/2/img_height width=(xmax-xmin)/img_width height=(ymax-ymin)/img_height txt_file.write(f"0{x_center}{y_center}{width}{height}\n")# 使用示例convert_xml_to_yolo('datasets/puddle_dataset/xml','datasets/puddle_dataset/yolo/labels')2. 创建数据集配置文件 (data.yaml)
创建一个data.yaml文件来配置数据集路径和类别信息。
3. 分割数据集
为了训练和验证，我们需要将数据集分割成训练集和验证集。
importosimportrandomfrompathlibimportPathimportshutildefsplit_dataset(data_dir,train_ratio=0.8):images=list(Path(data_dir).glob('*.jpg'))random.shuffle(images)num_train=int(len(images)*train_ratio)train_images=images[:num_train]val_images=images[num_train:]train_dir=Path(data_dir).parent/'train'val_dir=Path(data_dir).parent/'val'train_img_dir=train_dir/'images'train_label_dir=train_dir/'labels'val_img_dir=val_dir/'images'val_label_dir=val_dir/'labels'train_img_dir.mkdir(parents=True,exist_ok=True)train_label_dir.mkdir(parents=True,exist_ok=True)val_img_dir.mkdir(parents=True,exist_ok=True)val_label_dir.mkdir(parents=True,exist_ok=True)forimgintrain_images:label_path=img.with_suffix('.txt')shutil.copy(img,train_img_dir/img.name)shutil.copy(label_path,train_label_dir/label_path.name)forimginval_images:label_path=img.with_suffix('.txt')shutil.copy(img,val_img_dir/img.name)shutil.copy(label_path,val_label_dir/label_path.name)# 使用示例split_dataset('./datasets/puddle_dataset/images')4. 训练脚本
接下来是使用YOLOv8进行训练的脚本。
importtorchfromultralyticsimportYOLO# 设置随机种子以保证可重复性torch.manual_seed(42)# 定义数据集路径dataset_config='data.yaml'# 加载预训练的YOLOv8n模型model=YOLO('yolov8n.pt')# 训练模型results=model.train(data=dataset_config,epochs=100,imgsz=512,batch=16,name='puddle_detection',project='runs/train')# 评估模型metrics=model.val()# 保存最佳模型权重best_model_weights='runs/train/puddle_detection/weights/best.pt'print(f"Best model weights saved to{best_model_weights}")5. 推理脚本
以下是一个简单的推理脚本，用于测试模型的性能。
fromultralyticsimportYOLOimportcv2importnumpyasnpfromPILimportImage# 加载模型model=YOLO('runs/train/puddle_detection/weights/best.pt')# 图片检测函数defdetect_image(model,image_path,conf_threshold=0.5):results=model.predict(image_path,conf=conf_threshold)[0]annotated_frame=annotate_image(image_path,results,model)returnannotated_frame# 标注图像函数defannotate_image(image_path,results,model):frame=cv2.imread(image_path)forresultinresults.boxes.cpu().numpy():r=result.xyxy[0].astype(int)cls=int(result.cls[0])conf=result.conf[0]label=f"{model.names[cls]}{conf:.2f}"color=(0,255,0)cv2.rectangle(frame,(r[0],r[1]),(r[2],r[3]),color,2)cv2.putText(frame,label,(r[0],r[1]-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,color,2)returnframe# 测试一张图片test_image_path='./datasets/puddle_dataset/images/test_image.jpg'annotated_image=detect_image(model,test_image_path)cv2.imwrite('annotated_test_image.jpg',annotated_image)6. Streamlit 主界面
最后，我们创建一个简单的Streamlit界面来进行实时推理和显示结果。
importstreamlitasstfromultralyticsimportYOLOimportcv2importnumpyasnpfromPILimportImageimporttempfile# 加载模型@st.cache_resourcedefload_model(weights_path):model=YOLO(weights_path)returnmodel# 图片检测函数defdetect_image(model,image,conf_threshold):results=model.predict(image,conf=conf_threshold)[0]annotated_frame=annotate_image(image,results,model)returnannotated_frame# 视频检测函数defdetect_video(model,video_path,conf_threshold):cap=cv2.VideoCapture(video_path)whilecap.isOpened():ret,frame=cap.read()ifnotret:breakresults=model.predict(frame,conf=conf_threshold)[0]annotated_frame=annotate_image(frame,results,model)yieldannotated_frame cap.release()# 摄像头检测函数defdetect_camera(model,conf_threshold):cap=cv2.VideoCapture(0)whileTrue:ret,frame=cap.read()ifnotret:breakresults=model.predict(frame,conf=conf_threshold)[0]annotated_frame=annotate_image(frame,results,model)yieldannotated_frame cap.release()# 标注图像函数defannotate_image(image,results,model):forresultinresults.boxes.cpu().numpy():r=result.xyxy[0].astype(int)cls=int(result.cls[0])conf=result.conf[0]label=f"{model.names[cls]}{conf:.2f}"color=(0,255,0)cv2.rectangle(image,(r[0],r[1]),(r[2],r[3]),color,2)cv2.putText(image,label,(r[0],r[1]-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,color,2)returnimage# Streamlit 主界面defmain():st.title("Puddle Detection System")# 动态加载模型weights_options=['best_puddle_yolov8.pt']# 添加更多权重文件路径selected_weights=st.sidebar.selectbox("Select Model Weights",weights_options)model=load_model(selected_weights)# 动态调整置信度阈值conf_threshold=st.sidebar.slider("Confidence Threshold",min_value=0.0,max_value=1.0,value=0.5,step=0.01)# 输入方式选择input_type=st.sidebar.radio("Input Type",["Image","Video","Camera"])ifinput_type=="Image":uploaded_file=st.file_uploader("Upload an image...",type=["jpg","jpeg","png"])ifuploaded_fileisnotNone:image=Image.open(uploaded_file)image_np=np.array(image)annotated_image=detect_image(model,image_np,conf_threshold)st.image(annotated_image,channels="BGR",caption="Detected Image")# 统计检测到的物体数量results=model.predict(image_np,conf=conf_threshold)[0]class_counts={}forresultinresults.boxes.cpu().numpy():cls=int(result.cls[0])class_name=model.names[cls]ifclass_nameinclass_counts:class_counts[class_name]+=1else:class_counts[class_name]=1st.subheader("Detection Summary:")forclass_name,countinclass_counts.items():st.write(f"{class_name}:{count}")elifinput_type=="Video":uploaded_file=st.file_uploader("Upload a video...",type=["mp4","avi"])ifuploaded_fileisnotNone:tfile=tempfile.NamedTemporaryFile(delete=False)tfile.write(uploaded_file.read())tfpath=tfile.name cap=cv2.VideoCapture(tfpath)frame_placeholder=st.empty()forannotated_frameindetect_video(model,tfpath,conf_threshold):frame_placeholder.image(annotated_frame,channels="BGR",use_column_width=True)cap.release()os.remove(tfpath)elifinput_type=="Camera":frame_placeholder=st.empty()forannotated_frameindetect_camera(model,conf_threshold):frame_placeholder.image(annotated_frame,channels="BGR",use_column_width=True)if__name__=="__main__":main()文件结构
puddle_detection/ ├── main.py ├── datasets/ │ └── puddle_dataset/ │ ├── xml/ │ │ ├── image1.xml │ │ ├── image2.xml │ │ └── ... │ ├── yolo/ │ │ ├── labels/ │ │ │ ├── image1.txt │ │ │ ├── image2.txt │ │ │ └── ... │ │ ├── train/ │ │ │ ├── images/ │ │ │ │ ├── image1.jpg │ │ │ │ ├── image2.jpg │ │ │ │ └── ... │ │ │ └── labels/ │ │ │ ├── image1.txt │ │ │ ├── image2.txt │ │ │ └── ... │ │ └── val/ │ │ ├── images/ │ │ │ ├── image1.jpg │ │ │ ├── image2.jpg │ │ │ └── ... │ │ └── labels/ │ │ ├── image1.txt │ │ ├── image2.txt │ │ └── ... │ └── images/ │ ├── image1.jpg │ ├── image2.jpg │ └── ... ├── best_puddle_yolov8.pt └── requirements.txt安装依赖项
首先，确保你已经安装了所有必要的依赖项。你可以通过以下命令安装：
pipinstall-rrequirements.txtrequirements.txt内容如下：
streamlit==1.25.0 opencv-python torch==2.0 ultralytics运行步骤总结
克隆项目仓库（如果有的话）：
gitclone https://github.com/yourusername/puddle_detection.gitcdpuddle_detection安装依赖项：
pipinstall-rrequirements.txt数据集准备：
- 如果数据集尚未转换，请运行数据集转换脚本。
python convert_xml_to_yolo.py - 如果数据集尚未分割，请运行数据集分割脚本。
python split_dataset.py - 确保数据集目录结构正确。
- 如果数据集尚未转换，请运行数据集转换脚本。
训练模型：
python train.py评估模型：
在训练脚本中，模型会在训练结束后自动进行评估。推理测试：
python inference.py运行Streamlit应用：
streamlit run main.py
总结
以上是完整的基于YOLOv8的路面积水识别系统的项目介绍和代码实现。该项目支持图片、视频识别以及本地摄像头识别，并且可以通过UI界面动态调节模型置信度和选择不同的模型权重。希望这些详细的信息和代码能够帮助你顺利实施和优化你的路面积水识别系统。