Datawhale AI 从零上手CV竞赛Task1:YOLO方案笔记

1. 项目背景

本项目旨在通过使用YOLO(You Only Look Once)算法完成计算机视觉竞赛的Task1。YOLO是一种基于深度学习的目标检测算法,具有实时性强、准确率高的特点,适用于自动驾驶、安防监控等领域。本文档详细记录了该项目的实施过程,从数据集的准备到模型的训练和测试,并讨论了每个步骤中的关键技术细节。

2. 项目环境配置

在开始项目之前,首先需要配置开发环境。本项目使用了Python 3.8,并且依赖以下主要库:

  • torch:用于深度学习模型的构建和训练
  • numpy:用于数值计算
  • opencv-python:用于图像处理
  • matplotlib:用于可视化
!/opt/miniconda/bin/pip install opencv-python pandas matplotlib ultralytics
!apt install zip unzip -y
!apt install unar -y

这些命令确保你可以处理视频、图片、数据并运行YOLOv8模型。

然后导入库:

import os, sys
import cv2, glob, json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

然后,下载并解压训练集和测试集数据。

!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

3数据处理

接下来,我们读取训练集中的标注文件,并使用OpenCV读取视频帧。通过迭代每一帧,并将目标标注框绘制在视频帧上,可以可视化数据的标注情况。

train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()
bbox = [746, 494, 988, 786]
cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

4.数据集准备

为YOLOv8模型训练准备数据集,首先需要创建存放数据集的文件夹结构,然后编写yolo.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')

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格式的数据标签文件。每个标签文件中包含目标类别和对应的边界框位置。

category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]

# 处理训练集
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(f'./yolo-dataset/train/{anno_path.split("/")[-1][:-5]}_{frame_idx}.jpg', frame)
        if len(frame_anno) != 0:
            with open(f'./yolo-dataset/train/{anno_path.split("/")[-1][:-5]}_{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-dataset/val目录中。

5.模型训练

下载YOLOv8预训练模型,并利用自定义的数据集进行训练。训练过程包括设置训练参数,如epochs和batch size。

!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)
6.模型推理与结果生成

训练完成后,使用训练好的模型对测试集视频进行推理,提取每一帧中的目标,并保存检测结果为JSON文件。

if not os.path.exists('result/'):
    os.mkdir('result')

model = YOLO("runs/detect/train/weights/best.pt")

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(f'./result/{path.split("/")[-1][:-4]}.json', 'w', encoding='utf-8') as up:
        json.dump(submit_json, up, indent=4, ensure_ascii=False)

最后,压缩并导出推理结果。

!zip -r result.zip result/

项目总结

本项目主要围绕计算机视觉竞赛中的目标检测任务展开,使用了YOLOv8模型来实现视频中的特定目标检测和分类。整个项目从数据准备、模型训练到最终的推理结果生成,均进行了详细的实现与记录。

项目成果
  1. 数据准备

    • 下载并解压训练集和测试集。
    • 通过处理视频帧与对应的标注文件,成功构建了适用于YOLOv8的训练和验证数据集。
  2. 模型训练

    • 使用预训练的YOLOv8模型对定制的数据集进行训练,设定了训练参数(如epochs和batch size)并成功生成了模型权重文件。
  3. 模型推理

    • 利用训练好的模型对测试视频进行了推理,准确检测出了视频中的目标,生成了包含边界框坐标、类别和置信度的检测结果。
  4. 结果输出

    • 最终生成的推理结果以JSON文件格式保存,并压缩打包为一个ZIP文件,便于进一步提交和分析。
项目亮点
  • 自动化流程:从数据下载、解压、处理,到模型训练与推理,整个流程实现了较高的自动化,极大地提高了工作效率。
  • 数据处理的细致性:项目中对每一帧视频的处理都进行了详尽的标注和检测,保证了数据的准确性和一致性。
  • 模型性能:利用YOLOv8的高效性与精度优势,成功完成了复杂环境下的目标检测任务。
挑战与解决方案
  • 数据处理复杂性:处理包含大量帧与不同标注的训练数据时,确保每个标注的准确性是一个挑战。通过使用Pandas对标注文件进行精确处理,最终解决了这个问题。
  • 模型训练时间:在高分辨率的图片上进行训练时,训练时间较长。通过适当调整batch size和图像尺寸,确保了模型在合理时间内完成训练。
未来改进方向
  • 模型优化:可以尝试调节模型的超参数,如学习率和epochs,以进一步优化模型的检测精度和速度。
  • 数据增强:在数据集的构建过程中,可以引入数据增强技术,如随机裁剪、旋转、镜像等,增加数据的多样性,提高模型的泛化能力。
  • 多模型对比:除了YOLOv8外,还可以尝试其他目标检测模型,如Faster R-CNN、RetinaNet等,并进行性能对比分析。

通过本项目,进一步巩固了对YOLO模型的理解与应用,同时也为计算机视觉竞赛任务的实现积累了宝贵经验。这将为未来的相关研究和开发工作打下坚实基础。

                                                                                                                                        #啊岭撰

  • 11
    点赞
  • 11
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值