使用深度学习模型YOLOv8训练使用水面垃圾数据集 模型对水面及漂浮物垃圾进行检测 识别水面漂浮物13种垃圾

水面垃圾数据集

格式：yolo数据集格式 voc格式
类别：can bad fruit bottle glass bottle plastic box plastic bag milk box
leaf branch paper ball grass garbage

标注：labelimg
总共有3000+张标记图像，主要应用背景：水面垃圾检测

使用深度学习模型对水面垃圾进行检测，采用YOLOv8作为目标检测框架。由于你提到数据集已经提供了YOLO和VOC两种格式的标注文件，基于YOLO格式进行训练。以下是详细的步骤、关键文件及代码示例来指导你完成从准备数据集到训练模型的整个流程。

1. 数据集准备

假设你的数据集结构如下：

water_trash_detection/ ├── images/ │ ├── train/ │ │ ├── img1.jpg │ │ └── ... │ ├── val/ │ │ ├── img1.jpg │ │ └── ... │ └── test/ │ ├── img1.jpg │ └── ... └── labels/ ├── train/ │ ├── img1.txt │ └── ... ├── val/ │ ├── img1.txt │ └── ... └── test/ ├── img1.txt └── ... data_water_trash.yaml

data_water_trash.yaml文件内容示例：

train:./water_trash_detection/images/train/val:./water_trash_detection/images/val/nc:12# 类别数量：can, bad fruit, bottle, glass bottle, plastic box, plastic bag, milk box, leaf, branch, paper, ball, grass, garbagenames:['can','bad fruit','bottle','glass bottle','plastic box','plastic bag','milk box','leaf','branch','paper','ball','grass','garbage']

确保每个标签文件的内容符合YOLO格式，即每行包含一个对象的类别ID、中心点x, y坐标（归一化到0-1），宽度和高度（同样归一化）。

2. 安装依赖库

确保安装了必要的库：

pipinstallultralytics opencv-python-headless tensorboard

3. 模型训练

创建一个Python脚本来开始训练过程。这里我们以YOLOv8为例说明如何训练模型。

训练脚本

fromultralyticsimportYOLOdefmain_train():# 加载预训练的YOLOv8n模型或从头开始定义模型model=YOLO('yolov8n.yaml')# 或者直接加载预训练权重，如 'yolov8n.pt'results=model.train(data='./data_water_trash.yaml',epochs=100,# 根据需要调整imgsz=640,batch=16,project='./runs/detect',name='water_trash_detection',optimizer='SGD',device='0',# 使用GPU编号save=True,cache=True,verbose=True,)if__name__=='__main__':main_train()

4. 推理与结果可视化

训练完成后，我们可以利用训练好的模型对新图片进行预测，并将结果可视化。

推理脚本

importcv2fromPILimportImagefromultralyticsimportYOLO model=YOLO('./runs/detect/water_trash_detection/weights/best.pt')defdetect_objects(image_path):results=model.predict(source=image_path)img=cv2.imread(image_path)forresultinresults:boxes=result.boxes.numpy()forboxinboxes:r=box.xyxy x1,y1,x2,y2=int(r[0]),int(r[1]),int(r[2]),int(r[3])label_id=int(box.cls)label=result.names[label_id]confidence=box.confifconfidence>0.5:# 设置置信度阈值cv2.rectangle(img,(x1,y1),(x2,y2),(0,255,0),2)# 绘制矩形框cv2.putText(img,f'{label}{confidence:.2f}',(x1,y1-10),cv2.FONT_HERSHEY_SIMPLEX,0.9,(0,255,0),2)returnimg# 示例调用result_image=detect_objects('your_test_image.jpg')# 确保测试图像路径正确Image.fromarray(cv2.cvtColor(result_image,cv2.COLOR_BGR2RGB)).show()# 使用PIL显示图像

5. 可视化界面

为了监控训练过程，可以使用TensorBoard。在训练脚本中添加tensorboard=True参数，然后运行以下命令启动TensorBoard：

tensorboard--logdirruns/

然后在浏览器中访问http://localhost:6006查看训练进度和结果。

6. 运行步骤总结

1. 数据集准备：确认数据集已按要求组织好，并创建data_water_trash.yaml。
2. 安装依赖：通过提供的命令安装所需的Python库。
3. 训练模型：运行训练脚本开始训练过程。
4. 推理与可视化：使用训练好的模型进行推理，并可视化结果。
5. 监控训练：使用TensorBoard监控训练过程。

使用YOLOv8对水面垃圾检测数据集进行目标检测任务。请根据实际需求调整相关参数和代码。如果有任何问题或需要进一步的帮助，请随时提问！

此外，希望将其转换为YOLO格式，可以编写一个简单的脚本来实现这种转换。例如：

importxml.etree.ElementTreeasETimportosdefconvert_voc_to_yolo(xml_file,classes):tree=ET.parse(xml_file)root=tree.getroot()size=root.find('size')w=int(size.find('width').text)h=int(size.find('height').text)yolo_lines=[]forobjinroot.iter('object'):difficult=obj.find('difficult').text cls=obj.find('name').textifclsnotinclassesorint(difficult)==1:continuecls_id=classes.index(cls)xmlbox=obj.find('bndbox')b=(float(xmlbox.find('xmin').text),float(xmlbox.find('xmax').text),float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))bb=convert((w,h),b)yolo_lines.append(f"{cls_id}{' '.join(map(str,bb))}")txt_file=xml_file.replace('.xml','.txt')withopen(txt_file,'w')asf:f.write('\n'.join(yolo_lines))defconvert(size,box):dw=1./(size[0])dh=1./(size[1])x=(box[0]+box[1])/2.0-1y=(box[2]+box[3])/2.0-1w=box[1]-box[0]h=box[3]-box[2]x=x*dw w=w*dw y=y*dh h=h*dhreturn(x,y,w,h)# 使用方法classes=['can','bad fruit','bottle','glass bottle','plastic box','plastic bag','milk box','leaf','branch','paper','ball','grass','garbage']forxml_fileinos.listdir('path/to/voc/annotations'):ifxml_file.endswith('.xml'):convert_voc_to_yolo(os.path.join('path/to/voc/annotations',xml_file),classes)

这个脚本可以帮助你快速地将VOC格式的标注文件转换为YOLO格式，以