yolo方案baseline
Datawhale (linklearner.com) 飞书文档
apt install git-lfs
git lfs install
git clone https://www.modelscope.cn/datasets/Datawhale/AI_Camp5_baseline_CV.git
# 读取训练集视频
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)
# 计算yolo标注格式
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
这里用到了YOLO(You Only Look Once),这是一种流行的实时目标检测系统,YOLO模型的核心思想是将目标检测任务视为一个单一的回归问题,通过一个卷积神经网络(CNN)直接从图像像素到边界框坐标和类别概率的映射。YOLO模型经过了多次迭代,包括YOLOv2(YOLO9000)、YOLOv3和YOLOv4等版本,每个版本都在性能和速度上有所提升,同时也引入了一些新的技术,如更深的网络结构、更好的锚框机制、多尺度特征融合等。
这段代码是用于处理视频数据并将其转换为YOLO格式的标注文件的Python脚本。YOLO需要特定的标注格式来训练模型。以下是代码的主要功能和步骤:
使用zip
函数将训练集的标注文件路径(train_annos
)和视频文件路径(train_videos
)配对。
迭代前 5 对标注文件和视频文件路径。
读取每个标注文件(anno_path
)为Pandas DataFrame(anno_df
)。
使用OpenCV(cv2
)打开视频文件(video_path
)。
初始化frame_idx
为0,用于记录当前处理的视频帧索引。
使用while
循环读取视频中的所有帧。如果读取失败(ret
为False
),则跳出循环。
获取当前帧的高度和宽度。
将当前帧保存为图片文件,文件名基于标注文件名和帧索引。
检查当前帧是否有标注信息(frame_anno
)。
如果有标注信息,打开一个文本文件用于写入YOLO格式的标注数据。
对于每个标注,计算类别索引(category_idx
)和YOLO格式的坐标(x_center
, y_center
, width
, height
)。
检查x_center
是否大于1,如果是,则打印出边界框(bbox
)的值。这可能是一个错误检查,以确保坐标值在正确的范围内。
将YOLO格式的标注数据写入文本文件。
增加frame_idx
以处理下一帧。
Ultraalytics 是一个提供多种计算机视觉模型的库,包括 YOLO 系列。这段代码是一个简单的训练启动示例
from ultralytics import YOLO
# 设置模型版本
model = YOLO("yolov8n.pt")
# 设定数据集和训练参数
results = model.train(data="yolo-dataset/yolo.yaml", epochs=2, imgsz=1080, batch=16)
加载模型:使用 YOLO("yolov8n.pt")
创建一个 YOLO 模型实例。这里的 "yolov8n.pt"
是预训练模型的文件路径,yolov8n
表示 YOLOv8 的一个较小的模型版本,.pt
是 PyTorch 模型文件的扩展名。
训练模型:通过调用模型实例的 train
方法来训练模型。方法的参数如下:
data
:指定数据集配置文件的路径,这里是 "yolo-dataset/yolo.yaml"
。这个 YAML 文件应该包含了训练和验证数据集的路径、类别名称、锚框尺寸等信息。epochs
:训练的轮数,这里设置为 2 轮。imgsz
:输入图像的尺寸,这里设置为 1080 像素。这个尺寸应该与你的标注数据匹配。batch
:每个批次的图像数量,这里设置为 16。train
方法将返回一个包含训练结果的对象。这个对象可能包含训练过程中的各种统计信息,如损失值、精确度等。box_loss
是边界框回归损失,用于评估预测的边界框与真实边界框之间的差异。
cls_loss
是分类损失,用于评估类别预测的准确性。
dfl_loss
是防御性损失,用于提高模型的泛化能力。
这是baseline方案代码的学习,后面进阶方案会继续!
参与评论
手机查看
返回顶部