基于 Python 与 AI 的智能害虫识别系统实战

2. 搭建 CNN 模型（迁移学习版）

从零训练 CNN 模型对算力和数据要求较高。我们采用迁移学习策略，加载预训练的 ResNet50 权重，冻结主干网络参数，仅替换最后一层分类器以适应新的害虫类别数量。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models

# 加载预训练的 ResNet50
model = models.resnet50(pretrained=True)

# 冻结主干网络参数（只训练最后一层）
for param in model.parameters():
    param.requires_grad = False

# 替换最后一层分类器
num_classes = len(classes)  # 根据实际害虫类别数调整
model.fc = nn.Linear(model.fc.in_features, num_classes)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

# 设备配置（有 GPU 用 GPU，没有用 CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

3. 训练与评估模型

将预处理好的数据封装为 PyTorch 的 DataLoader，开始训练循环，并在验证集上监控精度变化。

from torch.utils.data import TensorDataset, DataLoader, random_split

# 转换为张量
X = torch.tensor(processed_data, dtype=torch.float32)
y = torch.tensor(labels, dtype=torch.long)
dataset = TensorDataset(X, y)

# 划分训练集和验证集
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

# 创建 DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

# 训练函数
def train_model(model, train_loader, val_loader, epochs=10):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        
        # 训练轮次
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        
        # 验证轮次
        model.eval()
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()
        
        print(f"Epoch {epoch+1}/{epochs} | Loss: {running_loss/len(train_loader):.4f} | Val Accuracy: {100* val_correct / val_total:.2f}%")
    return model

# 开始训练
trained_model = train_model(model, train_loader, val_loader, epochs=15)

# 保存模型
torch.save(trained_model.state_dict(), "pest_recognition_model.pth")

4. 使用 Gradio 部署成 Web 应用

训练完成后，需要一个交互界面来展示结果。Gradio 可以快速构建可视化页面，支持上传图片并实时返回识别结果。

import gradio as gr

# 加载训练好的模型
model.load_state_dict(torch.load("pest_recognition_model.pth"))
model.eval()

# 定义识别函数
def recognize_pest(img):
    # 预处理图片
    img = cv2.resize(img, (224, 224))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img / 255.0
    img = np.transpose(img, (2, 0, 1))
    img = torch.tensor(img, dtype=torch.float32).unsqueeze(0).to(device)
    
    # 模型预测
    with torch.no_grad():
        outputs = model(img)
        _, predicted = torch.max(outputs, 1)
        pest_name = classes[predicted.item()]
        confidence = torch.softmax(outputs, dim=1)[0][predicted.item()].item() * 100
    
    return f"识别结果：{pest_name} | 置信度：{confidence:.2f}%"

# 创建 Gradio 界面
interface = gr.Interface(
    fn=recognize_pest,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Textbox(),
    title="智能害虫识别助手",
    description="上传农作物害虫图片，自动识别害虫种类"
)

# 启动应用
interface.launch()

运行代码后，本地会生成访问链接。打开浏览器上传害虫图片，系统将在几秒内显示识别结果和置信度。

进阶优化方向

若需进一步提升效果，可考虑以下方向：

1. 数据增强：使用 torchvision.transforms 进行旋转、翻转、裁剪等操作，扩充数据集多样性。
2. 调参优化：调整学习率、批次大小（batch_size），或尝试不同的优化器（如 SGD）。
3. 多标签识别：如果单张图片包含多种害虫，修改模型输出层为多标签分类结构。
4. 部署优化：使用 ONNX 转换模型格式，或部署至云服务器实现远程访问。

总结

本项目涵盖了 AI 开发的核心流程：数据预处理、模型搭建、训练评估及部署应用。对于 Python 初学者而言，这是一个将理论转化为实践的优质案例；对于农业从业者，它提供了一种低成本的技术解决方案。通过几行代码，即可将 AI 能力落地到具体场景中。