DataWhale AI夏令营 2024大运河杯-数据开发应用创新赛
话不多说直接开始
baseline分析
这里我们忽略数据、模型下载的单元格
导入数据处理的一些包
import os, sys
import cv2, glob, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
读取下载的数据，并查看一下json的格式。
train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
train_anno[0], len(train_anno)
用pandas读取数据查看数据格式
pd.read_json('训练集(有标注第一批)/标注/45.json')
读取视频，使用VideoCapture对数据进行切帧处理。
video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
while True:
# 读取下一帧
ret, frame = cap.read()
if not ret:
break
break
根据json的信息，展示一张画框的图片
bbox = [746, 494, 988, 786]
pt1 = (bbox[0], bbox[1])
pt2 = (bbox[2], bbox[3])
color = (0, 255, 0)
thickness = 2 # 线条粗细
cv2.rectangle(frame, pt1, pt2, color, thickness)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
plt.imshow(frame)
截止到这里，上面其实都可以忽略，上面就是让大家看一下原始数据这个格式是什么样，大概该怎么处理这个数据。接下来开始构建YOLO所需的数据集。
构建YOLO数据集
yolo数据集的格式为一个data文件夹下包含三个内容,train; val; yolo.yaml,其中train和val不在介绍,yolo.yaml主要包含数据涉及到的标签信息。
我这里是吧数据放在/root/data文件夹下了，因为切帧的图片数据很多需要的空间后的云的系统盘空间不够。大家可以参考。
if not os.path.exists('/root/data/yolo-dataset/'):
os.mkdir('/root/data/yolo-dataset/')
if not os.path.exists('/root/data/yolo-dataset/train'):
os.mkdir('/root/data/yolo-dataset/train')
if not os.path.exists('/root/data/yolo-dataset/val'):
os.mkdir('/root/data/yolo-dataset/val')
dir_path = os.path.abspath('./') + '/'
# 需要按照你的修改path
with open('/root/data/yolo-dataset/yolo.yaml', 'w', encoding='utf-8') as up:
up.write(f'''
path: /root/data/yolo-dataset/
train: train/
val: val/
names:
0: 非机动车违停
1: 机动车违停
2: 垃圾桶满溢
3: 违法经营
''')
对获取的文件路径进行排序，以确保标注文件和视频文件按照相同顺序匹配。
train_annos = glob.glob('训练集(有标注第一批)/标注/*.json')
train_videos = glob.glob('训练集(有标注第一批)/视频/*.mp4')
train_annos.sort(); train_videos.sort();
category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]
我这里按照8:2划分训练集和验证集，一共应该是52组数据，划分之后是42:10
下面在代码中给出了逐行的注释，大家自行食用即可。
for anno_path, video_path in zip(train_annos[:42], train_videos[:42]):
print(video_path)
# 使用Pandas读取JSON格式的标注文件，返回一个DataFrame对象
anno_df = pd.read_json(anno_path)
# 使用OpenCV打开视频文件，准备逐帧读取
cap = cv2.VideoCapture(video_path)
frame_idx = 0
# 读取视频帧
while True:
ret, frame = cap.read()
if not ret:
break
# 获取当前帧的高度和宽度
img_height, img_width = frame.shape[:2]
# 从标注文件中提取当前帧的标注信息
frame_anno = anno_df[anno_df['frame_id'] == frame_idx]
# 将当前帧保存为JPEG图像文件
cv2.imwrite('/root/data/yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)
# 检查当前帧有没有标注信息
if len(frame_anno) != 0:
# 创建并打开一个与当前帧图像同名的文本文件，准备写入YOLO格式的标签
with open('/root/data/yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.txt', 'w') as up:
for category, bbox in zip(frame_anno['category'].values, frame_anno['bbox'].values):
# 获取当前标注对象类别的索引
category_idx = category_labels.index(category)
# 获取框的坐标
x_min, y_min, x_max, y_max = bbox
# 计算标注框的中心点横纵坐标，并归一化到 [0, 1] 之间
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
if x_center > 1:
print(bbox)
# 将YOLO格式的标注信息写入标签文件
up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
# 处理下一帧
frame_idx += 1
构建验证集，这部分代码直接看上一个即可一样的基本都是
for anno_path, video_path in zip(train_annos[-10:], train_videos[-10:]):
print(video_path)
anno_df = pd.read_json(anno_path)
cap = cv2.VideoCapture(video_path)
frame_idx = 0
while True:
ret, frame = cap.read()
if not ret:
break
img_height, img_width = frame.shape[:2]
frame_anno = anno_df[anno_df['frame_id'] == frame_idx]
cv2.imwrite('/root/data/yolo-dataset/val/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)
if len(frame_anno) != 0:
with open('/root/data/yolo-dataset/val/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.txt', 'w') as up:
for category, bbox in zip(frame_anno['category'].values, frame_anno['bbox'].values):
category_idx = category_labels.index(category)
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
up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
frame_idx += 1
开始训练
baseline使用的是yolov8n进行训练，在这里epoch代表训练的轮数，imgsz代表输入模型图像大小，batch代表一次梯度更新使用多少张图片
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="/root/data/yolo-dataset/yolo.yaml", epochs=15, imgsz=1080, batch=16)
创建结果目录
category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]
if not os.path.exists('result/'):
os.mkdir('result')
对测试集视频文件的处理，通过预训练的YOLO模型对每个视频的每一帧进行检测，并将检测结果保存为JSON格式的文件。
from ultralytics import YOLO
# 使用训练好的模型进行预测
model = YOLO("runs/detect/train/weights/best.pt")
import glob
for path in glob.glob('测试集/*.mp4'):
# 保存结果生成的json文件
submit_json = []
# 对视频文件进行推理，conf=0.05设置了最低置信度阈值
results = model(path, conf=0.05, imgsz=1080, verbose=False)
for idx, result in enumerate(results):
boxes = result.boxes # Boxes object for bounding box outputs
masks = result.masks # Masks object for segmentation masks outputs
keypoints = result.keypoints # Keypoints object for pose outputs
probs = result.probs # Probs object for classification outputs
obb = result.obb # Oriented boxes object for OBB outputs
if len(boxes.cls) == 0:
continue
# 获取检测框的坐标、类别、置信度
xywh = boxes.xyxy.data.cpu().numpy().round()
cls = boxes.cls.data.cpu().numpy().round()
conf = boxes.conf.data.cpu().numpy()
# 写入submit
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)
}
)
# 保存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)
120
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
