DataWhale AI夏令营 2024大运河杯-数据开发应用创新赛

本文主要是介绍DataWhale AI夏令营 2024大运河杯-数据开发应用创新赛，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

话不多说直接开始

baseline分析

这里我们忽略数据、模型下载的单元格
导入数据处理的一些包

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

读取下载的数据，并查看一下json的格式。

train_anno = json.load(open('训练集(有标注第一批)/标注/45.json', encoding='utf-8'))
train_anno[0], len(train_anno)

用pandas读取数据查看数据格式

pd.read_json('训练集(有标注第一批)/标注/45.json')

读取视频，使用VideoCapture对数据进行切帧处理。

video_path = '训练集(有标注第一批)/视频/45.mp4'
cap = cv2.VideoCapture(video_path)
while True:# 读取下一帧ret, frame = cap.read()if not ret:breakbreak  

根据json的信息，展示一张画框的图片

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所需的数据集。

构建YOLO数据集

yolo数据集的格式为一个data文件夹下包含三个内容,train; val; yolo.yaml,其中train和val不在介绍,yolo.yaml主要包含数据涉及到的标签信息。
我这里是吧数据放在/root/data文件夹下了，因为切帧的图片数据很多需要的空间后的云的系统盘空间不够。大家可以参考。

if not os.path.exists('/root/data/yolo-dataset/'):os.mkdir('/root/data/yolo-dataset/')
if not os.path.exists('/root/data/yolo-dataset/train'):os.mkdir('/root/data/yolo-dataset/train')
if not os.path.exists('/root/data/yolo-dataset/val'):os.mkdir('/root/data/yolo-dataset/val')dir_path = os.path.abspath('./') + '/'# 需要按照你的修改path
with open('/root/data/yolo-dataset/yolo.yaml', 'w', encoding='utf-8') as up:up.write(f'''
path: /root/data/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 = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]

我这里按照8:2划分训练集和验证集，一共应该是52组数据，划分之后是42:10
下面在代码中给出了逐行的注释，大家自行食用即可。

for anno_path, video_path in zip(train_annos[:42], train_videos[:42]):print(video_path)# 使用Pandas读取JSON格式的标注文件，返回一个DataFrame对象anno_df = pd.read_json(anno_path)   # 使用OpenCV打开视频文件，准备逐帧读取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]# 将当前帧保存为JPEG图像文件cv2.imwrite('/root/data/yolo-dataset/train/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)# 检查当前帧有没有标注信息if len(frame_anno) != 0:# 创建并打开一个与当前帧图像同名的文本文件，准备写入YOLO格式的标签with open('/root/data/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# 计算标注框的中心点横纵坐标，并归一化到 [0, 1] 之间x_center = (x_min + x_max) / 2 / img_widthy_center = (y_min + y_max) / 2 / img_height# 计算框的宽和高，并归一化width = (x_max - x_min) / img_widthheight = (y_max - y_min) / img_heightif x_center > 1:print(bbox)# 将YOLO格式的标注信息写入标签文件up.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')# 处理下一帧frame_idx += 1

构建验证集，这部分代码直接看上一个即可一样的基本都是

for anno_path, video_path in zip(train_annos[-10:], train_videos[-10:]):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:breakimg_height, img_width = frame.shape[:2]frame_anno = anno_df[anno_df['frame_id'] == frame_idx]cv2.imwrite('/root/data/yolo-dataset/val/' + anno_path.split('/')[-1][:-5] + '_' + str(frame_idx) + '.jpg', frame)if len(frame_anno) != 0:with open('/root/data/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 = bboxx_center = (x_min + x_max) / 2 / img_widthy_center = (y_min + y_max) / 2 / img_heightwidth = (x_max - x_min) / img_widthheight = (y_max - y_min) / img_heightup.write(f'{category_idx} {x_center} {y_center} {width} {height}\n')frame_idx += 1

开始训练

baseline使用的是yolov8n进行训练，在这里epoch代表训练的轮数，imgsz代表输入模型图像大小，batch代表一次梯度更新使用多少张图片

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="/root/data/yolo-dataset/yolo.yaml", epochs=15, imgsz=1080, batch=16)

创建结果目录

category_labels = ["非机动车违停", "机动车违停", "垃圾桶满溢", "违法经营"]if not os.path.exists('result/'):os.mkdir('result')

对测试集视频文件的处理，通过预训练的YOLO模型对每个视频的每一帧进行检测，并将检测结果保存为JSON格式的文件。

from ultralytics import YOLO
# 使用训练好的模型进行预测
model = YOLO("runs/detect/train/weights/best.pt")
import globfor path in glob.glob('测试集/*.mp4'):# 保存结果生成的json文件submit_json = []# 对视频文件进行推理，conf=0.05设置了最低置信度阈值results = model(path, conf=0.05, imgsz=1080,  verbose=False)for idx, result in enumerate(results):boxes = result.boxes  # Boxes object for bounding box outputsmasks = result.masks  # Masks object for segmentation masks outputskeypoints = result.keypoints  # Keypoints object for pose outputsprobs = result.probs  # Probs object for classification outputsobb = result.obb  # Oriented boxes object for OBB outputsif 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()# 写入submitfor 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)

优化思路

这篇关于DataWhale AI夏令营 2024大运河杯-数据开发应用创新赛的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---