ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

9类岩石检测数据集:YOLO训练配置、可视化脚本与避坑指南

9类岩石检测数据集:YOLO训练配置、可视化脚本与避坑指南 简介这份资源面向从事岩石识别、地质图像分析及目标检测实践的研究者与开发者提供一套可直接投入训练的YOLO格式岩石检测数据集省去自行采集与标注的繁琐流程。包内共2000个文件以1999个txt标注文件和1个Python可视化脚本为主压缩包约637MB标注采用YOLO相对坐标格式涵盖玄武岩、石灰岩、沉积岩等9个类别。数据按YOLOv5目录结构组织训练集含12501张图片及对应标签验证集含1104张图片及对应标签图像为640×640的RGB大分辨率样本并经过四图融合的马赛克增强边界框完整、每图含多个目标。配套的show.py脚本无需修改即可运行随机传入一张图片便能绘制并保存边界框便于快速核验标注质量。目前已有719人学习下载适合需要快速搭建岩石检测基线、验证模型效果或开展迁移学习实验的读者直接使用。1. 岩石检测数据集到底长什么样9 类样本、划分文件与可视化脚本一次说清做地质勘探、隧道掌子面编录或者矿山皮带异物分拣的同行大概率都动过「用 YOLO 自动识别岩石类别」的念头。真上手才发现卡住你的往往不是模型结构而是数据网上公开的岩石图像要么类别混乱要么没有标注要么标注格式和 YOLO 对不上。这个标题讲的就是一套已经整理好的 9 类岩石检测数据集——它包含划分好的训练/验证/测试集、一份类别 class 文件以及一个能把标注框画回原图的可视化脚本。它解决的是「从零标注到能开训」之间最耗时的那段脏活适合想快速跑通 YOLO 训练流程、验证自己改进思路的算法工程师也适合需要给岩石识别系统做原型的地质信息化团队。下面我按「数据长什么样 → 怎么接进 YOLO → 怎么验证没标错 → 坑在哪」的顺序把能直接抄的步骤和参数讲透。2. 拆开数据集目录结构、class 文件与标注格式的对应关系拿到一个检测数据集第一件事不是急着写训练脚本而是把目录结构和标注格式摸清楚。格式对不上后面 loss 不下降、mAP 恒为 0 的玄学问题会一个接一个。这一章先把数据集的骨架拆开让你知道每个文件是干什么的、YOLO 到底认哪种写法。2.1 典型目录结构与各文件职责一套划分好的 YOLO 岩石检测数据集常见做法是长这样rock_dataset/ ├── images/ │ ├── train/ # 训练集图像 │ ├── val/ # 验证集图像 │ └── test/ # 测试集图像 ├── labels/ │ ├── train/ # 与 images/train 一一对应的 txt 标注 │ ├── val/ │ └── test/ ├── classes.txt # 9 个类别名每行一个 └── visualize.py # 标注可视化脚本这里的关键约束是images/train/xxx.jpg必须对应labels/train/xxx.txt文件名不含扩展名完全一致。YOLO 训练时是按文件名去配对图像和标签的差一个字符这张图就会被当成无标签样本丢掉或者直接报找不到标签。classes.txt是类别索引的「字典」第 0 行对应类别 id 0第 1 行对应 id 1以此类推。很多人训练完发现预测框的类别名全是错的八成是classes.txt的顺序和标注时用的顺序对不上。提示划分好的数据集里train/val/test 三个子目录的图像数量比例常见是 7:2:1 或 8:1:1。如果 val 集太小比如不到 50 张验证指标会剧烈抖动别急着调模型先把 val 补到至少 10% 再看。2.2 YOLO 标注格式归一化坐标怎么算YOLO 检测的标签是每行一个目标格式为class_id x_center y_center width height五个值用空格分隔后四个都是归一化到 0~1的相对值不是像素。换算关系是# 假设图像宽 w1280高 h960 # 标注框左上角 (x1, y1)(320, 240)右下角 (x2, y2)(640, 480) x_center (320 640) / 2 / 1280 # 0.375 y_center (240 480) / 2 / 960 # 0.375 width (640 - 320) / 1280 # 0.25 height (480 - 240) / 960 # 0.25 # 最终写入: 0 0.375 0.375 0.25 0.25这段换算看着简单但它是后面所有排查的基准。归一化坐标的好处是图像缩放、letterbox 填充后标签不用改坏处是一旦你手动改过图像尺寸却没同步标签框就会整体偏移。岩石图像里目标往往占画面比例大、边界模糊标注时框稍微松一点紧一点对训练影响比清晰物体更明显所以坐标精度要盯紧。2.3 class 文件与类别不平衡的初步判断classes.txt里 9 个类别实际分布几乎不可能均匀。岩石类别里常见某几类比如花岗岩、砂岩样本多某几类比如某些变质岩样本少。开训前先统计一下每类框的数量import os from collections import Counter label_dir rock_dataset/labels/train counter Counter() for fname in os.listdir(label_dir): if not fname.endswith(.txt): continue with open(os.path.join(label_dir, fname)) as f: for line in f: line line.strip() if line: counter[int(line.split()[0])] 1 for cid in sorted(counter): print(fclass {cid}: {counter[cid]} boxes)跑完你会得到一张类别分布表。如果最多和最少差 10 倍以上直接开训大概率出现「多数类学得很好、少数类几乎不召回」。这时候要么在数据增强上给少数类加权要么用 mosaic、copy-paste 这类增强补样本别指望模型自己平衡。这一步花五分钟能省掉后面几小时的无效训练。3. 把数据集接进 YOLOdata.yaml 配置与训练命令数据摸清楚了接下来是让它真正跑起来。这一章给出一份能直接用的data.yaml讲清每个字段的含义再给一条最小训练命令和关键参数说明。不管你用的是 YOLOv5、YOLOv8 还是更新的版本配置思路是一致的差异主要在命令行的写法上。3.1 data.yaml 的字段与路径写法YOLO 系列训练时读的是一个 YAML 配置文件典型内容如下# rock_data.yaml path: /home/user/rock_dataset # 数据集根目录 train: images/train # 相对 path 的训练图像目录 val: images/val test: images/test nc: 9 # 类别数量必须和 classes.txt 行数一致 names: # 类别名顺序必须和标注里的 class_id 对应 0: 花岗岩 1: 砂岩 2: 石灰岩 3: 页岩 4: 大理岩 5: 片麻岩 6: 玄武岩 7: 石英岩 8: 板岩几个容易翻车的点path建议写绝对路径相对路径在不同工作目录下启动训练会找不到数据train/val是相对path的路径不是相对你当前终端nc和names的条目数必须严格相等少一个多一个都会在加载时直接报错。类别名用中文没问题但要注意 YAML 文件保存为 UTF-8否则可视化时类别名会变乱码。注意如果你的标注里 class_id 是从 1 开始而不是从 0 开始训练不会报错但所有类别会整体错位一位最后一类永远学不到。开训前务必抽查一个标签文件确认最小 class_id 是 0。3.2 最小训练命令与关键参数以常见的 YOLOv8 风格命令为例一条能跑通的最小训练命令是yolo detect train \ datarock_data.yaml \ modelyolov8n.pt \ epochs100 \ imgsz640 \ batch16 \ device0 \ projectruns/rock \ nameexp1逐项说明modelyolov8n.pt是从预训练权重起步小数据集上比从头训收敛快得多epochs100是上限实际看早停imgsz640是输入分辨率岩石图像如果目标很小可以提到 960 或 1280但显存占用会明显上升batch16要按显存调爆显存就减半device0指定第一块 GPU多卡用device0,1。如果你用的是 YOLOv5命令换成python train.py --data rock_data.yaml --weights yolov5n.pt --epochs 100 --img 640 --batch 16参数含义基本对应。训练启动后重点看三件事一是日志里train:和val:后面跟的图像数量应该和你划分的数量对得上对不上说明路径或配对有问题二是第一个 epoch 的 loss 是否在下降如果几个 epoch 都不动回去查标签格式三是验证集的 mAP 曲线正常是前期快速上升后趋缓。岩石类别之间纹理相似度高mAP 上升会比 COCO 那种差异大的数据集慢别看到 20 个 epoch 还没到 0.5 就慌。3.3 从预训练权重起步还是从头训小规模岩石数据集几千张量级强烈建议从预训练权重起步。原因很直接预训练模型已经学到了边缘、纹理这类通用特征岩石识别本质上吃的就是纹理和颜色分布迁移过来能省掉大量数据。从头训在几千张图上很容易过拟合训练集 loss 一路降、验证集 mAP 卡住不动这就是典型的过拟合信号。如果你要做的是「YOLO 改进」类工作比如换 backbone、加注意力模块那更要从预训练权重加载只让改动部分随机初始化其余部分用预训练参数。常见做法是加载权重时用strictFalse让新增模块的参数随机初始化、其余复用。这样对比实验才有意义否则你分不清涨点是来自改进还是来自训练策略。4. 可视化脚本把标注框画回原图验证有没有标错训练前最值钱的一步是把标注画回原图肉眼过一遍。标注错位、类别标反、框漏标这些问题在数字指标上很难第一时间发现但看图一眼就能看出来。这一章给一个能直接用的可视化脚本并讲清它怎么帮你定位问题。4.1 可视化脚本与逐行说明import os import cv2 IMG_DIR rock_dataset/images/train LBL_DIR rock_dataset/labels/train CLASSES open(rock_dataset/classes.txt, encodingutf-8).read().splitlines() COLORS [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255), (128, 0, 128), (0, 128, 255), (128, 128, 0)] for fname in os.listdir(IMG_DIR): if not fname.lower().endswith((.jpg, .png, .jpeg)): continue img_path os.path.join(IMG_DIR, fname) lbl_path os.path.join(LBL_DIR, os.path.splitext(fname)[0] .txt) img cv2.imread(img_path) if img is None: print(f[读取失败] {img_path}) continue h, w img.shape[:2] if not os.path.exists(lbl_path): print(f[缺标签] {fname}) continue with open(lbl_path, encodingutf-8) as f: for line in f: parts line.strip().split() if len(parts) ! 5: print(f[格式异常] {fname}: {line.strip()}) continue cid, xc, yc, bw, bh int(parts[0]), *map(float, parts[1:]) x1 int((xc - bw / 2) * w) y1 int((yc - bh / 2) * h) x2 int((xc bw / 2) * w) y2 int((yc bh / 2) * h) color COLORS[cid % len(COLORS)] cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) label CLASSES[cid] if cid len(CLASSES) else fid{cid} cv2.putText(img, label, (x1, max(y1 - 5, 15)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) cv2.imwrite(fvis_{fname}, img)逻辑说明脚本遍历图像目录按同名规则找标签文件把归一化坐标反算回像素坐标后画框。COLORS按类别 id 取色同一类颜色一致方便你一眼看出类别是否标反。三个打印分支分别对应「图像读不出来」「缺标签」「标签行格式不对」跑一遍就能把数据集里的脏样本筛出来。参数上IMG_DIR和LBL_DIR换成你自己的路径即可CLASSES直接读classes.txt保证类别名和训练时一致。4.2 从可视化结果里能看出什么画完图别只扫一眼就过重点看四类问题。第一类是框整体偏移通常是对图像做过 resize 或裁剪但没同步改标签第二类是框大小明显不合理比如把整张图框成一个目标多半是标注时偷懒第三类是类别标反同一块岩石在不同图里颜色不一致说明标注标准不统一第四类是漏标图里明显有多块岩石但只有一个框。这四类问题里漏标和类别标反对 mAP 的伤害最大而且靠指标很难定位只能靠看图。提示可视化输出建议单独放一个目录别和原图混在一起。跑完抽 30~50 张覆盖各个类别的图看一遍比随机看 200 张更有效。5. 避坑与排查岩石检测训练里最常见的 5 个翻车现场前面把流程走通了这一章专门收那些「跑起来了但结果不对」的问题。岩石检测因为纹理相似、类别边界模糊踩的坑和通用目标检测略有不同下面 5 条都是实操里反复出现的。5.1 现象训练 loss 正常下降但验证 mAP 一直是 0原因最常见的是验证集标签路径没配对YOLO 把 val 集当成全无标签计算 mAP 时没有真值可比。其次是data.yaml里val指向了空目录或路径写错。还有一种隐蔽情况是标签文件编码不是 UTF-8读取时静默失败。解决先看训练日志里val:后面的图像数量是否为 0 或异常小再手动跑一遍 4.1 的可视化脚本把IMG_DIR换成 val 目录能正常画出框说明标签没问题最后检查data.yaml的val字段是不是相对path的路径。三步基本能定位。5.2 现象某一类永远检测不出来其他类正常原因class_id 错位。标注时用了 1~9而data.yaml里names是 0~8导致最后一类没有对应标注第一类多了一堆错位样本。也可能是该类样本太少被多数类压制。解决抽查几个标签文件确认 class_id 范围是 0~8统计每类框数量如果某类少于总框数的 2%考虑增强或重采样。改完 class_id 后必须重新训练不能靠继续训练修正。5.3 现象验证集 mAP 波动很大每次跑结果差很多原因验证集太小。岩石数据集如果 val 只有几十张随机划分一次指标就跟着这几张图走波动是必然的。另一个原因是 batch 太小导致 BN 统计不稳定这个在 batch 小于 8 时比较明显。解决把 val 补到至少占总量的 10%并且划分时按类别分层抽样保证每个类别在 val 里都有代表。batch 尽量提到 16 以上显存不够就用梯度累积模拟大 batch。5.4 现象模型把背景当成岩石误检一堆原因负样本纯背景图太少或者标注时把一些模糊区域也框成了目标。岩石图像里背景和目标的纹理差异本来就小模型很容易学到「有纹理就是岩石」这种错误关联。解决往训练集里加一批纯背景图标签文件为空比例控制在 5%~10%同时复查标注把边界模糊、自己都拿不准的区域去掉宁可漏标不要错标。误检多的时候先怀疑数据再怀疑模型。5.5 现象换了自己的图像测试框位置整体偏移原因测试图像的预处理和训练时不一致。训练时 YOLO 默认做 letterbox 填充保持长宽比如果你自己推理时直接 resize 到 640×640坐标映射就会错。另一个原因是测试图像被 EXIF 旋转信息影响读进来方向就不对。解决推理时用和训练一致的预处理或者直接用官方predict接口它内部处理好了 letterbox。EXIF 问题用cv2.imread读之前先做方向校正或者统一用 PIL 读取并ImageOps.exif_transpose。6. 进阶用混淆矩阵和单类 AP 定位岩石类别的混淆对训练跑通、mAP 看着还行之后真正决定这套数据集能不能用于生产的是「哪些类别之间在互相混淆」。岩石类别里花岗岩和片麻岩、砂岩和石英岩纹理上本来就接近整体 mAP 会被这些混淆对拉低但你看总指标根本不知道问题出在哪。我的习惯是训练完先看混淆矩阵再逐类看 AP。YOLO 训练结束后会在runs/目录下生成confusion_matrix.png和confusion_matrix_normalized.png。归一化那张更值得看对角线是正确分类比例非对角线就是混淆。如果发现第 0 类和第 5 类互相混淆严重说明这两类在特征空间里没分开接下来有三个方向可以试。第一个方向是回到数据检查这两类的标注标准是否一致。很多时候混淆不是模型的问题是标注时这两类本身就标得含糊。第二个方向是针对性增强对混淆的两类做更强的颜色抖动和旋转逼模型学更鲁棒的特征。第三个方向才是改模型比如加注意力模块或者换更强的 backbone但这一步应该放在数据和增强都试过之后。逐类 AP 可以直接从验证输出里拿也可以自己写脚本统计# 伪代码按类别统计 AP定位短板类别 # 训练时加 --verbose 或从 results.csv 里读每类指标 # 重点关注 AP 最低的 2~3 类它们决定了整体上限我的血泪经验是岩石检测里整体 mAP 从 0.7 提到 0.8靠的往往不是换模型而是把混淆最严重的那两类的标注重新过一遍。数据质量的天花板比模型结构的天花板低得多也更容易够到。这套 9 类岩石数据集的价值就在于它把划分、class 文件、可视化脚本都给你备好了你可以把省下来的时间全砸在标注复查和类别平衡上。先跑通再看混淆矩阵再回头改数据这个顺序别反。希望帮到你。本文还有配套的精品资源点击获取
返回列表