
简介面向目标检测学习者和岩石识别应用开发者这份资源提供了一套可直接投入训练的YOLO格式岩石检测数据集省去自行采集、清洗与标注的繁琐流程。数据按YOLOv5目录结构组织图像统一为640×640的RGB大分辨率并经过四图融合的马赛克增强标注框完整每张图包含多个目标标注采用classes与x_centre、y_centre、w、h的相对坐标格式。数据集覆盖玄武岩、石灰岩、沉积岩等9个类别划分为12501张训练图和1104张验证图均配有对应txt标签另附9类别文本文件。压缩包共2000个文件以1999个txt标注与1个可视化py脚本为主整体约637MB。其中show.py可随机读取一张图片绘制边界框并保存到当前目录无需修改即可运行便于快速核验标注质量。目前已有719人学习下载适合需要现成数据开展岩石检测训练与验证的读者。1. 岩石检测数据集怎么选9 类岩石的 YOLO 落地为什么值得做做地质勘探、矿山巡检或者岩芯编录的团队几乎都会卡在同一个环节现场拍回来的岩石图像堆成山人工分类又慢又主观。9 种岩石检测这个方向之所以值得投入是因为它把「岩石识别」从实验室的岩相学搬到了工程现场——你需要的不是一篇论文而是一套能直接喂给 YOLO 训练的数据集外加类别定义和可视化脚本让模型跑起来之后能告诉你画面里到底是花岗岩还是玄武岩。这套方案适合三类人一是做地质 AI 应用但缺标注数据的开发者二是想把岩石分类自动化塞进巡检流程的工程师三是拿它当 YOLO 训练练手项目、顺便理解多类别检测全流程的学生。核心痛点从来不是模型选 YOLOv8 还是 YOLOv11而是数据集怎么划分、class 文件怎么写、可视化脚本怎么验证标注没跑偏。这三件事做对了后面训练和部署才有意义。2. 9 类岩石数据集的结构拆解与划分逻辑2.1 为什么岩石检测的数据集不能随便按 8:2 切岩石图像有个很坑的特性同一类岩石在不同光照、不同风化程度下纹理差异可能比两类岩石之间的差异还大。如果你直接按随机种子把全部图片打乱做 8:2 划分很可能出现训练集里全是新鲜面的花岗岩验证集里全是风化面的花岗岩模型学到的不是岩石类别而是风化程度。常见做法是按「采集批次」或「露头位置」做分组划分保证同一批次的图像要么全在训练集要么全在验证集。具体操作上我一般会先给每张图打一个 group_id比如site_A_001然后按 group 做 7:2:1 划分而不是按单张图。import os import random from collections import defaultdict # 假设图片按类别存放在 data/raw/类别名/ 下 # 先按采集批次分组这里用文件名前缀模拟 group_id def split_by_group(raw_dir, train_ratio0.7, val_ratio0.2): groups defaultdict(list) for cls_name in os.listdir(raw_dir): cls_dir os.path.join(raw_dir, cls_name) if not os.path.isdir(cls_dir): continue for img_name in os.listdir(cls_dir): # 文件名格式: siteA_001.jpg - group_id siteA group_id img_name.split(_)[0] groups[group_id].append((cls_name, img_name)) group_keys list(groups.keys()) random.shuffle(group_keys) n_total len(group_keys) n_train int(n_total * train_ratio) n_val int(n_total * val_ratio) train_groups group_keys[:n_train] val_groups group_keys[n_train:n_train n_val] test_groups group_keys[n_train n_val:] return train_groups, val_groups, test_groups这段代码的关键参数是train_ratio和val_ratio岩石检测场景下我通常设 0.7 和 0.2留 0.1 做测试。注意group_id的提取方式要根据你的实际命名规则改如果文件名里没有批次信息就得额外维护一个映射表。划分完之后一定要检查每个集合里 9 个类别的分布如果某个类别在验证集里只有个位数样本那这个划分就是失败的需要重新调随机种子或者手动干预。2.2 class 文件与 data.yaml 的对应关系YOLO 训练时最容易被忽略的就是 class 文件和 data.yaml 的对应。9 种岩石的类别名必须和标注文件里的 class_id 严格一致顺序错一个整个训练结果就是玄学。常见做法是单独维护一个classes.txt每行一个类别名行号就是 class_id。然后 data.yaml 里的names列表直接从这个文件读避免手写出错。# 生成 data.yaml classes [granite, basalt, sandstone, limestone, shale, gneiss, marble, quartzite, slate] with open(data.yaml, w) as f: f.write(fpath: ./dataset\n) f.write(ftrain: images/train\n) f.write(fval: images/val\n) f.write(ftest: images/test\n) f.write(fnc: {len(classes)}\n) f.write(fnames: {classes}\n)这里nc必须等于 9names的顺序就是 class_id 0 到 8 的顺序。如果你后面用可视化脚本画框框的颜色和标签文字都依赖这个顺序。我见过有人把names写成字典格式{0: granite, ...}YOLOv5 能读但 YOLOv8 会报错所以统一用列表最稳。另外path建议用相对路径方便整个数据集打包迁移。2.3 标注格式转换从 LabelImg 到 YOLO txt岩石检测的标注通常用 LabelImg 或 CVAT 做导出的是 PASCAL VOC 的 XML 或者 COCO JSON。YOLO 需要的是每张图一个 txt每行class_id x_center y_center width height全部归一化到 0-1。转换脚本网上一搜一大把但岩石图像有个坑很多标注工具默认坐标是左上角原点而 YOLO 要求的是中心点坐标转换时如果忘了减半宽高框会整体偏移。import xml.etree.ElementTree as ET import os def voc_to_yolo(xml_path, img_w, img_h, classes): tree ET.parse(xml_path) root tree.getroot() lines [] for obj in root.findall(object): cls_name obj.find(name).text if cls_name not in classes: continue cls_id classes.index(cls_name) bbox obj.find(bndbox) xmin float(bbox.find(xmin).text) ymin float(bbox.find(ymin).text) xmax float(bbox.find(xmax).text) ymax float(bbox.find(ymax).text) x_center (xmin xmax) / 2.0 / img_w y_center (ymin ymax) / 2.0 / img_h width (xmax - xmin) / img_w height (ymax - ymin) / img_h lines.append(f{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}) return linesimg_w和img_h必须从对应图片读取不能硬编码。岩石图像分辨率差异大有的 4000x3000有的 640x480硬编码会导致小图的框完全错位。转换完之后建议随机抽 20 张用可视化脚本画出来看一眼确认框的位置和标签都对。3. 数据可视化脚本怎么在训练前发现标注翻车3.1 用 OpenCV 画框并检查类别分布可视化脚本的核心作用不是好看而是让你在花几小时训练之前用几秒钟发现「框偏了」「类别标错了」「某类样本全是废图」这些问题。我一般会写一个脚本同时输出三样东西带框的图像、每个类别的样本数统计、以及框的宽高分布直方图。import cv2 import os import matplotlib.pyplot as plt from collections import Counter def visualize_yolo(img_dir, label_dir, classes, output_dir, sample_n20): os.makedirs(output_dir, exist_okTrue) class_counter Counter() all_widths, all_heights [], [] img_files [f for f in os.listdir(img_dir) if f.endswith((.jpg, .png))] for img_name in img_files[:sample_n]: img_path os.path.join(img_dir, img_name) label_path os.path.join(label_dir, img_name.rsplit(., 1)[0] .txt) img cv2.imread(img_path) h, w img.shape[:2] if not os.path.exists(label_path): print(f缺失标注: {img_name}) continue with open(label_path) as f: for line in f: parts line.strip().split() cls_id int(parts[0]) xc, yc, bw, bh map(float, parts[1:]) class_counter[classes[cls_id]] 1 all_widths.append(bw) all_heights.append(bh) x1 int((xc - bw/2) * w) y1 int((yc - bh/2) * h) x2 int((xc bw/2) * w) y2 int((yc bh/2) * h) cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) cv2.putText(img, classes[cls_id], (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) cv2.imwrite(os.path.join(output_dir, img_name), img) print(类别分布:, dict(class_counter)) return class_counter, all_widths, all_heightssample_n控制抽样的数量我一般设 20 到 30太多看不过来。class_counter能直接告诉你哪类样本少得可怜比如 slate 只有 15 个框那训练时这一类基本学不动。all_widths和all_heights可以画直方图如果发现大量框的宽高接近 1.0说明标注时把整张图框进去了这种脏数据必须删。3.2 用直方图判断锚框是否需要重新聚类YOLO 的默认锚框是基于 COCO 数据集聚类的岩石检测的框长宽比和 COCO 差别很大——岩石通常是块状宽高比接近 1:1而 COCO 里有很多细长物体。如果你直接拿默认锚框训练召回率会偏低。可视化脚本输出的宽高分布就是给你做 K-means 聚类用的。import numpy as np from sklearn.cluster import KMeans def cluster_anchors(widths, heights, n_clusters9): data np.array(list(zip(widths, heights))) kmeans KMeans(n_clustersn_clusters, random_state42).fit(data) anchors kmeans.cluster_centers_ # 按面积排序 areas anchors[:, 0] * anchors[:, 1] anchors anchors[np.argsort(areas)] print(推荐锚框 (width, height):) for a in anchors: print(f {a[0]:.4f} {a[1]:.4f}) return anchorsn_clusters设 9 是因为 YOLO 默认就是 9 个锚框分 3 个尺度各 3 个。聚类完之后把结果填到模型的 anchor 配置里通常能提升 2 到 5 个点的 mAP。注意这里用的是归一化后的宽高如果你用的是 YOLOv8它已经改成 anchor-free 了这一步可以跳过但宽高分布仍然值得看能帮你判断数据里有没有异常框。3.3 可视化脚本的批量运行与结果归档实际项目里不可能一张张手动跑我一般会把可视化脚本包成一个命令行工具支持按类别抽样、按划分集合抽样输出结果按vis/train/、vis/val/分目录存。这样每次数据集更新后重新跑一遍对比新旧可视化结果就能快速定位是哪个环节引入了脏数据。python visualize.py \ --img_dir dataset/images/train \ --label_dir dataset/labels/train \ --classes classes.txt \ --output_dir vis/train \ --sample_n 30参数--sample_n建议训练集抽 30、验证集抽 15因为验证集通常小。输出目录里除了画框图我还会存一个stats.json记录类别分布和宽高统计方便后面用脚本对比不同版本的数据集。这个习惯帮我省过好几次返工——有一次发现新加的 200 张玄武岩图片里有 40 张的标注框全部偏移了 10 个像素就是靠对比 stats.json 里的宽高分布异常发现的。4. 避坑与排查岩石检测数据集最常见的 5 个翻车现场4.1 现象训练 loss 正常下降但 mAP 一直为 0原因通常是 class 文件和 data.yaml 的类别顺序不一致。比如 classes.txt 里 granite 是第 0 类但 data.yaml 的 names 列表里 granite 写在了第 3 位模型学到的「0」对应的是 slate验证时自然全错。解决方法是写一个校验脚本读 classes.txt 和 data.yaml逐行对比顺序是否完全一致不一致就报错退出。4.2 现象可视化时框的位置整体偏移这是坐标转换时忘了归一化或者用错了图像尺寸。LabelImg 导出的 XML 里坐标是绝对像素值转换时如果用了错误的 img_w/img_h比如用了统一尺寸而不是每张图的实际尺寸框就会偏。解决方法是转换脚本里强制从图片文件读取实际宽高并且转换后随机抽 10 张画出来肉眼确认。4.3 现象某一类岩石的检测效果特别差先别急着改模型去可视化脚本的类别统计里看这一类有多少个标注框。如果少于 100 个大概率是样本不足。岩石检测里 shale 和 slate 容易混淆如果这两类的样本数差距大模型会偏向多的那类。解决办法是针对性补充少样本类的图像或者用 copy-paste 增强把少样本类扩到 200 个框以上。4.4 现象验证集 loss 震荡严重检查划分时有没有把同一张图的不同增强版本分到训练集和验证集。岩石图像经常做旋转、翻转增强如果增强后的图和原图分到了不同集合验证集就泄漏了。解决方法是增强只在训练时在线做磁盘上只存原图划分也按原图划分。4.5 现象训练到一半 BN 层崩溃loss 变 NaN这是 YOLO 训练中常见的数值不稳定问题岩石图像如果存在大量纯色背景比如实验室拍的岩芯batch 内方差过小会导致 BN 层出问题。解决办法是调小学习率到 0.001 以下或者把 batch size 调到 16 以上增加样本多样性。如果还不行在 data.yaml 里加rect: True做矩形训练减少 padding 带来的无效像素。5. 从 9 类岩石数据集到可复现训练我的参数习惯与验证套路数据集准备好之后训练本身反而简单了。我一般用 YOLOv8n 或 YOLOv8s 起步因为岩石检测的纹理特征比较明显不需要太大的模型。输入尺寸设 640batch size 根据显存来V100 上跑 32 没问题消费级卡就降到 16 或 8。学习率用默认的 0.01 配合余弦退火但如果发现 loss 震荡就降到 0.001。训练轮数先跑 100 轮看曲线如果验证集 mAP 在第 60 轮之后还在涨就加到 200 轮。验证阶段我有个习惯除了看 mAP50 和 mAP50-95一定会把验证集的预测结果用可视化脚本画出来和真实标注并排对比。重点看两类错误——把花岗岩预测成片麻岩类间混淆以及框只框住了岩石的一部分定位不准。前者说明需要补充混淆类别的样本后者说明锚框或输入尺寸需要调整。这个对比图我一般会存成val_compare/目录每次调参后重新生成肉眼扫一遍比看数字快得多。最后说一个我踩过的坑不要用测试集调参。测试集只在最终确定模型后跑一次用来报告最终指标。如果你用测试集反复调学习率、调数据增强那测试集就变成了验证集报出来的指标没有意义。我一般会把测试集单独锁在一个目录里训练脚本里根本不引用它直到最后评估才打开。这个习惯看起来简单但能帮你避免很多自欺欺人的「涨点」。希望帮到你。本文还有配套的精品资源点击获取