Datawhale AI 夏令营 cv task2笔记

# 安装和导入所需的库
!/opt/miniconda/bin/pip install opencv-python pandas matplotlib ultralytics
import os
import cv2
import glob
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# 读取和处理标注文件和视频
train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
pd.read_json('训练集(有标注第一批)/标注/45.json')

video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()  # 读取第一帧
frame.shape
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数据集目录结构
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('./') + '/'

# 生成YOLO格式的配置文件
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')
train_videos = glob.glob('训练集(有标注第一批)/视频/*.mp4')
train_annos.sort()
train_videos.sort()

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('./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

# 处理验证集数据
for anno_path, video_path in zip(train_annos[-3:], train_videos[-3:]):
    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

# 下载YOLOv8s的预训练权重
!wget http://mirror.coggle.club/yolo/yolov8s-v8.2.0.pt -O yolov8s.pt

# 设置CUDA环境和警告忽略
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import warnings
warnings.filterwarnings('ignore')

# 使用YOLOv8s模型进行训练
from ultralytics import YOLO
model = YOLO("yolov8s.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=30, imgsz=1080, batch=16)

# 保存训练好的模型
best_model_path = "runs/detect/train/weights/best.pt"

# 加载已经训练好的YOLO模型
model = YOLO(best_model_path)

# 确保结果目录存在
if not os.path.exists('result/'):
    os.mkdir('result')

# 遍历测试集目录中的所有MP4文件
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  # 获取边界框输出
        masks = result.masks  # 获取分割掩码输出（如果有）
        keypoints = result.keypoints  # 获取关键点输出（如果有）
        probs = result.probs  # 获取分类输出（如果有）
        obb = result.obb  # 获取定向边界框输出（如果有）

        # 如果没有检测到任何物体，则跳过此帧
        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()

        # 遍历每个检测结果，并添加到 JSON 列表中
        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)

# 压缩结果文件夹
shutil.make_archive('result', 'zip', 'result')

print("训练和推理完成，并生成result.zip文件。")

数据预处理与模型训练

本任务采用了一系列标准化的数据预处理步骤，以确保输入数据的质量和一致性。首先，利用OpenCV库对视频文件进行逐帧处理，并根据标注文件提取每一帧的目标边界框信息。我们从标注文件中读取目标的位置信息，并将其转化为YOLO模型所需的标注格式，包括类别标签、边界框中心坐标以及边界框的宽度和高度。处理后的帧图像和对应的标注信息分别存储在“训练”和“验证”两个目录下，以供后续模型训练使用。

为了构建YOLOv8模型的训练集，我们生成了YOLO格式的配置文件，该文件包括数据集的路径及类别名称。配置文件的生成确保了模型训练过程中对数据集的正确引用，同时也为模型推理阶段提供了基础。

YOLOv8模型的训练

在模型训练阶段，我们首先下载了YOLOv8s的预训练权重文件。此预训练模型为YOLOv8s变体，该变体在保持模型轻量化的同时具有较强的特征提取能力。训练过程在包含30个epoch的条件下进行，输入图像尺寸设置为1080，并采用批量大小为16的配置，以平衡训练速度与模型性能。

模型训练完成后，最佳的模型权重文件被保存下来，以备后续的推理任务使用。模型训练的整个过程使用了GPU加速，并且通过忽略不必要的警告信息来保证代码的整洁和执行效率。

模型推理与结果处理

推理阶段涉及对测试集中的视频文件逐帧进行处理。对于每一帧，YOLOv8模型将输出多个目标检测结果，包括边界框、类别标签、置信度分数等。这些信息被汇总到一个JSON文件中，便于后续的分析和可视化。为了确保推理结果的准确性，设置了置信度阈值，低于该阈值的检测结果将被忽略。

在处理每个视频文件时，我们逐帧收集检测结果，并将这些结果以JSON格式保存到本地。每个视频文件的推理结果被单独保存为一个JSON文件，文件名与原始视频文件名对应，以便于结果的组织和管理。