Task1:跑通YOLO方案baseline!
- 一、准备工作与环境搭建
Datawhale官方有提供详细的速通文档:Task1:跑通YOLO方案baseline!
按照上述文档可以速通baseline。
- 二、项目背景介绍:
利用图像处理和计算机视觉技术开发一套智能识别系统,自动检测和分类摄像头捕获的视频中,城市管理中的违规行为,包括:垃圾桶满溢、机动车违停、非机动车违停、违法经营等。视频数据为mp4格式,标注文件为json格式,每个视频对应一个json文件,json文件包含内容如下:
-
frame_id:违规行为出现的帧编号
-
event_id:违规行为ID
-
category:违规行为类别
-
bbox:检测到的违规行为矩形框的坐标,[xmin,ymin,xmax,ymax]形式
-
confidence:置信度,使用F1score、MOTA指标来评估模型预测结果
示例如下:
https://file.public.marsbigdata.com/2024/08/14/rksAIZHbVMj-3PAg.png -
三、部分代码解析
以下是你提供的代码的逐步解析和目的解释:
- 安装所需的Python包
!/opt/miniconda/bin/pip install opencv-python pandas matplotlib ultralytics
- 目的:使用
pip
安装必要的Python库,包括opencv-python
(用于图像处理)、pandas
(用于数据处理)、matplotlib
(用于数据可视化)和ultralytics
(用于YOLO模型的训练和推理)。
- 安装操作系统包
!apt install zip unzip -y
!apt install unar -y
- 目的:安装解压缩工具
zip
、unzip
和unar
,用于解压下载的数据集文件。
3. 下载并解压数据集
!wget "https://.../训练集(有标注第一批).zip" -O 训练集\(有标注第一批\).zip
!unar -q 训练集\(有标注第一批\).zip
!wget "https://.../测试集.zip" -O 测试集.zip
!unar -q 测试集.zip
- 目的:通过
wget
下载训练集和测试集,并使用unar
解压这些文件,以便后续处理和分析。
4. 加载并预览训练数据
train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
例如:
({‘frame_id’: 0,
‘event_id’: 1,
‘category’: ‘非机动车违停’,
‘bbox’: [746, 494, 988, 786]},
1688)
train_anno[0], len(train_anno)
pd.read_json('训练集(有标注第一批)/标注/45.json')
例如:
- 目的:使用
json
库加载训练集的标注数据,并使用pandas
预览JSON文件内容,查看标注信息和数量。
5. 读取视频帧并绘制矩形框
video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
while True:# 使用 cap.read() 从视频中读取下一帧。ret 是一个布尔值,表示帧是否成功读取;frame 则是读取的图像。读取第一帧后,循环立即中断,因此只读取了一帧。
ret, frame = cap.read()
if not ret:
break
break
frame.shape#获取并返回 frame 的形状,通常是(高度, 宽度, 通道数)的形式,比如(1080, 1920, 3) 表示1080p分辨率的RGB 图像
int(cap.get(cv2.CAP_PROP_FRAME_COUNT))#获取视频中的总帧数。
bbox = [746, 494, 988, 786]#定义一个包含四个数值的列表 bbox,表示目标的边界框。bbox 代表框的左上角 (x_min, y_min) 和右下角 (x_max, y_max) 的坐标。
pt1 = (bbox[0], bbox[1])
pt2 = (bbox[2], bbox[3])#将边界框的坐标分解为左上角 (pt1) 和右下角 (pt2) 的两个点。
color = (0, 255, 0)
thickness = 2#color 定义了矩形框的颜色,这里使用绿色 (0, 255, 0),thickness 定义了矩形框的线条粗细为 2 像素
cv2.rectangle(frame, pt1, pt2, color, thickness)#在帧 frame 上绘制矩形框,框的顶点由 pt1 和 pt2 决定,使用指定的颜色和线条粗细。
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)#将帧 frame 从 BGR 色彩空间转换为 RGB 色彩空间,这是因为 OpenCV 默认读取的图像是 BGR 格式,但 matplotlib 的 imshow 方法要求图像为 RGB 格式。
plt.imshow(frame)
- 目的:使用
OpenCV
读取视频帧并绘制一个绿色的矩形框,表示目标物体的边界框,然后使用matplotlib
显示带有框的图像。结果如下图所示:
6. 创建YOLO数据集文件夹和配置文件
if not os.path.exists('yolo-dataset/'):
os.mkdir('yolo-dataset/')
if not os.path.exists('yolo-dataset/train'):
os.mkdir('yolo-dataset/train')
if not os.path.exists('yolo-dataset/val'):
os.mkdir('yolo-dataset/val')
dir_path = os.path.abspath('./') + '/'
with open('yolo-dataset/yolo.yaml', 'w', encoding='utf-8') as up:
up.write(f'''
path: {dir_path}/yolo-dataset/
train: train/
val: val/
names:
0: 非机动车违停
1: 机动车违停
2: 垃圾桶满溢
3: 违法经营
''')
- 目的:创建用于YOLO模型训练的数据集文件夹,并生成YOLO配置文件
yolo.yaml
,定义数据集路径和类别标签。
7. 处理并保存训练集数据
train_annos = glob.glob('训练集(有标注第一批)/标注/*.json')
train_videos = glob.glob('训练集(有标注第一批)/视频/*.mp4')
train_annos.sort(); train_videos.sort();#对标注文件和视频文件的路径列表进行排序,确保标注和视频的对应关系一致。
category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]#定义一个 category_labels 列表,列出了四个类别的标签名称
for anno_path, video_path in zip(train_annos[:5], train_videos[:5]):
print(video_path)#使用 zip() 函数将 train_annos 和 train_videos 列表的前五个元素配对(即每个标注文件对应一个视频文件),即处理5个视频文本对。每个循环开始时,打印当前处理的视频路径。
anno_df = pd.read_json(anno_path)#使用 pandas 的 read_json() 函数读取当前标注文件,并将其转换为一个 DataFrame 对象 anno_df,以便于后续数据处理
cap = cv2.VideoCapture(video_path)#使用 cv2.VideoCapture 打开当前视频文件,创建一个 cap 对象,用于读取视频帧。
frame_idx = 0 #初始化 frame_idx,将其设置为 0,用于跟踪当前处理的视频帧的索引。
while True:
ret, frame = cap.read()
if not ret:
break #入一个 while 循环,使用 cap.read() 逐帧读取视频
img_height, img_width = frame.shape[:2]#获取当前帧的高度和宽度
frame_anno = anno_df[anno_df['frame_id'] == frame_idx]#从 anno_df 中筛选出当前帧对应的标注信息,即 frame_id 等于 frame_idx 的所有标注。frame_anno 是一个包含当前帧标注数据的 DataFrame
cv2.imwrite('./yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)#将当前帧保存为图像文件,文件名格式为 标注文件名_帧号.jpg,并保存在 yolo-dataset/train/ 目录下。
if len(frame_anno) != 0:
with open('./yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.txt', 'w') as up:#如果当前帧有标注信息(即 frame_anno 不为空),则打开一个与图像同名的 .txt 文件,用于保存 YOLO 格式的标注数据。
for category, bbox in zip(frame_anno['category'].values, frame_anno['bbox'].values):
category_idx = category_labels.index(category)#遍历当前帧的标注信息,将类别名称 category 和边界框 bbox 配对,找到该类别名称在 category_labels 中的索引,并赋值给 category_idx,这是 YOLO 格式所需的类别编号,例如0、1、2、3
x_min, y_min, x_max, y_max = bbox
x_center = (x_min + x_max) / 2 / img_width
y_center = (y_min + y_max) / 2 / img_height
width = (x_max - x_min) / img_width
height = (y_max - y_min) / img_height
#读取边界框坐标 bbox,计算边界框的中心坐标 (x_center, y_center),并将其归一化(相对于图像宽度和高度)。同时计算边界框的宽度和高度,并同样进行归一化。
if x_center > 1:
print(bbox)#检查计算出的 x_center 是否大于 1(这不应该发生,因为坐标已经归一化),如果是,打印出有问题的 bbox 坐标以便调试。
up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
#将 YOLO 格式的标注信息(类别编号、中心坐标、宽度、高度)写入 .txt 文件,格式为 类别编号 x_center y_center 宽度 高度。
frame_idx += 1 #将 frame_idx 增加 1,指向视频的下一帧,准备在下一次循环中处理。
- 目的:处理训练集中的前5个视频及其标注数据,提取每一帧图像并保存到YOLO数据集的
train
文件夹,同时生成对应的YOLO格式的标注文件。
9. 下载并配置YOLO模型
!wget http://mirror.coggle.club/yolo/yolov8n-v8.2.0.pt -O yolov8n.pt
!mkdir -p ~/.config/Ultralytics/
!wget http://mirror.coggle.club/yolo/Arial.ttf -O ~/.config/Ultralytics/Arial.ttf
- 目的:下载预训练的YOLOv8模型权重文件,并下载字体文件用于YOLO模型的输出显示。
10. 设置环境变量并训练YOLO模型
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=2, imgsz=1080, batch=16)
- 目的:设置GPU环境,忽略警告信息,然后加载YOLOv8模型并进行训练,使用之前创建的YOLO配置文件,设置训练参数如训练周期数(epochs)、图像大小(imgsz)和批量大小(batch)。
11. 使用训练好的模型进行推理并生成结果文件
category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]
#定义一个包含目标类别标签的列表,这些标签用于对检测到的目标进行分类。
if not os.path.exists('result/'):
os.mkdir('result')from ultralytics import YOLO
#检查是否存在名为 result 的目录,如果不存在则创建该目录,用于存放检测结果。
model = YOLO("runs/detect/train/weights/best.pt")
#加载经过训练的 YOLO 模型,模型权重文件位于 runs/detect/train/weights/best.pt
import glob
for path in glob.glob('测试集/*.mp4'):
#遍历 测试集 目录下所有 .mp4 文件,逐个进行处理。
submit_json = []
#初始化一个空列表 submit_json,用于存放每个视频的检测结果。
results = model(path, conf=0.05, imgsz=1080, verbose=False)
#使用 YOLO 模型对视频 path 进行目标检测,设置置信度阈值为 0.05,图像大小为 1080,且不输出详细日志。
for idx, result in enumerate(results):
boxes = result.boxes # 提取检测结果中的边界框信息。
masks = result.masks
keypoints = result.keypoints
probs = result.probs
obb = result.obb #分别提取检测结果中的掩码、关键点、分类概率和定向边界框(如果有的话)。
if len(boxes.cls) == 0:
continue!\rm result/.ipynb_checkpoints/ -rf
!\rm result.zip
!zip -r result.zip result/
xywh = boxes.xyxy.data.cpu().numpy().round()#将检测到的边界框坐标转换为 numpy 数组,并进行四舍五入。
cls = boxes.cls.data.cpu().numpy().round()#将检测到的类别索引转换为 numpy 数组,并进行四舍五入。
conf = boxes.conf.data.cpu().numpy()#将检测到的置信度转换为 numpy 数组。
for i, (ci, xy, confi) in enumerate(zip(cls, xywh, conf)):#遍历当前帧中所有检测到的目标,提取类别、边界框坐标和置信度。
submit_json.append(
{
'frame_id': idx,
'event_id': i+1,
'category': category_labels[int(ci)],
'bbox': list([int(x) for x in xy]),
"confidence": float(confi)
}
)#将每个检测目标的结果(包括帧编号、事件编号、类别、边界框和置信度)以字典形式添加到 submit_json 列表中。
with open('./result/' + path.split('/')[-1][:-4] + '.json', 'w', encoding='utf-8') as up:
json.dump(submit_json, up, indent=4, ensure_ascii=False)#将 submit_json 列表中的检测结果写入 JSON 文件,格式化输出并保持字符编码。
最后得到结果文件result.zip,代码运行完成!