Datawhale AI夏令营 第五期-CV方向Task2笔记(代码解读)

这次参加的比赛是2024“大运河杯”数据开发应用创新大赛——城市治理
赛事链接:https://www.marsbigdata.com/competition/details?id=3839107548872

Baseline代码实践解读

上回说到,我们用baseline代码跑通了整个流程,这次的任务是解读建模过程。
这段代码的主要目的是使用YOLOv8(You Only Look Once,第八版)深度学习模型来训练我们的赛题任务,然后使用训练好的模型来检测视频中的违规行为,并将结果保存为JSON格式的文件。

1. 安装和导入必要的库

!/opt/miniconda/bin/pip install opencv-python pandas matplotlib ultralytics
import os, sys
import cv2, glob, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
  • 安装Python库:opencv-pythonpandasmatplotlibultralytics
  • 导入必要的Python库,如oscv2(OpenCV)、globjsonnumpypandasmatplotlib

2. 下载并解压数据集

!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
  • 通过wget下载训练集和测试集的压缩包。
  • 使用unar解压下载的压缩文件。
  • 这里因为下载链接好像出bug了,所以我直接上传了之前下载好的数据集,如果有下载好的数据集,可以只进行解压操作。

3. 读取并可视化标注信息

train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()
frame.shape
  • 读取并加载一个标注文件(JSON格式)。
  • 读取对应的视频文件并获取其中一帧图像。
  • frame.shape返回图像的形状。

4. 在图像上绘制标注框

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)
  • 根据提供的bbox值绘制一个绿色的矩形框,以表示标注框,如下图。
  • image.png

5. 准备YOLO训练数据集

if not os.path.exists('yolo-dataset/'):
    os.mkdir('yolo-dataset/')
# 创建训练和验证数据集的目录结构
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,定义数据集路径和类别名称。
  • path: 指定了数据集的根目录路径,即所有数据子文件夹的上级目录。这里的路径是相对于当前配置文件的路径或者相对于执行训练脚本的工作目录。
  • train: 定义了训练集图像的相对路径。在训练模型时,程序会在指定的路径下查找图像文件。
  • val: 定义了验证集图像的相对路径。验证集用于在训练过程中评估模型性能,避免过拟合。
  • names: 是一个列表,包含了每个类别的名称。

6. 处理训练和验证数据

for anno_path, video_path in zip(train_annos[:5], train_videos[:5]):
    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
                    up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')
        frame_idx += 1
  • 读取每个标注文件和对应的视频文件。
  • 根据标注文件内容,将视频的每一帧保存为图像,并生成YOLO格式的标注文件。

7. 下载并训练YOLO模型

!wget http://mirror.coggle.club/yolo/yolov8n-v8.2.0.pt -O yolov8n.pt
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=2, imgsz=1080, batch=16)
  • 下载YOLOv8的预训练权重文件。
  • 使用YOLOv8进行模型训练,指定数据集路径和训练参数。
  • 这里可以考虑把yolov8n替换成不同的模型预训练权重,比如yolov8s、yolov8m、yolov8l、yolov8x。
  • 调整epochs,可以试试10,20,50等,batch大小可以调整16、32、64等,不过batch的调整可能会导致内存不够用,且看且调吧。

8. 使用训练好的模型进行预测

for path in glob.glob('测试集/*.mp4'):
    submit_json = []
    results = model(path, conf=0.05, imgsz=1080,  verbose=False)
    for idx, result in enumerate(results):
        boxes = result.boxes
        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()
        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)
                }
            )
    with open('./result/' + path.split('/')[-1][:-4] + '.json', 'w', encoding='utf-8') as up:
        json.dump(submit_json, up, indent=4, ensure_ascii=False)
  • 使用训练好的模型对测试集中的视频进行预测。
  • 将预测结果(类别、置信度、边界框)保存为JSON文件。

9. 压缩结果文件

!zip -r result.zip result/
  • 将预测结果压缩成一个ZIP文件,下载后便可到赛事官网提交结果。

这是我跑了10个epoch后的结果
image.png
其中:

train/box_loss衡量模型预测的边界框与真实边界框之间差异的损失值。
train/cls_loss衡量模型预测的类别与真实类别之间差异的损失值。
train/dfl_loss衡量模型对难分类样本的关注程度,减少易分类样本的影响。
metrics/precision(B)在训练过程中,预测为正类别中实际为正类别的比例。
metrics/recall(B)在训练过程中,所有实际正类别中被模型正确预测为正类别的比例。
metrics/mAP50(B)在50%的IoU阈值下计算,衡量模型的整体性能。
metrics/mAP50-95(B)在0.5到0.95的IoU阈值范围内计算,提供更全面的模型性能评估。
val/box_loss模型在未见过的验证集上的边界框损失,用于监控模型的泛化能力。
val/cls_loss模型在验证集上的分类损失,用于监控模型的泛化能力。
val/dfl_loss模型在验证集上的难易样本平衡损失。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值