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_TRAIN 和 DIR_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 重新组织批数据，适合处理包含多个组件（如图像和其标注信息）的数据结构。
• DataLoader：DataLoader 是 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目标检测入门案例