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文件中。
总的来说没什么复杂的地方,但是代码写的很漂亮。

posted @ 2024-08-22 21:12  蒋古诚  阅读(51)  评论(0)    收藏  举报