道路垃圾识别数据集 含pt模型&界面 18类 共7542张图片,xml和txt标签都有;
YOLO道路垃圾识别系统
系统概述
名称:YOLO道路垃圾识别系统
图片数量 :7542张
标注格式 :XML和TXT(支持YOLO格式)
类别 :Plastic Bottle, Face Mask, PaperBag, Plastic Cup, Paper Cup, Cardboard, Peel, Cans, Plastic Wrapper, Paperboard, Styrofoam, Tetra Pack, Colored Glass Bottles, Plastic Bag, Rags, Pile of Leaves, Glass Bottle
组件 :数据集、预训练模型(.pt文件)、可视化界面
用途:用于路面垃圾识别,垃圾检测,目标检测,适用于Python环境下的机器学习、深度学习项目,特别适合使用YOLO进行目标检测。
数据集特点
- 规模:包含7542张高分辨率图像,每张图像都带有详细的标注信息。
- 多样性:图像涵盖了不同的光照条件、天气状况、背景环境以及垃圾种类,以确保模型能够适应多样的实际场景。
- 标注质量:每张图像都有精确的手动标注,确保了高质量的训练数据。
- 标注格式 :
- VOC格式 (XML):符合Pascal VOC标准的XML文件,包含了图像的基本信息、对象类别以及边界框坐标。
- YOLO格式 (TXT) :每个目标用一个文本行表示,格式为
class_id x_center y_center width height
,所有坐标值都是归一化的。
标注信息
-
VOC格式:
<annotation> <folder>images</folder> <filename>image_0001.jpg</filename> <size> <width>800</width> <height>600</height> <depth>3</depth> </size> <object> <name>Plastic Bottle</name> <bndbox> <xmin>200</xmin> <ymin>100</ymin> <xmax>300</xmax> <ymax>200</ymax> </bndbox> </object> <object> <name>Face Mask</name> <bndbox> <xmin>400</xmin> <ymin>200</ymin> <xmax>500</xmax> <ymax>300</ymax> </bndbox> </object> </annotation>
-
YOLO格式:
0 0.25 0.1667 0.125 0.1667 1 0.5 0.3333 0.125 0.1667
解释:
0
表示塑料瓶,1
表示口罩。x_center
和y_center
是边界框中心点的归一化坐标,width
和height
是边界框的宽度和高度的归一化值。
应用领域
- 城市清洁:自动检测道路上的垃圾,帮助环卫工人高效清理。
- 环境保护:监测公共区域的垃圾分布情况,提高环保意识。
- 智能监控:结合视频监控系统,实时检测并报告垃圾情况。
- 自动驾驶:辅助自动驾驶车辆识别路面上的障碍物,提高行驶安全性。
获取方式
通常情况下,研究人员可以通过官方提供的链接或相关机构网站下载该数据集。请注意,使用时应遵循相应的许可协议和引用要求。
关键代码示例
1. 下载数据集
假设我们已经有了数据集的下载链接,可以使用 Python 的 requests
库来下载数据集:
import requests
import os
# 定义下载链接和保存路径
url = 'http://example.com/path/to/road_litter_dataset.zip' # 替换为实际的下载链接
save_path = './road_litter_dataset.zip'
# 检查是否已经下载过
if not os.path.exists(save_path):
print("Downloading dataset...")
response = requests.get(url, stream=True)
with open(save_path, 'wb') as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
print("Download complete.")
else:
print("Dataset already exists.")
# 解压数据集
import zipfile
with zipfile.ZipFile(save_path, 'r') as zip_ref:
zip_ref.extractall('./road_litter_dataset')
2. 解析 VOC 格式的标注文件
以下是一个解析 VOC 格式标注文件的函数:
import xml.etree.ElementTree as ET
def parse_voc_annotation(anno_file):
tree = ET.parse(anno_file)
root = tree.getroot()
annotations = []
for obj in root.findall('object'):
name = obj.find('name').text
bndbox = obj.find('bndbox')
xmin = int(bndbox.find('xmin').text)
ymin = int(bndbox.find('ymin').text)
xmax = int(bndbox.find('xmax').text)
ymax = int(bndbox.find('ymax').text)
annotations.append({
'class_name': name,
'bbox': [xmin, ymin, xmax, ymax]
})
return annotations
3. 加载图像并显示标注框
我们可以使用 OpenCV 来加载图像,并使用 Matplotlib 来显示图像及其标注框:
import cv2
import matplotlib.pyplot as plt
def load_image(image_path):
return cv2.imread(image_path)
def display_image_with_annotations(image, annotations):
fig, ax = plt.subplots()
ax.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
for anno in annotations:
bbox = anno['bbox']
rect = plt.Rectangle((bbox[0], bbox[1]), bbox[2] - bbox[0], bbox[3] - bbox[1],
fill=False, edgecolor='red', linewidth=2)
ax.add_patch(rect)
plt.show()
# 示例用法
image_path = './road_litter_dataset/images/image_0001.jpg'
anno_path = './road_litter_dataset/annotations/image_0001.xml'
image = load_image(image_path)
annotations = parse_voc_annotation(anno_path)
display_image_with_annotations(image, annotations)
4. 使用数据集进行训练
如果您打算使用这个数据集进行深度学习模型的训练,可以使用 PyTorch 或 TensorFlow 等框架。以下是一个简单的 PyTorch DataLoader 示例:
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
class RoadLitterDataset(Dataset):
def __init__(self, image_dir, anno_dir, transform=None):
self.image_dir = image_dir
self.anno_dir = anno_dir
self.transform = transform
self.images = os.listdir(image_dir)
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
img_name = self.images[idx]
image = Image.open(os.path.join(self.image_dir, img_name)).convert("RGB")
anno_name = img_name.replace('.jpg', '.xml')
anno_path = os.path.join(self.anno_dir, anno_name)
annotations = parse_voc_annotation(anno_path)
if self.transform:
image = self.transform(image)
return image, annotations
# 创建 DataLoader
dataset = RoadLitterDataset(image_dir='./road_litter_dataset/images',
anno_dir='./road_litter_dataset/annotations')
dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2)
# 遍历数据
for images, annotations in dataloader:
# 在这里进行模型训练
pass
YOLO模型训练
为了使用YOLO进行道路垃圾识别,您可以使用预训练的YOLO模型,并对其进行微调。以下是一个简单的YOLOv5训练示例:
-
安装YOLOv5:深色版本
train: ./road_litter_dataset/train/images
val: ./road_litter_dataset/val/images
nc: 17 # 类别数
names: ['Plastic Bottle', 'Face Mask', 'PaperBag', 'Plastic Cup', 'Paper Cup', 'Cardboard', 'Peel', 'Cans', 'Plastic Wrapper', 'Paperboard', 'Styrofoam', 'Tetra Pack', 'Colored Glass Bottles', 'Plastic Bag', 'Rags', 'Pile of Leaves', 'Glass Bottle'] # 类别名称
-
训练模型: 使用YOLOv5进行训练:
python train.py --img 640 --batch 16 --epochs 50 --data data.yaml --weights yolov5s.pt
-
评估模型: 训练完成后,可以使用验证集进行评估:
bash
深色版本
python val.py --data data.yaml --weights runs/train/exp/weights/best.pt
-
推理测试: 使用训练好的模型进行推理测试:
python detect.py --source ./road_litter_dataset/test/images --weights runs/train/exp/weights/best.pt --conf 0.4
可视化界面
为了提供一个用户友好的界面,您可以使用 Streamlit
或 Flask
等工具来构建一个简单的Web应用程序。以下是一个使用 Streamlit
的示例:
-
安装 Streamlit:
pip install streamlit
-
创建 Streamlit 应用程序:
import streamlit as st from PIL import Image import torch from yolov5.models.experimental import attempt_load from yolov5.utils.general import non_max_suppression from yolov5.utils.plots import plot_one_box import cv2 # 设置标题 st.title("YOLO道路垃圾识别系统") # 上传图片 uploaded_file = st.file_uploader("选择一张图片", type=["jpg", "jpeg", "png"]) # 加载预训练模型 model = attempt_load('runs/train/exp/weights/best.pt', map_location=torch.device('cpu')) # 假设模型在CPU上运行 model.eval() if uploaded_file is not None: # 读取并显示图片 image = Image.open(uploaded_file) st.image(image, caption="上传的图片", use_column_width=True) # 图像预处理 img = image.convert('RGB') img = np.array(img) img = cv2.resize(img, (640, 640)) # 调整大小以匹配模型输入 img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, HWC to CHW img = np.ascontiguousarray(img) img = torch.from_numpy(img).to(torch.device('cpu')) img = img.float() / 255.0 # 归一化 if img.ndimension() == 3: img = img.unsqueeze(0) # 进行推理 with torch.no_grad(): pred = model(img)[0] pred = non_max_suppression(pred, 0.4, 0.5) # 绘制结果 for i, det in enumerate(pred): # detections per image if len(det): for *xyxy, conf, cls in reversed(det): label = f'{model.names[int(cls)]} {conf:.2f}' plot_one_box(xyxy, img, label=label, color=(0, 255, 0), line_thickness=3) # 显示结果 result_image = Image.fromarray(cv2.cvtColor(img.squeeze().permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) st.image(result_image, caption="检测结果", use_column_width=True)
-
运行 Streamlit 应用程序:
streamlit run app.py
通过上述步骤,您将拥有一个完整的YOLO道路垃圾识别系统,包括数据集、预训练模型和用户友好的可视化界面。希望这些代码能帮助您更好地利用该系统!