Python+CNN火焰识别：从数据集到部署的图像分类实战

简介基于Python与CNN的火焰识别项目面向计算机视觉初学者及消防安全相关场景开发者提供从数据处理、模型训练到界面展示的完整工程。压缩包共302个文件体积11.72MB其中包含288张jpg图片和8张png图片构成带标签的火焰/非火焰数据集另有3个py脚本实现数据集文本生成、深度学习训练与PyQt交互界面3个txt文件则包含依赖环境说明等配置信息。资源对数据集进行了预处理包括在短边增加灰色边框统一为正方形、随机旋转增强样本多样性有助于提升模型泛化能力。配套代码基于PyTorch实现标签与训练集生成逻辑清晰适合想快速上手图像分类项目的学习者参考同时内置PyQt显示界面便于直观查看分类结果。目前已有75人浏览学习可作为火焰检测入门实践的参照模板。1. pythonCNN做火焰识别这个资源包里其实是一条完整图像分类流水线做安全监控的同事拿了一批现场图片问我能不能用 python CNN 深度学习直接判断画面里有没有火焰而不是等烟雾报警器触发后再人工盯屏。这个压缩包正好解决这个问题里面既有按类别整理好的火焰/非火焰数据集又有三个按顺序跑的脚本跑通之后就是一个从数据准备、模型训练到 PyQt 界面识别的完整图像分类项目。适合刚把 PyTorch 装好、想拿真实小数据集练手的同学也适合需要快速验证 CNN 火焰识别效果的现场工程师。下面我按实际拆项目的方式把数据标注、预处理、训练参数到界面部署逐个过一遍。2. CNN火焰识别项目拆解数据集、预处理和训练准备2.1 数据集结构类别文件夹、图片路径和标签生成的约定打开压缩包后先别急着跑代码。我拆这个项目的第一件事是把数据集文件夹的目录树拉出来看一眼。项目原文里写了“数据集文件夹存放了本次识别的各个类别图片”也就是说类别信息是用文件夹名表达的一个子文件夹代表一个类别里面放的是这个类别的 jpg 图片。实际操作时会发现火焰识别的分类任务通常只有两个类别有的项目叫fire和nofire有的叫has_fire和no_fire还有的直接用0和1做文件夹名。这个资源包里具体用哪两个名字以你解压后的实际目录为准因为代码是按“遍历数据集子目录”的方式自动读类别名的目录名本身不影响最终结果。这种按文件夹分类的组织方式好处是标签不用人工维护脚本读完子目录名就能直接映射成数字标签。坏处是如果现场图片混放或者一个文件夹里混进了其他格式的图片后续训练会翻车。所以建议拿到资源包后先做一次简单核验每个子文件夹下的图片数量是否均衡格式是否都是jpg有没有损坏的图片文件。2.2 为什么预处理要先补灰边再做旋转扩增这个项目在预处理阶段做了两个关键动作补灰边和旋转扩增。先看补灰边。CNN 模型输入一般是固定尺寸常见的是224×224、128×128或者96×96。直接把不同宽高比的图片resize成正方形会让火焰区域发生非等比压缩原本细长的火苗可能被压扁特征就变了。项目里采用的做法是在短边两侧补上灰色像素先把图片变成正方形再把正方形统一缩放成模型需要的尺寸。这样图片内容不发生形变模型学习到的火焰形状更接近真实。“如果图片原本是正方形则不会增加灰边”这句话说明代码里做了宽高判断。判断逻辑不复杂假如图片宽是 640高是 480短边是高 480那么上下不需要补左右各补(640-480)/2的灰边补完后变成 640×640然后再缩放到训练尺寸。补灰边的颜色选0或者127都常见选0就是补黑边选127是中灰具体要看代码里的实际值。旋转扩增属于数据增强。火焰图片有个特点火苗方向、形态、拍摄角度差异很大。为了让模型不依赖固定的拍摄角度项目会对训练图片做随机小角度旋转生成更多样化的训练样本。常见的旋转范围是±10°到±15°超过±20°可能会让火焰形状失真反而影响训练效果。旋转后边缘会产生空区域可以用fill参数填入灰边值和补灰边保持一致。这两个预处理放在一起相当于同时解决了输入尺寸统一和样本量不足两个问题这也是这个小项目比直接硬套ImageFolder更稳的原因。2.3 训练集和验证集按类别路径生成 txt 清单项目里没有用 PyTorch 默认的ImageFolder直接读目录而是先让01数据集文本生成制作.py把每个类别下的图片路径和对应标签写到文本文件里。摘要原话是“运行 01 数据集文本生成制作.py 文件会读取数据集下每个类别文件中的图片路径和对应的标签”。这样设计有个实在的好处训练之前可以肉眼检查训练集和验证集的划分发现某个类别的图片全跑到验证集里可以直接改 txt 文件不用重新移动图片。同时也方便后续在模型训练脚本里用自定义Dataset读取路径格式完全是可控的。文本文件通常每行是一条记录格式是“图片相对路径 标签整数”中间用空格或者\t分隔。训练集和验证集可以写在同一个文件里也可以分开写成train.txt和val.txt。这个项目提到“将 txt 文本中记录的训练集和验证集进行读取训练”按常见做法推断大概率是两个文件或者一个文件里按行标记了划分方式。到这里数据和预处理的框架已经清楚了。下一章我按 01 脚本的职责把目录遍历、标签生成、训练验证划分这三件事拆开讲并给出一份可以直接改用的生成代码。3. 生成数据集文本清单01脚本的核心逻辑与改造3.1 脚本流程拆解遍历目录、写路径、打标签01数据集文本生成制作.py这个脚本要做的事情按行拆开一共四步第一指定数据集根目录。项目里默认是数据集这个文件夹。第二遍历根目录下的所有子文件夹每找到一个子文件夹就当成一个类别并分配一个从0开始的整数标签。第三遍历该子文件夹下的所有图片把“相对路径 标签”拼成一行文本。第四把图片按比例随机划分成训练集和验证集分别写入两个 txt 文件。这里有一个容易被新手忽略的点写进 txt 里的图片路径最好是相对路径而不是绝对路径。如果写绝对路径换一台电脑或者把项目从D:\project挪到E:\work所有路径就全部失效训练脚本直接报“找不到文件”。用相对路径的话只要保持项目压缩包解压后的目录结构不变换目录也能跑。3.2 一份可改用的代码实现下面这份代码是我按项目描述重写的最小版本逻辑和 01 脚本保持一致但去掉了和具体目录名耦合的部分方便你直接复用。import os import random # 参数区 data_root 数据集 # 数据集根目录 train_txt train.txt # 训练集清单 val_txt val.txt # 验证集清单 val_ratio 0.2 # 验证集比例 random_seed 42 # 随机种子保证每次划分结果一致 # categories [d for d in os.listdir(data_root) if os.path.isdir(os.path.join(data_root, d))] categories.sort() print(检测到类别, categories) with open(train_txt, w, encodingutf-8) as f_train: with open(val_txt, w, encodingutf-8) as f_val: for label, cat in enumerate(categories): cat_dir os.path.join(data_root, cat) imgs [f for f in os.listdir(cat_dir) if f.lower().endswith((.jpg, .jpeg, .png))] if len(imgs) 0: print(f警告类别 {cat} 下没有图片已跳过) continue random.seed(random_seed) random.shuffle(imgs) val_size int(len(imgs) * val_ratio) for i, img in enumerate(imgs): # 这里写相对路径不写绝对路径 rel_path os.path.join(cat, img) line f{rel_path}\t{label}\n if i val_size: f_val.write(line) else: f_train.write(line) print(生成完成) print(f训练集总行数{sum(1 for _ in open(train_txt, encodingutf-8))}) print(f验证集总行数{sum(1 for _ in open(val_txt, encodingutf-8))})这段代码里有几个参数需要你实际改data_root是对应你自己数据集目录的名字如果你的目录叫fire_dataset就改成fire_dataset。val_ratio是验证集占比我这里给了 0.2也就是 20% 的图片进验证集。如果数据集总量很小比如只有两三百张建议改成 0.15 左右否则验证集会太少评估结果波动大。random_seed这个参数很有用。不设固定种子的话每次运行划分结果都不同训练脚本每次看到的验证集也不一样跑出来的准确率会有不小浮动。固定成42或者其他整数能保证每次复现结果一致这一点在排查“为什么昨天 95%今天只有 88%”的时候特别管用。写入行格式用了\t制表符分隔。如果你用空格分隔图片文件名里一旦出现空格训练脚本按空格切分就会把文件名切成两段。制表符能规避这个问题是我在类似小项目里的习惯。3.3 常见参数调整多类别、绝对路径/相对路径、验证集比例01 脚本最大的价值是把人工整理目录的过程自动化了。实际使用时有几种情况需要微调。如果项目从二分类扩展到多分类比如增加“烟雾”“明火”“无异常”三个文件夹这段代码不需要改动。因为它是按os.listdir动态读取子目录的新增一个文件夹标签会从0,1,2自动排下去只要保证每个文件夹命名规范就行。如果模型训练脚本要求绝对路径可以把rel_path os.path.join(cat, img)改成rel_path os.path.abspath(os.path.join(data_root, cat, img))。但我个人不建议这么做除非你有跨文件夹读取的强需求否则后面换机器全是坑。验证集比例是一个容易被调崩的参数。火焰识别这种二分类任务正负样本比例接近1:1时验证集 20% 足够。但如果某个类别图片特别少比如只有 30 张20% 就是 6 张验证图片准确率会被这几张图片的随机性主导。这时候要么调小val_ratio要么在收集图片时保证每个类别至少 50 张以上。4. 模型训练02脚本的参数、结构和保存策略4.1 CNN 模型结构卷积层、池化层、全连接层怎么定02深度学习模型训练.py负责真正的训练过程。这个项目没有用预训练模型而是自己搭了一个小型 CNN原因是火焰检测场景里火焰通常有比较明显的颜色特征和纹理特征数据量又不算大小网络更容易训练收敛也更好在 CPU 上跑。典型的 CNN 火焰识别结构是三到四个卷积块每个卷积块由卷积 - ReLU - 池化组成。比如输入是3×224×224的 RGB 图片第一层卷积用 32 个3×3卷积核得到32×112×112的特征图第二层用 64 个卷积核第三层用 128 个特征图尺寸逐层减半。最后接全局平均池化或者Flatten再连一到两层全连接输出两个类别得分。下面是一个可以直接替换到项目里的 PyTorch 模型定义和我平时的写法接近import torch.nn as nn class FireCNN(nn.Module): def __init__(self, num_classes2): super(FireCNN, self).__init__() # 三个卷积块每个都包含卷积 激活 池化 self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) # 全局平均池化避免全连接层输入尺寸写死 self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(128, 64), nn.ReLU(inplaceTrue), nn.Dropout(0.5), # 防止过拟合 nn.Linear(64, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x这段代码里最关键的是nn.AdaptiveAvgPool2d((1, 1))。它的作用是把任意尺寸的特征图强制压成1×1这样前面卷积层即使遇到不同的输入尺寸全连接层的参数个数也不会变。但在这个项目里由于预处理已经强制把图片变成正方形所以也可以不用它直接nn.MaxPool2d后计算全连接输入维度。用全局平均池化的好处是代码更省心不用手算踩坑。Dropout(0.5)放在全连接层前目的是让模型在训练时随机丢弃一半神经元减少对某一个特征的过拟合。火焰数据集如果只有几百张图片Dropout 基本是必需品否则验证集准确率会明显低于训练集。4.2 训练超参数batch size、学习率、epoch 怎么设训练火焰识别模型超参数比网络结构更影响最终效果。PyTorch 里数据加载要用DataLoader预处理方式必须和项目描述保持一致先补灰边成正方形再随机旋转然后统一尺寸。from torch.utils.data import Dataset, DataLoader from PIL import Image import torchvision.transforms as T class FireDataset(Dataset): def __init__(self, txt_path, transformNone): self.samples [] with open(txt_path, r, encodingutf-8) as f: for line in f: line line.strip() if not line: continue path, label line.split(\t) self.samples.append((path, int(label))) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label # 训练集和验证集使用同一套基础预处理 base_transform T.Compose([ T.Resize((224, 224)), # 图片之前已经补灰边成正方形 T.RandomRotation(10), # 旋转范围 ±10° T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])这里RandomRotation(10)是唯一一个和验证集不一致的变换。训练集可以做随机旋转验证集不应该做随机旋转否则评估结果不稳定。验证集只用Resize - ToTensor - Normalize就够了。训练超参数我一般这么起步batch_size 16或32。数据集小于一千张时16更稳梯度更新更频繁。learning_rate 0.001。用Adam优化器时这个数值是安全起点。epoch 30到50。小数据集 20 轮基本能看到收敛趋势超过 50 轮还没有明显提升就说明模型过拟合了。训练循环的逻辑是每个 epoch 遍历训练集计算损失反向传播更新参数每个 epoch 结束后在验证集上跑一次记录准确率。保存验证集准确率最高的那一版权重而不是最后一轮的权重。这一点很重要因为最后一轮往往已经在验证集上过拟合了。4.3 训练完成后保存什么模型权重、标签表、指标文件训练完成后的保存策略直接决定后面 PyQt 界面能不能顺利接上。建议至少保存三个东西第一best_model.pth只保存state_dict而不是整个模型对象。这样模型结构变了权重还能被新结构加载。第二classes.json保存类别对应关系。比如{0: fire, 1: nofire}界面加载后才知道标签 0 和 1 分别代表什么。第三train_history.json记录每个 epoch 的训练损失和验证准确率方便后续判断模型是欠拟合还是过拟合。import json import torch # 假设训练过程中保存了 best_model_state torch.save(model.state_dict(), best_model.pth) class_names {i: cat for i, cat in enumerate(categories)} with open(classes.json, w, encodingutf-8) as f: json.dump(class_names, f, ensure_asciiFalse, indent2)state_dict只存权重不存模型结构。加载时需要先用FireCNN(num_classes2)创建模型实例再load_state_dict。这个做法比torch.save(model)更稳因为后者会把类定义路径也存进去换电脑或改目录后经常加载失败。5. 从模型到界面PyQt识别部署与避坑实战5.1 03脚本的界面逻辑加载模型、选择图片、展示结果03pyqt_ui界面.py做的事就是打开一个图形窗口用户点按钮选择一张图片程序调用已经训练好的模型输出“有火焰”还是“无火焰”。PyQt 里最简版本只需要三个控件一个选图按钮、一个显示图片的QLabel、一个显示识别结果的文本标签。界面脚本的推理部分最容易出错的不是 PyQt 本身而是图片预处理。训练时你用了“补灰边 resize 归一化”推理时就必须原样复制这一套少一步都会导致识别精度明显下降。下面是我经常复用的推理函数骨架import torch from PIL import Image def predict_fire(model, img_path, device): from torchvision import transforms as T img Image.open(img_path).convert(RGB) w, h img.size # 1. 补灰边成正方形逻辑和训练预处理保持一致 if w h: result Image.new(RGB, (w, w), (0, 0, 0)) result.paste(img, (0, (w - h) // 2)) else: result Image.new(RGB, (h, h), (0, 0, 0)) result.paste(img, ((h - w) // 2, 0)) img result # 2. 缩放到模型输入尺寸转张量并归一化 transform T.Compose([ T.Resize((224, 224)), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img transform(img).unsqueeze(0).to(device) # 3. 推理并取最大得分 model.eval() with torch.no_grad(): outputs model(img) pred torch.argmax(outputs, dim1).item() return pred这段代码里补灰边用的填充色是(0,0,0)也就是黑边。如果你训练时用的是灰边(127,127,127)这里必须改成(127,127,127)。建议直接去打开 01 或 02 脚本确认补边时用的具体数值再原样抄过来不要凭感觉猜。PyQt 界面部分只需要把predict_fire的返回值映射成中文显示。比如pred0显示“未检测到火焰”pred1显示“检测到火焰”。如果要做得更直观还可以把预测得分softmax后的概率也显示出来这样现场操作的人看到的不只是一个人为设定的阈值还能判断模型是否在犹豫。5.2 避坑指南火焰识别项目的高频踩坑记录这个项目我在类似数据集上跑过几轮下面几条是最容易翻车的位置按“现象 → 原因 → 解决”的方式写出来可以直接对着排查。现象一运行 01 脚本后生成的 txt 内容为空或者只有几行。原因是数据集根目录名写错了或者图片文件后缀大小写不一致。写的是数据集但实际目录名是datasetos.listdir就找不到任何类别。解决方法是先print(os.listdir(data_root))确认目录存在再确认图片后缀是否都在.jpg, .png的过滤范围内。现象二02 训练时报错RuntimeError: CUDA out of memory。原因很直接显卡显存不够火焰图片分辨率高、batch_size又偏大图像预处理后的张量把显存占满了。解决方法是把batch_size从 32 改成 16 或 8同时把训练尺寸从224降到192。如果项目支持可以在训练脚本里加一句torch.backends.cudnn.benchmark False降低显存峰值。现象三训练时准确率一直停在 50% 附近怎么调都上不去。原因是标签和图片可能对不上。比如文件夹名排序后某两个类别的图片混在一起或者 01 脚本读取时os.listdir顺序随机导致标签映射不稳定。解决方法是打印categories列表确认排序结果再打开train.txt前几行手动用 PIL 打开图片核对图片内容和标签是否一致。现象四界面推理结果很差训练时准确率很高。这是最典型的血泪经验推理预处理和训练预处理不一致。训练时做了随机旋转推理时也做了训练时补灰边推理时直接拉伸。这些不一致会让模型看到的图片分布和训练时完全不同。解决方法是把推理的预处理代码和训练的预处理代码并排对比逐行保证一致包括填充值、尺寸、归一化均值和标准差。现象五换一台电脑后界面能打开但选图片后报错找不到模型权重。原因是模型路径写死了绝对路径比如C:/Users/xxx/Desktop/project/best_model.pth换目录后失效。解决方法是把模型路径改成相对路径或者通过os.path.join(os.path.dirname(__file__), best_model.pth)动态拼接。这一点在压缩包解压到不同位置时尤其重要。这些坑我基本都踩过一遍。其中预处理不一致那个最隐蔽因为训练时不会报警界面也能正常跑只有看到识别结果才会发现模型好像变成了“瞎子”。6. 验证火焰识别效果单图预测、混淆矩阵和二次调优技巧模型训练完第一步不是直接上界面而是先用命令行脚本验证一张真实图片确认权重加载没问题。我习惯在项目根目录放一个predict_demo.py只做一件事接收图片路径打印类别概率。import torch from PIL import Image model FireCNN(num_classes2) model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) pred predict_fire(model, test_fire.jpg, cpu) print(预测类别, pred)这个脚本跑通后再拿几张没有参与训练的现场图测试重点看两类误报一是把红色塑料袋、红色灯光误识别成火焰二是把小火苗漏检。火焰识别的难点不在训练而在误报边界因为黄色、橙色、红色在自然场景里太常见了。进阶一点的做法是统计混淆矩阵看看模型到底在哪个类别上出错。对二分类任务矩阵只有四格真火判为火、真火判为无火、无火判为火、无火判为无火。如果第二格多说明模型对暗火和远距离小火不敏感解决方法是增加这些难样本的图片或者把训练尺寸从 224 提升到 256如果第三格多说明模型学到的可能是颜色而不是纹理解决方法是减少训练集里的暖色调干扰图或者加强旋转和亮度扰动。如果想在现有结构上快速提升精度还有一个技巧把模型第一层卷积的输出提取出来做个可视化。火焰图片经过第一个卷积层后边缘和颜色斑块的特征会被激活如果可视化后完全看不出轮廓说明预处理或学习率有问题如果轮廓清晰再决定是否要加深网络。这个方法比盲目调参快得多也不用换成 ResNet 那种大模型。从那以后我每次做这种图像识别小项目都会强制走一遍“单图预测 → 多看难样本 → 可视化中间特征”的验证流程确认模型学会了真的火焰特征而不是撞对了训练集的几张图。希望帮到你。本文还有配套的精品资源点击获取