简介:西瓜病识别图像分类数据集是一套面向图像分类入门与农业病害识别应用的高质量标注资源,适合深度学习初学者与研究者快速搭建分类模型。资源共2000个文件,含1998张JPG图片,覆盖花叶病毒、健康、炭疽病和霜霉病四个类别,已划分训练集与测试集并按类别分目录存放;另附JSON标注文件和Python可视化脚本,可快速查看图像与标签。压缩包约156.94MB,结构清晰,下载后即可接入PyTorch/TensorFlow训练。目前已有59人浏览学习。作者还提供CNN分类网络改进的配套专栏,结合资源中的训练集与可视化脚本,读者可完成从数据处理、模型训练到精度优化的完整流程,适合课程设计、科研预实验或项目练手。
1. 拿到约5,700张已标注的西瓜病识别数据集,先别急着开训
做图像分类的同行应该都有体会:找一个“已标注、能直接用”的垂直领域数据集有多难。西瓜病识别这个方向尤为典型——公开数据集里大多是叶片病害,果实表面的病斑识别要么样本少,要么标注只给到病害大类。手上这份正好反过来:约5,700张数据,已经完成标注,类目落在西瓜果实常见病害和正常果上,拿来跑图像分类模型基本不需要再做标注层面的返工。这篇文章想和你聊清楚三件事:这5,700张适合做什么量级的任务、从标签解析到训出第一个模型的完整路径、以及真正会让结果翻车的边界条件。适合谁看?做课程设计、搞论文实验、或者想练手垂直领域分类落地的工程师,这个规模不会大到让你等训练等到下班,也不会小到训完毫无说服力。
2. 先读标注再谈训练:目录结构与标签解析
2.1 数据集最常见的两种组织方式
这类已标注图像分类数据集,拿到手先看目录,不要急着写训练脚本。常见做法有两种:第一种是images目录存放全部图片,另配一个JSON或CSV文件记录文件名和类别;第二种是直接按类别建子目录,每个文件夹名就是标签。两种组织方式决定了后面的数据读取代码完全不一样。约5,700张的规模不算大,但如果你一开始就假设错了标签格式,后面的清洗步骤全要返工。
判断方式很简单:打开根目录看有没有标注文件。有JSON或者CSV就走映射方案,纯文件夹结构就走目录读取方案。下面这段代码用Python同时兼容两种方案,省得你反复改:
import os import json import csv from collections import Counter dataset_root = "./watermelon_disease" # 方案A:有标注文件(JSON/CSV) if os.path.exists(os.path.join(dataset_root, "annotations.json")): with open(os.path.join(dataset_root, "annotations.json"), "r", encoding="utf-8") as f: ann = json.load(f) # 常见格式之一:{"images": [{"file_name": "001.jpg", "label": "anthracnose"}], "categories": [...]} items = [(d["file_name"], d["label"]) for d in ann["images"]] elif os.path.exists(os.path.join(dataset_root, "labels.csv")): with open(os.path.join(dataset_root, "labels.csv"), "r", encoding="utf-8") as f: reader = csv.DictReader(f) items = [(row["file_name"], row["label"]) for row in reader] # 方案B:按类别子目录组织 else: items = [] for class_name in os.listdir(dataset_root): class_dir = os.path.join(dataset_root, class_name) if not os.path.isdir(class_dir): continue for fname in os.listdir(class_dir): if fname.lower().endswith((".jpg", ".jpeg", ".png")): items.append((os.path.join(class_name, fname), class_name)) counter = Counter(label for _, label in items) print(f"总样本数: {len(items)}") for label, cnt in counter.most_common(): print(f"{label}: {cnt}")这段代码的逻辑是先探测标注文件,再兜底读目录,最后统计类别分布。参数说明:annotations.json和labels.csv是常见的两种标注文件名,实际以你拿到的文件名为准,如果叫别的名字就把os.path.exists里的路径改掉。encoding="utf-8"在Windows下最好加上,否则中文类别名容易读成乱码。跑完这段,你会得到两个关键信息:类别数和各类别样本量——这两个数字直接决定后续要不要做重采样。
2.2 把类别分布打出来看,第一轮坑就藏在里面
统计结果不要只看总数,要把每个类别的数量打印出来。约5,700张听上去挺均衡,实际垂直领域数据集的类别分布往往偏得离谱:健康果可能占一半,某种病害只有三百张。如果某类占比低于5%,直接进训练的话,模型大概率把这类全部忽略。
下面这个表是这类数据集常见的分布形态,具体数字以你拿到手的实际统计为准:
| 类别 | 典型占比 | 风险提示 |
|---|---|---|
| 健康果 | 40%-50% | 数量充足,注意与病斑果的特征区分 |
| 炭疽病 | 15%-25% | 病斑特征明显,容易学偏到果皮反光 |
| 疫病 | 10%-15% | 与日灼病在颜色上容易混淆 |
| 日灼病 | 8%-12% | 颜色偏浅黄,和健康果亮度接近 |
| 其他/少见病害 | 5%-10% | 样本太少,需要过采样或者干脆合并类别 |
看到少数类只有两三百张的时候,先别急着删。图像分类里有一种做法叫类别合并——把症状相似、样本都少的病害合并成一个“其他病害”类。约5,700张的数据规模,合并后每类至少能有500张以上,训练才稳。这一步的取舍会在最终精度上体现得非常明显,我后面讲避坑章节时会再展开。
3. 数据清洗与划分:约5,700张里藏着多少重复和错标
3.1 去重脚本:哈希比对一步到位
已标注数据集不等于干净数据集。垂直领域数据集经常是从多个来源凑来的,同一个果实的不同角度照片、同一张图被缩放过两次、甚至同一张图在不同类目下各出现一次,这些情况都不少见。约5,700张的规模,人工翻一遍不现实,写脚本做哈希比对是标准做法。
import hashlib from collections import defaultdict def file_md5(path, chunk_size=8192): h = hashlib.md5() with open(path, "rb") as f: while chunk := f.read(chunk_size): h.update(chunk) return h.hexdigest() hash_map = defaultdict(list) duplicates = [] for file_path, label in items: # items 来自上一节的解析结果 full_path = os.path.join(dataset_root, file_path) if not os.path.exists(full_path): continue digest = file_md5(full_path) hash_map[digest].append((file_path, label)) for digest, file_list in hash_map.items(): if len(file_list) > 1: duplicates.extend(file_list) print(f"疑似重复文件组数: {len([v for v in hash_map.values() if len(v) > 1])}") for dup in duplicates: print(dup)这段脚本用MD5做文件级去重,全等重复(同一文件复制改名)一定会被抓到。参数说明:chunk_size=8192是分块读取的缓冲区大小,约5,700张图片全量算一遍MD5也就几秒到十几秒。抓到重复文件后,建议保留其中标注信息更详细的那条,另一条直接从items里剔除。
但要注意,MD5只能抓全等重复。同一张图被人为加了水印、调了亮度、改了尺寸,MD5就对不上了。这时候要用感知哈希(pHash)或者直接用图像缩放后做像素差比对。约5,700张的规模下pHash全量比对是秒级到分钟级的事,值得跑一遍:
import cv2 import numpy as np def perceptual_hash(image_path, size=(32, 32)): img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) img = cv2.resize(img, size, interpolation=cv2.INTER_AREA) # 用离散余弦变换的低频分量做特征 img_float = np.float32(img) / 255.0 dct = cv2.dct(img_float) # 取左上角8x8低频块,去掉第一个直流分量 low_freq = dct[:8, :8].flatten()[1:] median_val = np.median(low_freq) return (low_freq > median_val).astype(np.uint8) # 对items抽样计算,两两比较汉明距离,距离小于阈值视为疑似重复这段不做全量两两对比——约5,700张两两比对的复杂度是O(n^2),没必要。常见做法是抽样几百张做快速筛查,确认数据源风格差异大不大。如果这步发现大量近重复图,我一般会写一个窗口式的增量去重,把汉明距离小于10的图片打印出来人工二次确认。注意:去重和后面的数据划分是两件事,先去重再划分,顺序不能反过来。
3.2 按“来源组”而不是“单张图”划分训练验证集
划分训练集、验证集、测试集,约5,700张的常规比例是7:1.5:1.5,也就是训练约3,990张,验证和测试各约855张。但真正影响模型可信度的不是比例,而是划分的单位。如果你拿到的是多角度拍摄的数据,同一颗西瓜出现在多个文件里,随机按单张图去划分,训练集和验证集里会出现同一个果实的相似照片,验证集精度虚高,落地实测立刻打回原形。
正确做法是先看文件名或者目录结构里有没有“来源分组”的信息。有的数据集文件名类似field1_20230912_001.jpg,前半段就是分组键;有的按采集批次分了子目录。按组划分再随机抽组,代码如下:
import random from collections import defaultdict group_items = defaultdict(list) for file_path, label in items: # 按文件名前缀分组,常见分隔符是 "_" 或 "-" group_key = file_path.split("_")[0] group_items[group_key].append((file_path, label)) groups = list(group_items.keys()) random.seed(42) random.shuffle(groups) train_groups = groups[: int(len(groups) * 0.7)] val_groups = groups[int(len(groups) * 0.7) : int(len(groups) * 0.85)] test_groups = groups[int(len(groups) * 0.85) :] train_items = [item for g in train_groups for item in group_items[g]] val_items = [item for g in val_groups for item in group_items[g]] test_items = [item for g in test_groups for item in group_items[g]] print(f"训练集: {len(train_items)}, 验证集: {len(val_items)}, 测试集: {len(test_items)}")这段代码的核心逻辑是:先提取分组键,对组做随机化,再按组划分。参数说明:split("_")[0]是常见的分组键提取方式,实际要以文件名的命名规则为准,可能是前两个字段、也可能是中间某段,看了文件名格式再定。random.seed(42)固定随机种子,保证实验可复现。如果文件名看不出分组信息,退而求其次的做法是检查图片的采集时间戳或EXIF信息,但这步比较费工时,不是所有数据集都具备条件。
4. 训练第一个模型:用预训练图像分类模型做迁移学习
4.1 为什么约5,700张不要从零训练
在图像分类这个任务上,约5,700张数据做从零训练是笔亏本买卖。ResNet从零训到收敛需要百万级数据支撑,5,700张连正则化都压不住,模型直接过拟合到训练集。现在主流的图像分类模型都在迁移学习路径上:用ImageNet预训练权重做初始化,冻结前几层,只微调后面几层和分类头。这是垂直领域小数据集的标准打法,也是工业界最稳的做法。
模型怎么选?ResNet18和EfficientNet-B0是两类代表:ResNet系列结构简单、显存占用小、在CPU上也能做推理;EfficientNet-B0精度略高且参数更少,但对输入分辨率更敏感,默认的224x224输入表现不错。如果你只想快速在西瓜病识别上跑通流程,ResNet18足够;在意的指标是Top-1准确率而不是推理速度,就上EfficientNet-B0。至于ViT这类强Transformer架构,在约5,700张规模下预训练权重的作用没有CNN阵营那么强,容易在微调后期精度爬不上去。
import torch import torch.nn as nn import torchvision.models as models from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image class WatermelonDataset(Dataset): def __init__(self, items, dataset_root, transform=None): self.items = items # [(file_path, label), ...] self.root = dataset_root self.transform = transform self.class_to_idx = self._build_class_index() def _build_class_index(self): unique_labels = sorted(set(label for _, label in self.items)) return {label: idx for idx, label in enumerate(unique_labels)} def __len__(self): return len(self.items) def __getitem__(self, idx): file_path, label = self.items[idx] img = Image.open(os.path.join(self.root, file_path)).convert("RGB") if self.transform: img = self.transform(img) return img, self.class_to_idx[label] transform_train = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p=0.5), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_eval = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) num_classes = len(set(label for _, label in items)) model.fc = nn.Linear(model.fc.in_features, num_classes) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = model.to(device) train_loader = DataLoader( WatermelonDataset(train_items, dataset_root, transform_train), batch_size=32, shuffle=True, num_workers=4 ) val_loader = DataLoader( WatermelonDataset(val_items, dataset_root, transform_eval), batch_size=32, shuffle=False, num_workers=4 ) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) for epoch in range(20): model.train() running_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() # 每个epoch后验证 model.eval() correct, total = 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(f"Epoch {epoch+1:02d} | Loss: {running_loss/len(train_loader):.4f} | Val Acc: {val_acc:.4f}")这段代码是完整的训练骨架,几个关键参数值得细说:
models.ResNet18_Weights.IMAGENET1K_V1:加载ImageNet预训练权重。这里用带权重的枚举写法,不推荐pretrained=True这种旧方式,因为新版torchvision已经把参数名改了。model.fc = nn.Linear(model.fc.in_features, num_classes):替换最后一层全连接,输出维度改成实际类别数。分类头的初始化会用默认方式,不需要额外处理。batch_size=32:约5,700张训练集,一个epoch大约125个batch,20个epoch是2,500次迭代,在单张消费级显卡上几分钟就能跑完一轮实验。lr=3e-4:迁移学习的典型学习率区间是1e-4到5e-4,3e-4是一个安全中间值。如果发现Loss震荡剧烈,降一个数量级到3e-5再试。
几个训练时的注意点:ColorJitter对西瓜病识别很敏感——过度调饱和度会让病斑颜色失真,建议把saturation调低到0.1;RandomCrop配合Resize是先放大到256再随机裁剪到224,相当于引入轻微的尺度扰动,对病斑大小变化有抑制作用。训练过程中盯着验证集Loss,如果验证Loss在第10个epoch之后不降反升,大概率是过拟合了,有两种处理:提前停止,或者把学习率降到1e-5继续跑几个epoch。
4.2 类别不均衡时的损失函数调整
如果第2章的统计结果显示少数类占比很低,直接用CrossEntropyLoss会让模型偏向多数类。这时候可以给损失函数加类别权重,权重和样本量成反比:
from collections import Counter class_counts = Counter(label for _, label in train_items) total = sum(class_counts.values()) class_weights = torch.tensor( [total / (len(class_counts) * class_counts[label]) for label in sorted(class_counts.keys())], dtype=torch.float32 ).to(device) criterion = nn.CrossEntropyLoss(weight=class_weights)这个权重公式的含义是:每个类的权重 = 总样本数 / (类别数 × 该类样本数)。样本量少的类别权重放大,样本量多的类别权重缩小。比如某类只有300张,权重算出来是约3.2,另一类有2,000张,权重约0.48,训练时模型会“更在意”样本少的类。
加了类别权重之后,验证集指标不要只看整体准确率,要看每一类的召回率。约5,700张这个规模,小类可能只有几十个样本,整体准确率85%但小类召回率为0的情况在西瓜病识别里太常见了,我下面会把这个问题单独放进避坑清单。
5. 避坑指南:西瓜病识别训练里最常见的四个翻车现场
5.1 验证集精度高,实际测试却翻车
做图像分类的人一定见过这个场景:训练时验证集精度一路爬到92%,拿到新的西瓜图片一测,识别结果完全不对。发现这类问题要先去查数据划分方式。之前说过,如果同一来源(同一颗西瓜、同一块地、同一天拍摄)的近重复图片被拆进了训练集和验证集,验证时模型见过了近似图像,成绩自然虚高。解决方法是回到第3章的按组划分思路,重新生成三个集合后再训练。另外一个隐蔽原因是测试图片的来源分布和训练集差异太大,比如训练集全是田间拍摄,测试却用了仓库里的照片,光照、背景、拍摄距离全部变化,模型见过的特征对不上。处理手段是手动把测试集里识别错误的图挑出来看,确认是特征分布漂移还是模型本身学错了。
5.2 病斑太小:模型学的是果皮反光,不是病斑
西瓜病斑在整张图中的占比往往很小,尤其是早期病斑,可能只有十几个像素。模型在这种数据上学到的特征可能是果皮的高光区域、叶片的绿色背景、甚至照片角落的水印,而不是真正的病斑纹理。我遇到过一版模型看起来精度很高,用热力图一看,关注的全是果实边缘的反光区域。解决办法有三个:一是先做目标裁剪,把果实区域从背景中分离出来,再送入分类模型,常见做法是先用目标检测器定位果实,或者人工裁剪后重新训练;二是把输入分辨率从224提高到384,给病斑多留一些像素;三是用Grad-CAM这类可视化工具逐张检查,发现学错特征就去调整预处理或者换模型。在约5,700张的规模下,优先尝试后两种方案,成本最低见效最快。
5.3 少数类完全学不动
训练结束时少数类准确率依然是0的问题,多数来自训练过程中模型从未把该类的loss降到有效范围。排查手段是在每个epoch打印loss per class,如果某个类平均loss一直不下降,就是没学进去。原因通常有两个:一是类别权重加得不够,把权重调大一倍再试;二是该类的训练样本之间差异太大,病斑形态、角度、光照各不相同,模型通过有限的50张样本学不到稳定的类别特征,这时候需要检查标注是否有误,把同类样本铺开人工过一遍。如果标注没问题,另一个方向是把类别合并,同类病害合并成一个大类,虽然细粒度信息丢了,但至少模型在这个类上有区分度。
5.4 训练过程中Loss出现NaN
约5,700张的小数据集上用Adam优化器,Loss出现NaN的概率不大,但也不是没有。常见诱因是学习率过大导致梯度爆炸,或者输入图像中有损坏的图片文件(截断的JPEG),数据读取时返回了空张量。一个排查思路是先把学习率降到1e-5跑几个epoch看是否复现,不复现就是学习率问题;另一个做法是所有训练图片做一遍完整性校验——用PIL打开并转换RGB,抛异常的图片直接从items里剔除。另外,使用混合精度训练(AMP)时也要注意,小数据集上精度缩放策略设置不当也会导致NaN,通常把torch.cuda.amp.GradScaler的init_scale参数降低即可。
6. 最后的进阶技巧:用Grad-CAM验证模型到底在看什么地方
训练完模型,先别急着验收精度指标。我现在的习惯是花半小时做一轮可视化归因,这个动作帮我躲开了至少三次误报。Grad-CAM能生成热力图,标记出模型做决策时关注的图像区域,这对病斑类任务非常关键——判断模型看的是“病斑”还是“背景”一目了然。
from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from pytorch_grad_cam.utils.image import show_cam_on_image import matplotlib.pyplot as plt model.eval() cam_model = GradCAM(model=model, target_layers=[model.layer4[-1]]) # 随便挑一张测试集里的病斑图 img_path, true_label = test_items[0] img = Image.open(os.path.join(dataset_root, img_path)).convert("RGB") input_tensor = transform_eval(img).unsqueeze(0).to(device) target = [ClassifierOutputTarget(class_idx)] grayscale_cam = cam_model(input_tensor=input_tensor, targets=target)[0] visualization = show_cam_on_image( np.float32(np.array(img.resize((224, 224)))) / 255.0, grayscale_cam, use_rgb=True ) plt.imshow(visualization) plt.axis("off") plt.show()这段代码把热力图叠加在原始图像上,逻辑说明:target_layers=[model.layer4[-1]]选择ResNet最后一个卷积块作为特征提取层,病斑的位置信息在这一层仍然保留;ClassifierOutputTarget指定要看哪个类别的激活区域——传入错误类别,就能看到模型在判断这个类时看的是哪里。如果热力图中心偏离病斑区域、落在果皮反光或背景上,说明模型没学到该学的特征。
配合一个简单的量化检查:手动框出测试图像上的病斑区域(用标注框或者简单抠图),计算热力图高亮区域和病斑区域的IoU。如果IoU低于0.3,这张图的预测结果就要打上问号。用这个方法在约5,700张的验证集或测试集上抽100张左右做一次抽样检查,基本就能判断模型学的是“西瓜病”还是“西瓜反光”了。这个习惯帮我筛掉过很多表面高精度、实际不可用的模型权重。希望帮到你。
本文还有配套的精品资源,点击获取