MARS城市治理比赛baseline解读
数据解读
官方提供训练集和数据集,其中训练集包含训练所用的所有视频数据,每个视频样本以mp4格式储存,还包含json文件,文件中储存的是标注信息;测试集包含测试所用的视频数据。
在json文件中,每个样本的标注信息各变量含义如下:
- frame_id:违规行为出现的帧编号
- event_id:违规行为ID
- category:违规行为类别
- bbox:检测到的违规行为矩形框的坐标,[xmin,ymin,xmax,ymax]形式
这些信息官方都有作介绍。
下面对baseline中的代码作解读
导入相关的库
import os, sys
import cv2, glob, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
下面的代码是linux命令,在jupyter中也可以执行。其主要作用就是从对应的链接中下载数据集并将数据集解压。
# !apt install zip unzip -y
# !apt install unar -y
# !wget "https://comp-public-prod.obs.cn-east-3.myhuaweicloud.com/dataset/2024/%E8%AE%AD%E7%BB%83%E9%9B%86%28%E6%9C%89%E6%A0%87%E6%B3%A8%E7%AC%AC%E4%B8%80%E6%89%B9%29.zip?AccessKeyId=583AINLNMLDRFK7CC1YM&Expires=1739168844&Signature=9iONBSJORCS8UNr2m/VZnc7yYno%3D" -O 训练集\(有标注第一批\).zip
# !unar -q 训练集\(有标注第一批\).zip
# !wget "https://comp-public-prod.obs.cn-east-3.myhuaweicloud.com/dataset/2024/%E6%B5%8B%E8%AF%95%E9%9B%86.zip?AccessKeyId=583AINLNMLDRFK7CC1YM&Expires=1739168909&Signature=CRsB54VqOtrzIdUHC3ay0l2ZGNw%3D" -O 测试集.zip
# !unar -q 测试集.zip
查看数据。下面的代码主要是从标注json文件中取出一个,查看其格式信息。
train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
train_anno[0], len(train_anno)
用pandas也可以读取json文件,并转化为DataFrame格式,方便查看。
pd.read_json('训练集(有标注第一批)/标注/45.json')
从训练集中取出一个视频样本查看其基本信息。可以看到,视频中的每一帧都是三通道图像,而这个视频中一共有422帧。
video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
while True:
# 读取下一帧
ret, frame = cap.read()
if not ret:
break
break
frame.shape # (1080, 1920, 3)
int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # frame.shape
前面以及对json文件中储存的数据的格式有了个初步的了解,这种格式是不能被用于训练yolo的,所以我们还要将json文件转化为yolo可以读取的yaml文件。首先创建文件夹用于储存转换格式后的yaml文件;随后我们就开始修改文件。
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('./') + '/'
# 需要按照你的修改path
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: 违法经营
''')
train_annos = glob.glob('训练集(有标注第一批)/标注/*.json')这行代码会找到 训练集(有标注第一批)/标注 这个文件夹中的所有的json文件,并返回其路径。
train_annos = glob.glob('训练集(有标注第一批)/标注/*.json')
train_videos = glob.glob('训练集(有标注第一批)/视频/*.mp4')
train_annos.sort(); train_videos.sort();
category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]
下面的代码是取前5个视频,将其转化为yaml格式的训练数据用于训练模型
for anno_path, video_path in zip(train_annos[:5], train_videos[:5]):
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('./yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)
if len(frame_anno) != 0:
with open('./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
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)
up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
frame_idx += 1
获得验证集
for anno_path, video_path in zip(train_annos[-3:], train_videos[-3:]):
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('./yolo-dataset/val/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)
if len(frame_anno) != 0:
with open('./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
下载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
用os设置默认训练器,0表示选择电脑上的第一块显卡。接着选择忽略警告。然后导入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)
后面的代码就是在测试集上测试训练的模型,并输出结果到result.zip文件中。
总的来说没什么复杂的地方,但是代码写的很漂亮。

浙公网安备 33010602011771号