简介一套基于 PyTorch 实现的岩石图像分类深度学习入门资源面向计算机视觉初学者或地质相关专业学生帮助读者完成从数据集制作、模型训练到 GUI 界面展示的完整流程。资源共 398 个文件其中 392 张 JPG 图片按类别存放于数据集文件夹3 个 Python 脚本分别承担数据预处理与标注、模型训练、PyQt 交互界面三大环节另有 3 个 txt 文件用于记录文件清单与环境依赖压缩包整体约 27MB。代码对数据集进行了针对性增强在短边补灰边统一为正方形并结合旋转角度扩充样本训练后模型可保存在本地。已有 264 人学习。这套资源能帮你直接理解图像分类项目的工程结构适合作为课程设计或入门实践参考环境安装细节也可按描述中给出的指引完成。1. 岩石识别项目拆解一个能直接跑通的小型图像分类闭环做地质相关的图像处理时最烦的不是模型选型而是手里一堆玄武岩、花岗岩照片靠肉眼分起来又慢又容易翻车。这个「python深度学习识别岩石」项目就是一个完整的小型图像分类闭环数据预处理、数据集文本生成、PyTorch模型训练、PyQt5界面推理都齐了还自带一份岩石图像数据集。适合刚入门深度学习、想从 MNIST 那种玩具数据集跳到真实图像分类的读者也适合需要快速搭一个岩石/矿物分类演示系统的从业者。它不追求 SOTA但能让你在今天内看到训练曲线掉下来、模型保存成功、界面弹出识别结果这对建立完整认知比堆理论有用得多。2. 环境与项目结构从 requirement.txt 到三条主流程2.1 先看 requirement.txtPyTorch 环境到底要装什么拿到压缩包后第一件事不是跑代码而是看根目录下的 requirement.txt。这个文本文件里列出的是运行三个脚本所需的 Python 库包括 torch、torchvision、opencv-python、numpy、Pillow、PyQt5 这些。它只负责告诉你「装什么」不负责帮你装好环境需要自行配置。我一般建议用 conda 先建一个干净环境避免把系统 Python 搞乱。常见做法是conda create -n rock python3.9 conda activate rock pip install -r requirement.txt如果是 NVIDIA 显卡用户装完 CPU 版的 torch 后训练也能跑但速度会让你怀疑人生。正确的顺序是先确定自己的 CUDA 版本再安装对应版本的 PyTorch。例如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121提示装完之后务必在 Python 里执行import torch; print(torch.cuda.is_available())返回 True 才说明 GPU 真的可用。这一步能省掉后面所有「训练莫名慢」的排查时间。很多新手在这一步翻车原因是装了 CPU 版 torch 却以为自己在用 GPU。2.2 三个 Python 脚本的分工与执行顺序环境装好后整个项目的运转逻辑非常清晰依次运行三个脚本就能走完数据准备、模型训练、界面推理的完整流程脚本名职责产出01数据集文本生成制作.py读取 data 目录下的类别图片做预处理和数据增强生成训练/验证用的 TXT 文本train.txt / val.txt以及增强后的图片02深度学习模型训练.py读取 TXT 文本中记录的路径和标签训练模型并保存权重本地模型文件 .pth03pyqt_ui界面.py加载训练好的模型弹出图形界面选择图片输出识别结果桌面端识别工具这三个脚本是串行关系01 不跑02 没有输入02 不跑03 没有模型可加载。刚开始接触这个项目的读者最容易犯的错就是跳步——直接运行 03pyqt_ui 界面弹窗报「模型文件不存在」。这不是代码 bug是执行顺序没遵守。2.3 数据集目录与标签一张图对应一个类别数据集文件夹存放了本次识别的各个类别图片目录结构大致是这样的data/ ├── Basalt/ │ ├── Basalt rock1.jpg │ ├── Basalt rock4.jpg │ └── ... └── Granite/ ├── granite rock4.jpg ├── granite rock41_rotated45.jpg ├── granite rock41_flip.jpg └── ...从文件名能看出很多信息granite rock41_rotated45.jpg 是旋转 45 度后的增强图granite rock41_flip.jpg 是水平翻转后的增强图。也就是说数据集本身已经包含了增强样本01 脚本还会在此基础上再做一轮自动化处理。每个子文件夹名就是类别名脚本会按文件夹遍历图片把「图片路径 类别序号」写进 TXT。理解这个结构后你想加自己的岩石类别直接往 data 下新建文件夹丢图片即可。注意目录名和文件名尽量不要带中文和空格。Windows 下带空格的路径虽然能读但遇到某些 OpenCV 版本会偶发路径解析问题属于典型的「你排查半天以为是代码问题其实是命名问题」。3. 数据预处理与增强短边加灰边转正方形再旋转扩增3.1 为什么非要把图片变成正方形这个项目的数据预处理有一个关键步骤通过在较短边增加灰边使得图片变为正方形如果图片原本就是正方形则不会增加灰边。很多读者会问直接 resize 成正方形不就行了为什么非要用灰边补直接 resize 会把原本的长宽比例破坏掉。岩石纹理在拉伸后颗粒结构、层理特征都会变形模型学到的是「被压扁的岩石」而不是「岩石」。而补灰边相当于把图片放进一个正方形画布原图内容比例完全不变只是周围多了一圈灰色区域。灰色一般是 RGB 114,114,114 附近在归一化后接近中性值不会给模型引入额外偏差。这是目标检测里常见 letterbox 预处理思路在分类任务上的应用。另一个好处是适配预训练模型。torchvision 自带的 ResNet、VGG 等模型预训练权重是基于 ImageNet 的 224×224 输入训练的喂给它的图片最好先规范化到接近这个尺寸。如果图片本身不是正方形resize 到 224×224 时会畸变先补成正方形再 resize则只做等比缩放信息损失小得多。3.2 01 脚本的核心实现灰边填充与旋转增强01脚本里核心的补边逻辑用 OpenCV 写出来大致是这样import cv2 def pad_to_square(img): h, w img.shape[:2] if h w: return img side max(h, w) top (side - h) // 2 bottom side - h - top left (side - w) // 2 right side - w - left # 灰色填充等价于 YOLO 系列常用的 borderValue return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value(114, 114, 114)) img cv2.imread(data/Basalt/Basalt rock1.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # OpenCV 默认 BGR img_square pad_to_square(img)这段代码的逻辑是先取图像的高和宽如果相等就直接返回否则取最长边作为正方形边长计算上下左右各需要补多少像素然后使用copyMakeBorder填充。//整除是为了让上下、左右补的像素数尽量对称避免图片内容偏向一侧。value(114,114,114)是灰边颜色值之所以选灰色而不是黑色或白色是因为黑色会拉低归一化后的均值白色会拉高灰色最中性。旋转增强的代码也很短def rotate_45(img): h, w img.shape[:2] M cv2.getRotationMatrix2D((w // 2, h // 2), 45, 1.0) rotated cv2.warpAffine(img, M, (w, h), borderModecv2.BORDER_CONSTANT, borderValue(114, 114, 114)) return rotated这里getRotationMatrix2D的第二个参数是旋转角度正值代表逆时针第三个参数是缩放比例1.0 代表不缩放。旋转中心取图像中心。warpAffine输出尺寸保持和输入一致旋转后四个顶角会空出来同样用灰边填充。3.3 训练集 / 验证集文本怎么生成01脚本最终的目标是生成文本文件记录图片路径和标签。运行后会在根目录得到类似下面的格式data/Basalt/Basalt rock1.jpg 0 data/Basalt/Basalt rock4.jpg 0 data/Granite/granite rock4.jpg 1 data/Granite/granite rock41_rotated45.jpg 1每行前半部分是图片相对路径后半部分是类别序号。类别序号由脚本遍历文件夹时按顺序生成0 对应 Basalt1 对应 Granite。生成这段文本的逻辑是import os data_dir data label_map {} lines [] for idx, cls_name in enumerate(os.listdir(data_dir)): cls_path os.path.join(data_dir, cls_name) if not os.path.isdir(cls_path): continue label_map[cls_name] idx for fname in sorted(os.listdir(cls_path)): if fname.lower().endswith((.jpg, .jpeg, .png)): lines.append(f{os.path.join(cls_path, fname)} {idx}) with open(train_val.txt, w, encodingutf-8) as f: f.write(\n.join(lines))这段代码做的事情是遍历 data 目录下的所有子目录每个子目录名对应一个类别并分配序号再遍历该目录下所有图片文件把「相对路径 空格 序号」写入文本。enumerate保证类别序号从 0 开始连续编号sorted保证文件顺序稳定避免每次运行生成的文本顺序不同导致验证集分布漂移。验证集的划分在 02 训练脚本里通常会按比例随机切分常见做法是train_val.txt全部读进来然后用random.shuffle后按 8:2 切分。如果你想要更严谨的划分也可以在 01 脚本里直接输出 train.txt 和 val.txt 两个文件这样训练和验证的数据边界更清晰。提示03pyqt_ui 界面推理时模型输出的类别序号必须和这里的 label_map 顺序一致。你在这个脚本里把 Basalt 放在第 0 位、Granite 放在第 1 位训练和推理就都用这个顺序别在推理时重新排序否则识别结果会驴唇不对马嘴。4. 模型训练从 TXT 路径到保存本地模型4.1 自定义 Dataset按文本路径读图02 训练脚本的核心是读取 01 生成的 TXT 文件按路径加载图片并训练模型。PyTorch 里自定义数据读取需要继承torch.utils.data.Dataset这是整个训练链路里最容易被写错的一环。常见写法如下from torch.utils.data import Dataset from PIL import Image class RockDataset(Dataset): def __init__(self, txt_path, transformNone): self.samples [] self.transform transform with open(txt_path, r, encodingutf-8) as f: for line in f: line line.strip() if not line: continue path, label line.split() self.samples.append((path, int(label))) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) # 统一转 RGB防止灰度图出问题 if self.transform: img self.transform(img) return img, label这里__getitem__返回的是「图像张量 标签」的元组。Image.open(path).convert(RGB)很关键如果数据集中混有灰度图或 RGBA 图不统一通道数训练时会直接报张量维度不匹配的错误。TXT 每行按空格拆分成path和labelint(label)确保标签是数值类型不然 CrossEntropyLoss 会报类型错误。有了 Dataset再包一层 DataLoader 就能batch迭代from torch.utils.data import DataLoader dataset RockDataset(train_val.txt, transformtrain_transform) train_size int(0.8 * len(dataset)) val_size len(dataset) - train_size train_ds, val_ds torch.utils.data.random_split(dataset, [train_size, val_size]) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers2) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers2)random_split按固定比例切训练集和验证集shuffleTrue保证每个 epoch 数据顺序都不同防止模型记住样本顺序。num_workers是数据加载的并行进程数Windows 上设成 2 即可设太高容易踩到多进程 bug。4.2 训练超参轮次、学习率、优化器怎么定岩石图像分类属于典型的细粒度图像识别类别间差异主要体现在纹理和颜色上不是结构差异。因此不需要从零训练一个模型直接用 ImageNet 预训练权重做迁移学习是最划算的做法。常见方案是加载 ResNet18 或 ResNet50替换最后一层全连接import torchvision.models as models import torch.nn as nn num_classes 2 # Basalt 和 Granite model models.resnet18(pretrainedTrue) model.fc nn.Linear(model.fc.in_features, num_classes)model.fc是 ResNet 的最后分类层原版输出 1000 类这里改成自己的类别数。pretrainedTrue会下载 ImageNet 权重第一次运行需要联网。训练超参的常见设置是import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size5, gamma0.5)学习率 1e-4 比默认的 1e-3 更稳。迁移学习场景下预训练权重已经是很优的解学习率太大容易把前面层的特征冲掉。StepLR每 5 个 epoch 把学习率乘 0.5让训练后期在更小的步长下微调。训练轮次一般设 20 到 50 轮岩石数据集小的花30 轮以内就能收敛到不错的效果。训练循环的标准写法for epoch in range(30): model.train() for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() model.eval() total, correct 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total print(fEpoch {epoch1}/30, Loss: {loss.item():.4f}, Val Acc: {val_acc:.4f}) scheduler.step()model.train()和model.eval()必须成对出现。train()模式下 BatchNorm 和 Dropout 是生效的eval()模式下这些层会切换到推理行为。很多新手忘记切回train()导致每个 epoch 的验证精度和训练精度对不上然后开始怀疑数据有问题。4.3 模型保存与加载的边界训练结束后脚本会把模型权重保存到本地。保存方式不同后续加载方式也不同# 只保存权重推荐 torch.save(model.state_dict(), rock_model.pth)state_dict只保存模型参数不保存模型结构。之后加载时必须先定义好同样的模型结构再load_state_dictmodel models.resnet18(pretrainedFalse) model.fc nn.Linear(model.fc.in_features, num_classes) model.load_state_dict(torch.load(rock_model.pth, map_locationcpu)) model.eval()map_locationcpu这一步在 PyQt5 推理时尤其重要。如果训练是在 GPU 上跑的保存的权重会带有 CUDA 张量标记而界面推理的机器不一定有显卡不加map_location就会报 RuntimeError。03 脚本里加载模型的代码如果是在 CPU 环境跑务必保留这个参数。5. 训练与推理避坑岩石识别项目里最容易翻车的五个点5.1 训练极慢但代码不报错现象训练能正常启动loss 也在下降但一个 epoch 要跑十几分钟GPU 显存占用却是 0。原因安装的是 CPU 版 PyTorch或者装了 GPU 版但 CUDA 不可用to(device)把模型和数据的 device 都设成了 GPU但 PyTorch 底层没有 CUDA 支持。解决按 2.1 节的方法先torch.cuda.is_available()验证确认安装的是与 CUDA 版本匹配的 PyTorch。装完之后用nvidia-smi看显卡驱动版本再对照 PyTorch 官网的 CUDA 对应表选安装命令。5.2 图片路径带空格导致读取失败现象TXT 已经生成但训练时报 FileNotFoundError而且报错的文件名看起来是对的。原因我见过实际数据集里文件名是「granite rock4.jpg」这种带空格的TXT 里按空格拆分路径和标签时路径被拆成了两段。解决生成 TXT 时用制表符\t分隔路径和标签读取时用split(\t)或者干脆在预处理阶段把文件名里的空格替换成下划线。别小看这个问题Windows 下的 OpenCV 和 PIL 对带空格路径的兼容性完全靠运气。5.3 旋转增强后图片内容被裁掉现象训练集里旋转 45 度的图片看起来四个角全是灰边岩石主体变小了验证精度上不去。原因warpAffine输出尺寸固定为原图宽高旋转 45 度后原图的四个顶点会被切掉内容信息丢失。这不是本项目代码特有的问题而是所有用 OpenCV 做旋转增强的人都会踩的坑。解决旋转前先计算外接矩形尺寸把输出尺寸扩到side * sqrt(2)左右或者不用旋转 45 度改用 15 度、30 度这样的小角度信息损失小得多。如果你追求稳妥宁愿用随机裁剪加翻转也不要依赖大角度旋转。5.4 PyQt5 界面推理结果总是同一类现象03脚本能跑通但不管选哪张图界面都输出同一个类别而且置信度接近 1。原因加载模型后没有调用model.eval()模型处于训练模式。训练模式下 Dropout 层会随机丢弃部分神经元BatchNorm 会使用 batch 统计量而不是全局统计量导致输出不稳定且偏向某个类别。解决加载权重后立刻加一行model.eval()。另外检查预处理是否和训练一致——推理时也要先补灰边转正方形再 resize 到 224×224不能直接拿原图喂给模型。5.5 验证集精度高但实测效果差现象训练时验证集精度 95% 以上但用手机拍一张岩石照片识别就乱套。原因数据集图片和真实场景存在分布差异。训练集里的岩石照片偏暗、背景单一而真实照片可能有手指、硬币、桌面反光。解决在 01 脚本里增加背景扰动增强把岩石随机贴在不同背景上或者推理前对输入图片做一次自适应直方图均衡降低光照差异。这是一种常被忽略的「训练和推理分布对齐」问题属于典型的模型玄学但根因其实是数据太干净。提示以上五条里第 5.1 和第 5.4 可以提前用一句检查代码规避掉第 5.2 靠改变 TXT 分隔符解决第 5.3 靠改旋转输出尺寸解决第 5.5 靠数据增强多样性解决。所以训练脚本开头可以统一加一段自检逻辑把 device 可用性和模型状态打印出来。6. 把模型装进 PyQt5 界面推理脚本与预处理一致性03pyqt_ui 界面脚本做的事不复杂加载模型权重、弹出文件选择框、对选中的图片做和训练一致的预处理、前向推理输出类别和置信度。但它是最容易暴露前面所有坑的环节因为训练时你面对的是处理干净的 TXT 数据而界面里面对的是用户随手挑的一张照片。界面推理的核心代码结构如下import cv2 import torch from PyQt5.QtWidgets import QFileDialog device torch.device(cpu) model.eval() def pad_to_square(img): h, w img.shape[:2] side max(h, w) top (side - h) // 2 bottom side - h - top left (side - w) // 2 right side - w - left return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value(114, 114, 114)) def preprocess(path): img cv2.imread(path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img pad_to_square(img) img cv2.resize(img, (224, 224)) img img.astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) img (img - mean) / std tensor torch.from_numpy(img.transpose(2, 0, 1)).unsqueeze(0) return tensor path, _ QFileDialog.getOpenFileName(None, 选择岩石图片, , Images (*.jpg *.png *.jpeg)) tensor preprocess(path) with torch.no_grad(): outputs model(tensor) prob torch.softmax(outputs, dim1) cls_id torch.argmax(prob, dim1).item() confidence prob[0, cls_id].item() print(f类别序号: {cls_id}, 置信度: {confidence:.4f})这段代码里有个细节我反复强调过预处理必须和训练时完全一致。训练时的RandomHorizontalFlip、RandomRotation是随机增强推理时不能再用但补灰边、resize、归一化的均值和标准差必须逐字节一致。很多从其他项目抄推理代码的人经常在这里翻车——训练用的归一化是(0.485, 0.456, 0.406)推理时抄成了(0.5, 0.5, 0.5)然后精度断崖式下跌还以为是模型没训练好。推理输出的是类别序号要和 label_map 对应起来展示成中文名称。界面上建议把置信度也显示出来这是判断模型是否「硬猜」的关键指标。如果一张玄武岩图片跑出置信度只有 0.55说明模型其实没把握如果置信度稳定在 0.9 以上才说明它真的学到了岩石纹理特征。这套流程我从拆包到跑通只花了不到一小时最深刻的教训是这个项目的三个脚本之间的依赖关系远比想象中脆弱01 脚本生成的 TXT 格式、02 脚本的模型分类层定义、03 脚本的类别映射任何一环不一致最后界面输出的都是错的结果。从那以后我每次拿到这种完整项目都强制走一遍「生成数据 → 训练 → 推理」全链路并且第一步先确认设备可用性和文件路径规范再去看模型结构。希望帮到你。本文还有配套的精品资源点击获取