病虫害识别以及分类的实现

import numpy as np
import pandas as pd
import os
import random
from tqdm import tqdm
from textwrap import wrap

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

import cv2
import matplotlib.pyplot as plt
import seaborn as sns

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

import timm

# 读取类别文件
f = open('../input/ip02-dataset/classes.txt')
label = []
name = []
for line in f.readlines():
    label.append(int(line.split()[0]))
    name.append(' '.join(line.split()[1:]))
classes = pd.DataFrame([label, name]).T
classes.columns = ['label','name']
classes

# 读取训练集、测试集和验证集
train_df = pd.read_csv('../input/ip02-dataset/train.txt',sep=' ',header=None, engine='python')
train_df.columns = ['image_path','label']

test_df = pd.read_csv('../input/ip02-dataset/test.txt',sep=' ',header=None, engine='python')
test_df.columns = ['image_path','label']

val_df = pd.read_csv('../input/ip02-dataset/val.txt',sep=' ',header=None, engine='python')
val_df.columns = ['image_path','label']

train_df.head()

# 定义目录和超参数
TRAIN_DIR = '../input/ip02-dataset/classification/train'
TEST_DIR = '../input/ip02-dataset/classification/test'
VAL_DIR = '../input/ip02-dataset/classification/val'
LR = 2e-5
BATCH_SIZE = 8
EPOCH = 2

device = torch.device('cuda')

# 可视化部分训练数据
fig, axs = plt.subplots(10,11,figsize=(30,30))
images = []
for i in classes.label:
    random_img = random.choice(train_df[train_df.label==i-1].image_path.values)
    label = classes.name[i-1]
    img = plt.imread(os.path.join(TRAIN_DIR,str(i-1),random_img))
    images.append(img)

[ax.imshow(image) for image,ax in zip(images,axs.ravel())]
[ax.set_title("\n".join(wrap(label,20))) for label,ax in zip(list(classes.name),axs.ravel())]
[ax.set_axis_off() for ax in axs.ravel()]
plt.show()

# 定义模型
class InsectModel(nn.Module):
    def __init__(self,num_classes):
        super(InsectModel, self).__init__()
        self.num_classes = num_classes
        self.model = timm.create_model('vit_base_patch16_224',pretrained=True,num_classes=num_classes)
    def forward(self, image):
        return self.model(image)

# 定义数据增强
def train_transform():
    return A.Compose([
        A.HorizontalFlip(),
        A.RandomRotate90(),
        A.RandomBrightnessContrast(),
        A.Resize(224, 224),
        ToTensorV2()])

def valid_transform():
    return A.Compose([
        A.Resize(224,224),
        ToTensorV2()])

def collate_fn(batch):
    return tuple(zip(*batch))

# 定义数据集
class InsectDataset(Dataset):
    def __init__(self, image, image_dir, transforms=None):
        self.image_info = image
        self.transforms = transforms
        self.imgdir = image_dir
    def __len__(self):
        return self.image_info.shape[0]
    def __getitem__(self, index):
        image_info = self.image_info[index]
        image = cv2.imread(os.path.join(self.imgdir,str(image_info[1]),image_info[0]),cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        image /= 255.

        if self.transforms is not None:
            image = self.transforms(image = image)['image']

        label = image_info[1]

        image = torch.as_tensor(image, dtype=torch.float32)
        label = torch.as_tensor(label, dtype=torch.long)
        return image, label

# 创建数据加载器
train_dataset = InsectDataset(image=train_df.values, 
                              image_dir=TRAIN_DIR, 
                              transforms=train_transform())
train_data_loader = DataLoader(train_dataset,
                              batch_size=BATCH_SIZE,
                              shuffle=True,
                              num_workers=2)
val_dataset = InsectDataset(image=val_df.values,
                            image_dir=VAL_DIR,
                            transforms=valid_transform())
val_data_loader = DataLoader(val_dataset,
                             batch_size=BATCH_SIZE,
                             shuffle=True,
                             num_workers=2)

# 定义平均计量器
class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.loss = 0
        self.correct = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, loss,correct, n=1):
        self.loss = loss
        self.correct += correct
        self.sum += loss * n
        self.count += n

        self.avg = self.sum / self.count
        self.acc = self.correct / self.count

class Accuracy(object):
    def __init__(self):
        self.reset

# 定义训练函数
def train_fn(data_loader, model, criterion, device, optimizer, epoch):
    model.train()
    criterion.train()

    summary = AverageMeter()
    tk0 = tqdm(data_loader, total=len(data_loader))
    for step, (images, labels) in enumerate(tk0):
        images = images.to(device, non_blocking = True).float()
        labels = labels.to(device, non_blocking = True).long()

        output = model(images)
        loss = criterion(output, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        preds = output.softmax(1).argmax(1)
        correct = (preds == labels).sum().item()

        summary.update(loss.item(),correct, BATCH_SIZE)
        tk0.set_postfix(loss=summary.avg, acc=summary.acc, epoch=epoch+1)
    return summary

# 定义验证函数
def eval_fn(data_loader, model, criterion, device, epoch):
    model.eval()
    criterion.eval()

    summary = AverageMeter()
    tk0 = tqdm(data_loader, total=len(data_loader))
    with torch.no_grad():
        for step, (images, labels) in enumerate(tk0):
            images = images.to(device, non_blocking = True).float()
            labels = labels.to(device, non_blocking = True).long()

            output = model(images)
            loss = criterion(output, labels)

            preds = output.softmax(1).argmax(1)
            correct = (preds == labels).sum().item()

            summary.update(loss.item(), correct, BATCH_SIZE)
            tk0.set_postfix(loss=summary.avg, acc=summary.acc, epoch=epoch+1)
    return summary

os.environ['WANDB_CONSOLE'] = 'off'

# 运行训练和验证
def run():
    model = InsectModel(num_classes=102)
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    criterion = criterion.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
    best_loss = 10**5
    for epoch in range(0, EPOCH):
        train_loss = train_fn(train_data_loader, model, criterion, device, optimizer, epoch)
        val_loss = eval_fn(val_data_loader, model, criterion, device, epoch)
        if val_loss.avg < best_loss:
            best_loss = val_loss.avg
            torch.save(model.state_dict(), f'vit_best.pth')
        print(f'Epoch {epoch+1+0:03}: | Train Loss: {train_loss.avg:.5f} | Val Loss: {val_loss.avg:.5f}')

run()

# 加载最佳模型并进行预测
model = InsectModel(num_classes=102)
model.load_state_dict(torch.load("./vit_best.pth"))
images, labels = next(iter(val_data_loader))
preds = model(images).softmax(1).argmax(1)

# 可视化预测结果
fig, axs = plt.subplots(2,4,figsize=(13,8))
[ax.imshow(image.permute((1,2,0))) for image,ax in zip(images,axs.ravel())]
[ax.set_title("\n".join(wrap(f'实际: {classes.name[label.item()]} 预测: {classes.name[pred.item()]}',30)),color = 'g' if label.item()==pred.item() else 'r') for label,pred,ax in zip(labels,preds,axs.ravel())]
[ax.set_axis_off() for ax in axs.ravel()]
plt.show()

import cv2
import torch
import time
import numpy as np
import pandas as pd
import timm
import torch.nn as nn
from textwrap import wrap

# 视频流地址
url = "http://192.168.2.8/stream"

# 打开摄像头
cap = cv2.VideoCapture(url)

# 加载模型
class InsectModel(nn.Module):
    def __init__(self, num_classes):
        super(InsectModel, self).__init__()
        self.num_classes = num_classes
        self.model = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=num_classes)
    def forward(self, image):
        return self.model(image)

try:
    model = InsectModel(num_classes=102)
    model.load_state_dict(torch.load("./vit_best.pth"))
    model.eval()
    print("模型加载成功")
except Exception as e:
    print(f"模型加载失败: {e}")

# 定义图像预处理函数
def preprocess(image):
    try:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        image /= 255.0
        image = cv2.resize(image, (224, 224))
        image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
        print("图像预处理成功")
        return image
    except Exception as e:
        print(f"图像预处理失败: {e}")
        return None

# 读取类别名称
try:
    f = open('classes.txt')
    label = []
    name = []
    for line in f.readlines():
        label.append(int(line.split()[0]))
        name.append(' '.join(line.split()[1:]))
    classes = pd.DataFrame([label, name]).T
    classes.columns = ['label', 'name']
    print("类别名称读取成功")
except Exception as e:
    print(f"类别名称读取失败: {e}")

last_recognition_time = time.time()

while True:
    # 读取摄像头的帧
    ret, frame = cap.read()

    if not ret:
        print("无法读取视频流，尝试重新连接...")
        cap.release()
        cap = cv2.VideoCapture(url)
        continue

    # 在窗口中显示帧
    cv2.imshow('Camera', frame)

    # 每隔10秒进行一次识别
    if time.time() - last_recognition_time > 10:
        image = preprocess(frame)
        if image is not None:
            with torch.no_grad():
                output = model(image)
                pred = output.softmax(1).argmax(1).item()
                if pred < len(classes):
                    label_name = classes.name[pred]
                # else:
                #     label_name = "无法识别"   暂时不起作用
                print(f"识别结果: {label_name}")
        last_recognition_time = time.time()

    # 按下'q'键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()