FasterRCNN入门案例水稻图像目标检测新手友好入门案例

目录

依赖环境

代码概述

引用库

读取数据指定目录

数据集划分

数据集加载Dataset类

特征增强处理

预训练模型定义

评估指标定义

实例化训练集和测试集

设置硬件调取一个batch

可视化

​编辑

激活选定硬件,初始化损失函数参数

模型训练

模型测试和验证

配套数据集和源码下载

pytorch实现FasterRCNN目标检测入门案例


依赖环境

opencv、numpy、pandas、pillow、albumentations、pytorch、torchvision、matplotlib

代码概述

引用库

import pandas as pd  # 用于数据处理和分析的库
import numpy as np  # 数值计算库,提供高效的多维数组操作
import cv2  # OpenCV库,用于进行图像处理和计算机视觉任务
import os  # 提供与操作系统交互的功能,如文件路径操作
import re  # 用于执行正则表达式,进行字符串匹配和处理

from PIL import Image  # 用于图像打开、处理和保存

import albumentations as A  # 图像增强库,用于提高模型的泛化能力
from albumentations.pytorch.transforms import ToTensorV2  # 将图像数据转换为PyTorch张量

import torch  # PyTorch基础库,提供多维数组(张量)和自动求导功能
import torchvision  # 提供图像视频处理工具,和模型训练或加载预训练模型

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor  # Faster R-CNN的预测器部分
from torchvision.models.detection import FasterRCNN  # Faster R-CNN模型
from torchvision.models.detection.rpn import AnchorGenerator  # 生成锚点用于RPN(区域提议网络)

from torch.utils.data import DataLoader, Dataset  # 数据加载和批处理
from torch.utils.data.sampler import SequentialSampler  # 数据采样器,按顺序采样

from matplotlib import pyplot as plt  # 绘图库,用于数据可视化

读取数据指定目录

DIR_INPUT = '.'  # 设置基本目录为当前目录
DIR_TRAIN = f'{DIR_INPUT}/train'  # 设置训练数据集的目录
DIR_TEST = f'{DIR_INPUT}/test'  # 设置测试数据集的目录

train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')  # 使用pandas加载训练数据集的CSV文件
train_df.shape  # 显示加载的DataFrame的形状(即,行数和列数)
  • DIR_INPUT - 指定了数据的基本输入目录,通常设置为项目的根目录或数据存储的目录。
  • DIR_TRAINDIR_TEST - 这两行代码分别设置了训练和测试数据集的文件路径。这有助于在代码中引用这些目录时更有条理和清晰。
  • train_df - 这行代码通过读取CSV文件加载训练数据到train_df这个DataFrame中。DataFrame是pandas库中用于存储和操作结构化数据的主要数据结构。
  • train_df.shape - 这行代码用于输出DataFrame的维度,即里面有多少行(样本)和列(特征)。这通常用来快速获取数据集的大小和结构。

这部分代码是数据预处理和初步检查的基础部分,确保数据正确加载并且可以进一步用于训练模型。如果你需要更进一步的帮助,例如数据的可视化或预处理,请继续提供相关的代码或需求。

数据集划分

# 在DataFrame中初始化四个新列,用于存储边界框的x, y坐标和宽度、高度
train_df['x'] = -1
train_df['y'] = -1
train_df['w'] = -1
train_df['h'] = -1

# 定义一个函数,用于从字符串中提取边界框数据
def expand_bbox(x):
    r = np.array(re.findall("([0-9]+[.]?[0-9]*)", x))  # 使用正则表达式提取数字
    if len(r) == 0:  # 如果没有找到任何数字,返回默认值
        r = [-1, -1, -1, -1]
    return r

# 将'bbox'列的字符串转换为具体的x, y, w, h数值,并更新到对应列
train_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))
train_df.drop(columns=['bbox'], inplace=True)  # 删除原始的'bbox'列
train_df['x'] = train_df['x'].astype(np.float32)  # 转换数据类型为浮点数
train_df['y'] = train_df['y'].astype(np.float32)
train_df['w'] = train_df['w'].astype(np.float32)
train_df['h'] = train_df['h'].astype(np.float32)

# 获取所有唯一的图像ID
image_ids = train_df['image_id'].unique()
valid_ids = image_ids[-665:]  # 划分最后665个ID为验证集
train_ids = image_ids[:-665]  # 剩余的为训练集

# 创建验证集和训练集的DataFrame
valid_df = train_df[train_df['image_id'].isin(valid_ids)]
train_df = train_df[train_df['image_id'].isin(train_ids)]

# 输出验证集和训练集的形状(即,行数和列数)
valid_df.shape, train_df.shape
  • 初始化列 - 首先为边界框的坐标和尺寸初始化列,设置默认值为-1。
  • expand_bbox函数 - 这个函数用于解析边界框数据,它从格式化的字符串中提取边界框的坐标和尺寸。
  • 数据类型转换 - 将提取的字符串转换为浮点数,以便于后续的数值计算。
  • 数据集划分 - 根据图像ID划分数据为训练集和验证集。这里的分割是基于图像ID的唯一性,确保同一图像的所有数据要么在训练集,要么在验证集中。

这部分代码主要负责数据预处理和训练验证集的分割,这是模型训练前的重要步骤,确保数据的正确格式和合适的分布。

数据集加载Dataset类

class WheatDataset(Dataset):
    # 初始化方法
    def __init__(self, dataframe, image_dir, transforms=None):
        super().__init__()
        self.image_ids = dataframe['image_id'].unique()  # 获取所有唯一的图像ID
        self.df = dataframe  # 数据帧
        self.image_dir = image_dir  # 图像文件的目录
        self.transforms = transforms  # 图像变换(如增强)

    # 获取单个样本
    def __getitem__(self, index: int):
        image_id = self.image_ids[index]  # 根据索引获取图像ID
        records = self.df[self.df['image_id'] == image_id]  # 获取该图像ID的所有记录

        # 读取图像并转换颜色空间
        image = cv2.imread(f'{self.image_dir}/{image_id}.jpg', cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        image /= 255.0  # 归一化图像数据

        # 计算边界框并调整格式
        boxes = records[['x', 'y', 'w', 'h']].values
        boxes[:, 2] = boxes[:, 0] + boxes[:, 2]  # 转换宽度为右下角x坐标
        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]  # 转换高度为右下角y坐标
        
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        area = torch.as_tensor(area, dtype=torch.float32)  # 计算每个框的面积

        labels = torch.ones((records.shape[0],), dtype=torch.int64)  # 所有物体的标签为1(只有一类)
        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)  # 假设没有物体是重叠的

        # 构建目标字典,存储边界框等信息
        target = {}
        target['boxes'] = boxes
        target['labels'] = labels
        target['image_id'] = torch.tensor([index])
        target['area'] = area
        target['iscrowd'] = iscrowd

        # 如果有变换应用变换
        if self.transforms:
            sample = {
                'image': image,
                'bboxes': target['boxes'],
                'labels': labels
            }
            sample = self.transforms(**sample)
            image = sample['image']
            target['boxes'] = torch.stack(tuple(map(torch.tensor, zip(*sample['bboxes'])))).permute(1, 0)

        return image, target, image_id

    # 返回数据集大小
    def __len__(self) -> int:
        return self.image_ids.shape[0]
  • init 方法:初始化数据集类,设置图像ID、数据帧、图像目录和变换。
  • getitem 方法:用于按索引加载和处理单个图像及其标注数据,适用于训练时每次迭代获取数据。
  • len 方法:返回数据集中的图像总数,使得 PyTorch 的 DataLoader 可以知道有多少项数据。

这个类提供了从指定目录加载图像和标注,应用预处理和变换,以及格式化输出以供模型训练使用的完整工作流。

特征增强处理

# 定义训练时使用的图像变换
def get_train_transform():
    return A.Compose([
        A.HorizontalFlip(p=0.5),  # 以50%的概率水平翻转图像
        ToTensorV2(p=1.0)  # 将图像转换为PyTorch的Tensor格式
    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})
    # bbox_params指定了边界框的格式和相关标签字段

# 定义验证时使用的图像变换
def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0)  # 验证时只进行Tensor转换,不进行其他增强
    ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})
    # 同样设置边界框的格式为Pascal VOC
  • get_train_transform():这个函数用于创建训练阶段使用的变换。包括水平翻转图像的操作,这可以增加模型的泛化能力,因为模型需要学习识别不同方向的对象。ToTensorV2 用于确保图像数据以适合 PyTorch 处理的形式输入模型。
  • get_valid_transform():验证阶段的函数只包含将图像转换为张量的操作,通常不包括任何形式的数据增强。这是因为验证的目的是评估模型在未修改过的数据上的表现。

bbox_params 参数确保在进行图像变换时,相应的边界框(bbox)也被正确地变换。'format': 'pascal_voc' 表示使用 Pascal VOC 格式(xmin, ymin, xmax, ymax),这是物体检测中常用的边界框格式。'label_fields': ['labels'] 告诉 Albumentations 库标签数据在哪里,以便在处理时保持与边界框同步。

预训练模型定义

# 加载一个预训练的Faster R-CNN模型,使用的是ResNet-50作为backbone,带有特征金字塔网络(FPN)
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

num_classes = 2  # 定义类别数:1个小麦类别加上背景

# 获取分类器的输入特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features

# 使用新的分类头来替换原来的预训练分类头
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
  • 模型加载:使用 torchvision.models.detection 模块来加载一个预训练的 Faster R-CNN 模型,这里使用的是基于 ResNet-50 和 FPN 的结构。这种结构提供了良好的特征提取能力,非常适合复杂的视觉任务,如物体检测。
  • 类别数定义:这里定义的 num_classes 为 2,考虑到了一个目标类别(小麦)加上一个背景类别。在 Faster R-CNN 中,背景被视为一个独立的类别。
  • 自定义头部:模型的 ROI (Region of Interest) head 包括一个用于分类的部分,这里通过读取原有模型头部的输入特征数(in_features),并以此创建一个新的 FastRCNNPredictor,将其用作新的分类器头部。这样做的目的是将模型的输出调整为适应当前任务的类别数。

这些更改确保了模型能够适应于特定的检测任务,同时也利用了预训练模型在通用任务(如 COCO 数据集上的物体检测)上获得的丰富特征和经验,这有助于提高模型在特定任务上的性能。

评估指标定义

class Averager:
    # 初始化方法
    def __init__(self):
        self.current_total = 0.0  # 初始化总和为0
        self.iterations = 0.0  # 初始化迭代次数为0

    # 向平均器发送新的值
    def send(self, value):
        self.current_total += value  # 将新值加到总和中
        self.iterations += 1  # 迭代次数加1

    # 计算当前的平均值
    @property
    def value(self):
        if self.iterations == 0:
            return 0  # 如果没有迭代,返回0防止除以0
        else:
            return 1.0 * self.current_total / self.iterations  # 计算平均值

    # 重置平均器
    def reset(self):
        self.current_total = 0.0  # 重置总和为0
        self.iterations = 0.0  # 重置迭代次数为0
  • init:初始化函数,设置总和和迭代次数的初始值。这确保了每次创建 Averager 实例时,都从零开始。
  • send:此方法用于向总和中添加新的值,并递增计数器。它是在每次计算如损失或任何其他需要平均的度量后调用。
  • value:这是一个属性,通过装饰器 @property 实现,它使得每次访问 .value 时,都会动态计算当前的平均值。这样做的好处是可以随时获取最新的平均值,而不必显式调用一个方法。
  • reset:重置方法使得 Averager 实例可以在不同阶段(例如,新的训练周期开始时)重新使用,而无需创建新的实例。

Averager 类是一个简洁而有效的工具,适用于训练过程中监控和报告指标的平均值,如每个批次或每个周期的平均损失。这对于调试和优化模型非常有帮助。

实例化训练集和测试集

# 自定义的数据批处理函数
def collate_fn(batch):
    return tuple(zip(*batch))  # 将批量数据中的元素按组件重新组织

# 创建训练数据集和验证数据集实例
train_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())
valid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())

# 随机打乱索引并转为列表形式
indices = torch.randperm(len(train_dataset)).tolist()

# 创建训练数据加载器
train_data_loader = DataLoader(
    train_dataset,
    batch_size=4,  # 每批处理4个图像
    shuffle=False,  # 数据不进行额外的混洗(已经预先打乱)
    # num_workers=1,  # 线程数(这行被注释,使用默认设置)
    collate_fn=collate_fn  # 使用自定义的批处理函数
)

# 创建验证数据加载器
valid_data_loader = DataLoader(
    valid_dataset,
    batch_size=4,  # 同样每批4个图像
    shuffle=False,  # 验证数据通常不需要混洗
    # num_workers=1,  # 线程数(这行也被注释)
    collate_fn=collate_fn  # 同样使用自定义的批处理函数
)
  • collate_fn:这是一个重要的函数,用于决定如何将多个数据样本组合成一个批次。这里使用的 collate_fn 函数通过 zip 重新组织批数据,适合处理包含多个组件(如图像和其标注信息)的数据结构。
  • DataLoaderDataLoader 是 PyTorch 中用于加载数据的一个工具,它可以处理并行加载和数据批处理,使数据输入模型时更加高效。这里创建了两个数据加载器,一个用于训练,一个用于验证。设置 batch_size 为 4 表示每个批次处理 4 个图像。

这些设置是进行有效训练和验证的关键,因为它们确保数据以适当的方式被批处理和提供给模型,同时也利用了 DataLoader 的多线程能力(虽然在这里具体的 num_workers 被注释了,可能是为了避免在某些环境下的问题,如单线程环境或者与特定硬件的兼容性问题)。

设置硬件调取一个batch

# 确定使用的设备,如果GPU可用则使用GPU,否则使用CPU
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 从训练数据加载器中获取第一个批次的数据
images, targets, image_ids = next(iter(train_data_loader))

# 将图像列表中的每张图像转移到设定的设备(GPU或CPU)
images = list(image.to(device) for image in images)

# 将每个目标字典中的所有张量也转移到设定的设备
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

# 从第三个目标(索引为2)中提取边界框,并将其从GPU移动到CPU,并转换为整数类型的NumPy数组
boxes = targets[2]['boxes'].cpu().numpy().astype(np.int32)

# 将第三张图像(索引为2)的维度顺序从[C, H, W]调整为[H, W, C],并转移到CPU,转换为NumPy数组
sample = images[2].permute(1,2,0).cpu().numpy()
  • 设备选择:这段代码首先检查是否有 GPU 可用,并据此设置使用的设备。这是高效运行深度学习模型的关键步骤,因为 GPU 提供了显著更快的计算能力。
  • 数据提取和转移:代码从数据加载器中提取一个批次的图像、目标和图像ID。然后,它将图像和目标中的所有张量移动到先前确定的设备上,这是进行模型训练前的必要步骤。
  • 处理特定数据:选择特定图像和目标以进行进一步处理,例如可视化或特定操作。示例中提取了特定目标的边界框和对应图像的调整,以便可以直接处理这些数据。

这种数据处理方式典型地出现在准备数据进行模型训练或评估之前,确保数据在适当的设备上并以适合处理的格式。

可视化

# 创建一个图和一个子图,设置图的大小为16x8英寸
fig, ax = plt.subplots(1, 1, figsize=(16, 8))

# 遍历边界框数组boxes,每个box代表一个边界框
for box in boxes:
    # 使用OpenCV在图像sample上绘制矩形(边界框)
    cv2.rectangle(sample,
                  (box[0], box[1]),  # 矩形左上角
                  (box[2], box[3]),  # 矩形右下角
                  (220, 0, 0), 3)  # 矩形颜色为红色,线宽为3

# 隐藏坐标轴
ax.set_axis_off()

# 在子图上显示图像
ax.imshow(sample)
  • 图像和子图的创建:使用 matplotlib.pyplot.subplots 创建一个图和一个子图。这里指定了子图的大小,使得图像可以更加清晰地展示。
  • 边界框的绘制:通过遍历每一个边界框,使用 cv2.rectangle 在图像上绘制。这个函数接受左上角和右下角的坐标,并允许自定义边界框的颜色和线宽。
  • 图像显示:使用 imshow 方法将图像显示在子图上。set_axis_off 方法用于隐藏坐标轴,使图像显示更为美观,专注于内容本身。

这种方式通常用于数据探索和结果展示阶段,可以直观地评估目标检测模型识别出的边界框是否准确,或者是否需要调整模型的训练过程中的参数。

激活选定硬件,初始化损失函数参数

# 将模型移至之前选择的设备上(GPU或CPU)
model.to(device)

# 选择模型中所有需要梯度更新(可训练)的参数
params = [p for p in model.parameters() if p.requires_grad]

# 为选定的参数创建一个随机梯度下降优化器
# 设置学习率为0.005,动量为0.9,权重衰减为0.0005
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)

# 学习率调度器,这里被注释掉了,意味着学习率将保持不变
# lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)
lr_scheduler = None

# 设置训练的总轮数为2
num_epochs = 2
  • 模型迁移:将模型转移到适当的计算设备上是为了利用该设备的计算资源,例如GPU的并行处理能力。
  • 优化器设置:通过为模型的可训练参数配置优化器,可以控制参数的更新方式。使用SGD(随机梯度下降)优化器是深度学习中常见的选择,适合大多数任务。
  • 学习率调度器:通常用于调整训练过程中的学习率,以优化训练结果和加快收敛速度。这里虽然被设置为 None,通常可以根据需要激活和配置。

整体而言,这些设置对模型的训练过程至关重要,它们影响模型训练的效率和最终的性能。正确配置这些元素是确保模型能够有效学习和泛化的关键。

模型训练

# 创建一个用于跟踪平均损失的Averager实例
loss_hist = Averager()
itr = 1  # 初始化迭代计数器

# 循环执行指定的训练轮次
for epoch in range(num_epochs):
    loss_hist.reset()  # 重置平均损失计算器
    
    # 从数据加载器中迭代每个批次
    for images, targets, image_ids in train_data_loader:
        
        # 将图像和目标张量移至设定的设备
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        # 计算模型对当前批次的损失
        loss_dict = model(images, targets)

        # 将损失字典中的所有损失相加得到总损失
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()  # 将损失转为Python浮点数

        # 向平均损失计算器发送当前损失
        loss_hist.send(loss_value)

        # 清除旧的梯度
        optimizer.zero_grad()
        # 反向传播损失以计算梯度
        losses.backward()
        # 根据梯度更新模型参数
        optimizer.step()

        # 每50次迭代输出一次当前的损失
        if itr % 50 == 0:
            print(f"Iteration #{itr} loss: {loss_value}")

        itr += 1  # 更新迭代计数器
    
    # 如果有设置学习率调度器,则更新学习率
    if lr_scheduler is not None:
        lr_scheduler.step()

    # 输出每个轮次的平均损失
    print(f"Epoch #{epoch} loss: {loss_hist.value}")  
  • 损失计算:模型在每个批次上的输出用于计算损失,损失反映了当前模型性能的好坏。这一步是训练过程中至关重要的,因为它指导了模型的优化方向。
  • 优化过程:使用梯度下降法(通过 optimizer.step())更新模型的参数,目的是最小化损失函数。
  • 输出和监控:通过打印语句定期输出训练过程中的损失值,可以监控训练过程并进行调试。
  • 学习率调整:如果配置了学习率调度器,会在每个训练轮次后调整学习率,帮助模型更好地收敛。

这些步骤确保了训练过程的有效性和可追踪性,是实现有效深度学习模型训练的关键环节。

模型测试和验证

# 从验证数据加载器中获取一批数据
images, targets, image_ids = next(iter(valid_data_loader))

# 将图像和目标数据移动到设定的设备(GPU或CPU)
images = list(img.to(device) for img in images)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

# 提取特定目标的边界框并转换为NumPy数组
boxes = targets[1]['boxes'].cpu().numpy().astype(np.int32)
# 调整图像数据的维度顺序,准备显示
sample = images[1].permute(1,2,0).cpu().numpy()

# 设置模型为评估模式
model.eval()
cpu_device = torch.device("cpu")  # 创建一个CPU设备对象

# 在当前批次的图像上运行模型进行预测
outputs = model(images)
# 将输出数据移动到CPU
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]

# 创建一个图和一个子图,设置图的大小
fig, ax = plt.subplots(1, 1, figsize=(16, 8))

# 遍历边界框并在图像上绘制
for box in boxes:
    cv2.rectangle(sample,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (220, 0, 0), 3)  # 使用红色绘制边界框

# 隐藏坐标轴
ax.set_axis_off()
# 显示图像
ax.imshow(sample)
  • 数据提取与处理:从验证数据加载器中提取图像和目标,这些数据被送到指定的设备上。
  • 模型预测:模型被设置为评估模式,这是进行预测前的重要步骤,它确保了模型中的所有层都按预测模式运行(例如,关闭dropout层)。
  • 输出处理:模型的输出被转移到CPU设备,方便后续处理和可视化。
  • 可视化:使用 matplotlib 和 OpenCV 可视化图像及其边界框,这有助于直观地评估模型的性能和检测结果的准确性。

这个过程是评估深度学习模型在实际任务中表现的常用方法,通过直接观察预测结果,可以更好地理解模型的行为和潜在的改进方向。

配套数据集和源码下载

pytorch实现FasterRCNN目标检测入门案例

相关推荐

最近更新

  1. TCP协议是安全的吗?

    2024-06-08 03:06:02       16 阅读
  2. 阿里云服务器执行yum,一直下载docker-ce-stable失败

    2024-06-08 03:06:02       16 阅读
  3. 【Python教程】压缩PDF文件大小

    2024-06-08 03:06:02       15 阅读
  4. 通过文章id递归查询所有评论(xml)

    2024-06-08 03:06:02       18 阅读

热门阅读

  1. 【Vue】sync修饰符

    2024-06-08 03:06:02       9 阅读
  2. springboot获取当前数据库连接

    2024-06-08 03:06:02       10 阅读
  3. Lambda表达式与函数式接口

    2024-06-08 03:06:02       8 阅读
  4. Nginx实现缓存

    2024-06-08 03:06:02       6 阅读
  5. c# 身份证信息验证

    2024-06-08 03:06:02       9 阅读
  6. oracle常用经典SQL查询

    2024-06-08 03:06:02       9 阅读
  7. 理论学习-自动控制

    2024-06-08 03:06:02       7 阅读
  8. Python笔记 - generator方法

    2024-06-08 03:06:02       12 阅读
  9. Elixir学习笔记——模块和函数

    2024-06-08 03:06:02       10 阅读
  10. Skins

    Skins

    2024-06-08 03:06:02      9 阅读