train_annos = glob.glob('训练集(有标注第一批)/标注/*.json')

train_videos = glob.glob('训练集(有标注第一批)/视频/*.mp4')

train_annos.sort(); train_videos.sort();

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

glob.glob 是 Python 中 glob 模块提供的一个函数，它的作用是文件名模式匹配，即搜索文件系统中与指定模式（模式字符串）匹配的所有路径名。

这里就是把所有文件提取出来，然后按文件名（字母顺序）排序。这在windows系统操作者看来是理所应当的事情，但在python中，你仍然需要一些步骤去完成它们以便应对更复杂的情况。

（3）批量读入全部数据

读入训练集

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

喂喂喂，这真的好吗，不知道你发现没有，我发现这里读数据的时候baseline代码没有读取全部数据，要是你不好好学习直接跑baseline的话，你的分数肯定不会很高啊，原来datawhale官方为了让大家认真学习真的用心良苦啊！

不过不要紧，你还可以看小虢的解说来发现这一点，然后修改代码跑一个不错的分数，截止到笔记发表的时候，排行榜有将近100人没有发现这一点，看了小虢的解说，你的名次可以前进100名！（手动狗头）

上面是开个玩笑，baseline这么写主办方主要是时间和内存方面的考量，毕竟视频处理起来还是非常耗费时间的啦！另外提醒一下，如果你在使用主办方提供算力资源中的最基础版本，大概30批视频就已经够呛了，请朋友们自行留意硬件使用限制与资费情况，合理控制训练集与测试集的规模，合理划分训练集与测试集。

好了，话归正题，这段代码主要做了这些工作：

1.读取训练数据：通过zip函数同时获取训练注释（标注）路径train_annos和训练视频路径train_videos。

2.读取标注数据：使用pd.read_json读取每个视频对应的注释文件，存储在anno_df变量中。

3.使用cv2.VideoCapture打开视频文件，并逐帧读取。

4.循环读取帧：使用while循环和cap.read()函数读取视频的每一帧，如果读取失败（即ret为False），则跳出循环。

5.保存帧图像：将读取的帧保存为JPEG图像文件，路径为./yolo-dataset/train/，文件名由原注释文件名和帧编号组成。

6.处理标注：如果当前帧存在注释（即frame_anno非空），则执行以下操作：打开一个文本文件用于写入处理后的注释数据；遍历注释中的每一类别和边界框；将边界框的坐标转换为YOLO格式（类别索引、中心点的x和y坐标、宽高）；写入转换后的注释数据到文本文件；将原始的边界框坐标(x_min, y_min, x_max, y_max)转换为中心点坐标(x_center, y_center)和宽高(width, height)，这些值被归一化到[0, 1]区间。

7.错误检查：如果检测到x_center的值大于1，打印出边界框的坐标值，归一化后的坐标不应该大于1。

8.帧编号递增：在处理完当前帧的注释后，递增frame_idx以准备处理下一帧。

4.训练模型

1.加载模型

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

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

os.mkdir('result')

from ultralytics import YOLO

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

import glob

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 # Boxes object for bounding box outputs

masks = result.masks # Masks object for segmentation masks outputs

keypoints = result.keypoints # Keypoints object for pose outputs

probs = result.probs # Probs object for classification outputs

obb = result.obb # Oriented boxes object for OBB outputs

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:code>

json.dump(submit_json, up, indent=4, ensure_ascii=False)

下次笔记，我将带着大伙一起炼丹，争取获得更好的分数和排名！

笔者水平有限，欢迎各位大佬批判指正！