目录
依赖环境
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 可视化图像及其边界框，这有助于直观地评估模型的性能和检测结果的准确性。
这个过程是评估深度学习模型在实际任务中表现的常用方法，通过直接观察预测结果，可以更好地理解模型的行为和潜在的改进方向。