[Python图像识别] 四十八.Pytorch构建Faster-RCNN模型实现小麦目标检测

该系列文章是讲解Python OpenCV图像处理知识，前期主要讲解图像入门、OpenCV基础用法，中期讲解图像处理的各种算法，包括图像锐化算子、图像增强技术、图像分割等，后期结合深度学习研究图像识别、图像分类应用。希望文章对您有所帮助，如果有不足之处，还请海涵~

上一篇文章主要通过Keras深度学习构建CNN模型识别阿拉伯手写文字图像，一篇非常经典的图像分类文字。本文将详细讲解Pytorch构建Faster-RCNN模型实现小麦目标检测，主要参考kaggle大佬和刘兄的模型，推荐大家关注。这是一篇非常经典的图像识别文字，希望您喜欢，且看且珍惜。

第二阶段我们进入了Python图像识别，该部分主要以目标检测、图像识别以及深度学习相关图像分类为主，将会分享近50篇文章，感谢您一如至往的支持。作者也会继续加油的！

• https://maoli.blog.csdn.net/article/details/117688738

同时，该部分知识均为作者查阅资料撰写总结，并且开设成了收费专栏，为小宝赚点奶粉钱，感谢您的抬爱。如果有问题随时私聊我，只望您能从这个系列中学到知识，一起加油。代码下载地址（如果喜欢记得star，一定喔）：

• https://github.com/eastmountyxz/ImageProcessing-Python

图像识别：

• [Python图像识别] 四十五.对象检测案例入门及ImageAI基础用法
• [Python图像识别] 四十六.图像预处理之图像去雾详解（ACE算法和暗通道先验去雾算法）
• [Python图像识别] 四十七.Keras深度学习构建CNN识别阿拉伯手写文字图像
• [Python图像识别] 四十八.Pytorch构建Faster-RCNN模型实现小麦目标检测

图像处理：

• [Python图像处理] 一.图像处理基础知识及OpenCV入门函数
• [Python图像处理] 二.OpenCV+Numpy库读取与修改像素
• [Python图像处理] 三.获取图像属性、兴趣ROI区域及通道处理
• [Python图像处理] 四.图像平滑之均值滤波、方框滤波、高斯滤波及中值滤波
• [Python图像处理] 五.图像融合、加法运算及图像类型转换
• [Python图像处理] 六.图像缩放、图像旋转、图像翻转与图像平移
• [Python图像处理] 七.图像阈值化处理及算法对比
• [Python图像处理] 八.图像腐蚀与图像膨胀
• [Python图像处理] 九.形态学之图像开运算、闭运算、梯度运算
• [Python图像处理] 十.形态学之图像顶帽运算和黑帽运算
• [Python图像处理] 十一.灰度直方图概念及OpenCV绘制直方图
• [Python图像处理] 十二.图像几何变换之图像仿射变换、图像透视变换和图像校正
• [Python图像处理] 十三.基于灰度三维图的图像顶帽运算和黑帽运算
• [Python图像处理] 十四.基于OpenCV和像素处理的图像灰度化处理
• [Python图像处理] 十五.图像的灰度线性变换
• [Python图像处理] 十六.图像的灰度非线性变换之对数变换、伽马变换
• [Python图像处理] 十七.图像锐化与边缘检测之Roberts算子、Prewitt算子、Sobel算子和Laplacian算子
• [Python图像处理] 十八.图像锐化与边缘检测之Scharr算子、Canny算子和LOG算子
• [Python图像处理] 十九.图像分割之基于K-Means聚类的区域分割
• [Python图像处理] 二十.图像量化处理和采样处理及局部马赛克特效
• [Python图像处理] 二十一.图像金字塔之图像向下取样和向上取样
• [Python图像处理] 二十二.Python图像傅里叶变换原理及实现
• [Python图像处理] 二十三.傅里叶变换之高通滤波和低通滤波
• [Python图像处理] 二十四.图像特效处理之毛玻璃、浮雕和油漆特效
• [Python图像处理] 二十五.图像特效处理之素描、怀旧、光照、流年以及滤镜特效
• [Python图像处理] 二十六.图像分类原理及基于KNN、朴素贝叶斯算法的图像分类案例
• [Python图像处理] 二十七.OpenGL入门及绘制基本图形（一）
• [Python图像处理] 二十八.OpenCV快速实现人脸检测及视频中的人脸
• [Python图像处理] 二十九.MoviePy视频编辑库实现抖音短视频剪切合并操作
• [Python图像处理] 三十.图像量化及采样处理万字详细总结（推荐）
• [Python图像处理] 三十一.图像点运算处理两万字详细总结（灰度化处理、阈值化处理）
• [Python图像处理] 三十二.傅里叶变换（图像去噪）与霍夫变换（特征识别）万字详细总结
• [Python图像处理] 三十三.图像各种特效处理及原理万字详解（毛玻璃、浮雕、素描、怀旧、流年、滤镜等）
• [Python图像处理] 三十四.数字图像处理基础与几何图形绘制万字详解（推荐）
• [Python图像处理] 三十五.OpenCV图像处理入门、算数逻辑运算与图像融合（推荐）
• [Python图像处理] 三十六.OpenCV图像几何变换万字详解（平移缩放旋转、镜像仿射透视）
• [Python图像处理] 三十七.OpenCV和Matplotlib绘制直方图万字详解（掩膜直方图、H-S直方图、黑夜白天判断）
• [Python图像处理] 三十八.OpenCV图像增强万字详解（直方图均衡化、局部直方图均衡化、自动色彩均衡化）
• [Python图像处理] 三十九.Python图像分类万字详解（贝叶斯图像分类、KNN图像分类、DNN图像分类）
• [Python图像处理] 四十.全网首发Python图像分割万字详解（阈值分割、边缘分割、纹理分割、分水岭算法、K-Means分割、漫水填充分割、区域定位）
• [Python图像处理] 四十一.Python图像平滑万字详解（均值滤波、方框滤波、高斯滤波、中值滤波、双边滤波）
• [Python图像处理] 四十二.Python图像锐化及边缘检测万字详解（Roberts、Prewitt、Sobel、Laplacian、Canny、LOG）
• [Python图像处理] 四十三.Python图像形态学处理万字详解（腐蚀膨胀、开闭运算、梯度顶帽黑帽运算）
• 万字长文告诉新手如何学习Python图像处理 (上篇完结 四十四)

---

一.Pytorch安装

Pytorch安装需要在官网选择对应的环境，接着按自动生成的安装命令执行。

• 官网：https://pytorch.org/

选择与自己相匹配的版本，这里显示是我安装的选择。

安装代码：

• pip3 install torch torchvision torchaudio
• conda install pytorch torchvision torchaudio cpuonly -c pytorch

同时安装扩展包albumentations。

• pip install albumentations

---

二.数据集描述

1.Kaggle赛题

数据集是来自Kaggle的——全球小麦检测数据，题目是“您能使用图像分析帮助识别小麦吗？”

• https://www.kaggle.com/c/global-wheat-detection

题目介绍：
打开您的食品储藏室，您很可能会找到几种小麦产品。事实上，您的早餐吐司或麦片可能依赖于这种常见的谷物。它作为一种流行的食物和作物，让小麦得到了广泛的研究。为了获得有关全球麦田的大量准确数据，植物科学家使用“小麦头”的图像，检测含有谷物的植物顶部尖峰。这些图像用于估计不同品种小麦的密度和大小。农民在他们的田地做出管理决策时，可以使用这些数据来评估其健康和成熟度。

然而，在室外田间图像中准确检测麦头在视觉上具有挑战性。密密麻麻的小麦植株经常重叠，风会模糊照片。两者都使识别单个头部变得困难。此外，外观因成熟度、颜色、基因类型和头部方向而异。最后，由于小麦在世界范围内种植，因此必须考虑不同的品种、种植密度、模式和田间条件。小麦开发模型需要在不同的生长环境之间进行概括。当前的检测方法涉及一级和二级检测器（Yolo-v3 和 Faster-RCNN），但即使使用大型数据集进行训练，对训练区域的偏差仍然存在。

在全球小麦头数据集是由来自七个国家的九个研究机构主导，包括东京大学等。此后，许多机构都加入了他们追求准确检测小麦头部的行列，包括全球食品安全研究所、DigitAg、Kubota 和 Hiphen。在本次比赛中，您将从小麦植物的室外图像中检测小麦头，包括来自全球的小麦数据集。使用全球数据，您将专注于通用解决方案来估计小麦头的数量和大小。为了更好地衡量未知基因型、环境和观察条件的性能，训练数据集涵盖多个区域。您将使用来自欧洲（法国、英国、瑞士）和北美（加拿大）的 3,000 多张图像。测试数据包括来自澳大利亚、日本和中国的约 1,000 张图像。

小麦是全球的主食，这就是为什么这种竞争必须考虑到不同的生长条件。为小麦表型开发的模型需要能够在环境之间进行概括。如果成功，研究人员可以准确估计不同品种小麦头的密度和大小。通过改进的检测，农民可以更好地评估他们的作物，最终将谷物、烤面包和其他喜爱的菜肴带到您的餐桌上。有关数据采集和过程的更多详细信息，请访问：

• https://arxiv.org/abs/2005.02162

---

2.数据集介绍

我们应该期望数据格式是什么？
数据是麦田的图像，每个识别的麦头都有边界框，并非所有图像都包含小麦头/边界框。这些图像被记录在世界各地的许多地方。

• CSV 数据很简单：图像 ID 与给定图像的文件名相匹配，并且包含图像的宽度和高度以及边界框（见下文）。train.csv每个边界框都有一行，并非所有图像都有边界框。大多数测试集图像是隐藏的，包含一小部分测试图像供您在编写代码时使用。

我们在预测什么？
您正在尝试预测图像中每个小麦头周围的边界框。如果没有小麦头，则必须预测没有边界框。

数据集包含四个文件

• train.csv - 训练数据
• sample_submission.csv - 格式正确的示例提交文件
• train.zip - 训练图像
• test.zip - 测试图像

数据集如下图所示：

文件夹中包含小麦图像，名称是其ID。

train.csv中对应五列结果，分别是：

• image_id - 唯一的图像 ID
• width - 图像的宽度
• height - 图像的高度
• bbox - 一个边界框，格式为 [xmin, ymin, width, height] 的 Python 样式列表
• source - 图像对应的类别

训练集中各小麦类型分布如下图所示：

整个小麦预测的大致流程如下图所示：

模型评估参数如下所示，推荐大家阅读kaggle官网介绍。

提交格式需要以空格分隔的一组边界框。例如：

• ce4833752, 0.5 0 0 100 100
  表示图像ce4833752有一个边界框，aconfidence为 0.5，在x== 0 和y== 0，awidth和height为 100。

该文件应包含标题并具有以下格式，您提交的每一行都应包含给定图像的所有边界框。

image_id,PredictionString
ce4833752,1.0 0 0 50 50
adcfa13da,1.0 0 0 50 50
6ca7b2650,
1da9078c1,0.3 0 0 50 50 0.5 10 10 30 30
7640b4963,0.5 0 0 50 50

---

三.代码实现

下面我们参考Kaggle Peter老师的代码，来复现Faster-RCNN模型。

• https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train

模型的框架如下图所示，相信大家都比较熟悉，也推荐大家使用并深入了解背后的原理。

1.读取小麦数据

读取小麦数据集的代码如下：

# -*- coding: utf-8 -*-
"""
Created on Fri Oct 29 13:42:38 2021
@author: xiuzhang
"""
import os
import re
import cv2
import pandas as pd
import numpy as np
from PIL import Image
import albumentations as A
from matplotlib import pyplot as plt
from albumentations.pytorch.transforms import ToTensorV2

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import SequentialSampler

from dataset import WheatDataset

#-----------------------------------------------------------------------------
#第一步 函数定义
#----------------------------------------------------------------------------
#提取box的四个坐标
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

#训练图像增强 Albumentations
def get_train_transform():
    return A.Compose([
        A.Flip(0.5),
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

#验证图像增强
def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

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

#-----------------------------------------------------------------------------
#第二步 定义变量并读取数据
#-----------------------------------------------------------------------------
DIR_INPUT = 'data'
DIR_TRAIN = f'{
       DIR_INPUT}/train'
DIR_TEST = f'{
       DIR_INPUT}/test'
train_df = pd.read_csv(f'{
       DIR_INPUT}/train.csv')
print(train_df.shape)

train_df['x'] = -1
train_df['y'] = -1
train_df['w'] = -1
train_df['h'] = -1

#读取box四个坐标
train_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))
train_df.drop(columns=['bbox'], inplace=True)
train_df['x'] = train_df['x'].astype(np.float)
train_df['y'] = train_df['y'].astype(np.float)
train_df['w'] = train_df['w'].astype(np.float)
train_df['h'] = train_df['h'].astype(np.float)

#获取图像id
image_ids = train_df['image_id'].unique()
valid_ids = image_ids[-665:]
train_ids = image_ids[:-665]
valid_df = train_df[train_df['image_id'].isin(valid_ids)]
train_df = train_df[train_df['image_id'].isin(train_ids)]
print(valid_df.shape, train_df.shape)
print(train_df.head())

显示结果如下图所示，分别获取图像id和数据，并划分为train（训练）和valid（验证）。

其中，dataset.py文件代码如下：

• 获取图像id
• 获取图像像素值并归一化处理
• 获取图像对应的边界（x | y | w | h）

# -*- coding: utf-8 -*-
"""
Created on Fri Oct 29 13:42:38 2021
@author: xiuzhang
"""
import numpy as np
import cv2
import torch
from torch.utils.data import Dataset


class WheatDataset(Dataset):

    def __init__(self, dataframe, image_dir, transforms=None):
        super().__init__()

        self.image_ids = dataframe['image_id'].unique()
        self.df = dataframe
        self.image_dir = image_dir
        self.transforms = transforms

    def __getitem__(self, index: int):
        image_id = self.image_ids[index]
        records = self.df[self.df['image_id'] == image_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]
        boxes[:, 3] = boxes[:, 1] + boxes[:, 3]

        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        area = torch.as_tensor(area, dtype=torch.float32)

        # there is only one class
        labels = torch.ones((records.shape[0],), dtype=torch.int64)

        # suppose all instances are not crowd
        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)

        target = {
     }
        target['boxes'] = boxes
        target['labels'] = labels
        # target['masks'] = None
        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]

---

2.可视化展示

接下来我们对小麦图像进行简单的可视化操作，代码如下：

# -*- coding: utf-8 -*-
"""
Created on Fri Oct 29 13:42:38 2021
@author: xiuzhang
"""
import os
import re
import cv2
import pandas as pd
import numpy as np
from PIL import Image
import albumentations as A
from matplotlib import pyplot as plt
from albumentations.pytorch.transforms import ToTensorV2

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import SequentialSampler

from dataset import WheatDataset

#-----------------------------------------------------------------------------
#第一步 函数定义
#----------------------------------------------------------------------------
#提取box的四个坐标
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

#训练图像增强 Albumentations
def get_train_transform():
    return A.Compose([
        A.Flip(0.5),
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

#验证图像增强
def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

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

#-----------------------------------------------------------------------------
#第二步 定义变量并读取数据
#-----------------------------------------------------------------------------
DIR_INPUT = 'data'
DIR_TRAIN = f'{
       DIR_INPUT}/train'
DIR_TEST = f'{
       DIR_INPUT}/test'
train_df = pd.read_csv(f'{
       DIR_INPUT}/train.csv')
print(train_df.shape)

train_df['x'] = -1
train_df['y'] = -1
train_df['w'] = -1
train_df['h'] = -1

#读取box四个坐标
train_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))
train_df.drop(columns=['bbox'], inplace=True)
train_df['x'] = train_df['x'].astype(np.float)
train_df['y'] = train_df['y'].astype(np.float)
train_df['w'] = train_df['w'].astype(np.float)
train_df['h'] = train_df['h'].astype(np.float)

#获取图像id
image_ids = train_df['image_id'].unique()
valid_ids = image_ids[-665:]
train_ids = image_ids[:-665]
valid_df = train_df[train_df['image_id'].isin(valid_ids)]
train_df = train_df[train_df['image_id'].isin(train_ids)]
print(valid_df.shape, train_df.shape)
print(train_df.head())

#-----------------------------------------------------------------------------
#第三步 加载数据
#-----------------------------------------------------------------------------
train_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())
valid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())

train_data_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=False,
    num_workers=0,
    collate_fn=collate_fn
)

valid_data_loader = DataLoader(
    valid_dataset,
    batch_size=2,
    shuffle=False,
    num_workers=0,
    collate_fn=collate_fn
)

#-----------------------------------------------------------------------------
#第四步 数据可视化
#-----------------------------------------------------------------------------
#提取训练数据和类别
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
images, targets, image_ids = next(iter(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]
boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)
sample = images[0].permute(1, 2, 0).cpu().numpy()

fig, ax = plt.subplots(1, 1, figsize=(10, 8))

#绘制小麦目标识别box
for box in boxes:
    cv2.rectangle(sample,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (255, 0, 0), 3)

    ax.text(box[0], 
            box[1] - 2, 
            '{:s}'.format('wheat'), 
            bbox=dict(facecolor='blue', alpha=0.5),
            fontsize=12, 
            color='white')

ax.set_axis_off()
ax.imshow(sample)
plt.show()

输出结果如下图所示，按照train.csv定义好的边界我们绘制了wheat红色框，将小麦标记。最终的测试集，我们希望能自动预测小麦的边界，从而有效识别小麦区域和数量。

警告：

• Clipping input data to the valid range for imshow with RGB data ([0…1] for floats or [0…255] for integers).

---

3.构建Faster-RCNN模型

接下来构造Faster-RCNN模型，这是目标检测的经典模型。其核心代码如下所示：

• model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

#-----------------------------------------------------------------------------
#第五步 模型构建
#-----------------------------------------------------------------------------
num_classes = 2  #1 class (wheat) + background
lr_scheduler = None
num_epochs = 1
itr = 1

class Averager:
    def __init__(self):
        self.current_total = 0.0
        self.iterations = 0.0

    def send(self, value):
        self.current_total += value
        self.iterations += 1

    @property
    def value(self):
        if self.iterations == 0:
            return 0
        else:
            return 1.0 * self.current_total / self.iterations

    def reset(self):
        self.current_total = 0.0
        self.iterations = 0.0

#load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

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

#replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

#参数设置
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
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)

loss_hist = Averager()
print("Start training....")

# 迭代训练
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]
        for t in targets:
            t['boxes'] = t['boxes'].float()
        
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()
        loss_hist.send(loss_value)
        print("loss is :",loss_value)

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        if itr % 50 == 0:
            print(f"Iteration #{
       itr}/{
       len(train_data_loader)} loss: {
       loss_value}")
        itr += 1

    #更新学习率
    if lr_scheduler is not None:
        lr_scheduler.step()
    print(f"Epoch #{
       epoch} loss: {
       loss_hist.value}")

torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')
print("Next Test....")

运行过程如下图所示：

Downloading: "https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth" 
to C:\Users\xxx/.cache\torch\hub\checkpoints\fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
100%|██████████| 160M/160M [04:06<00:00, 679KB/s]  

---

4.模型预测

增加模型预测的最终完整代码如下：

• 第一步：函数定义
• 第二步：定义变量并读取数据
• 第三步：加载数据
• 第四步：数据可视化
• 第五步：Faster-RCNN模型构建
• 第六步：模型测试

# -*- coding: utf-8 -*-
"""
Created on Fri Oct 29 13:42:38 2021
@author: xiuzhang
"""
import os
import re
import cv2
import pandas as pd
import numpy as np
from PIL import Image
import albumentations as A
from matplotlib import pyplot as plt
from albumentations.pytorch.transforms import ToTensorV2

import torch
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import SequentialSampler

from dataset import WheatDataset

#-----------------------------------------------------------------------------
#第一步 函数定义
#----------------------------------------------------------------------------
#提取box的四个坐标
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

#训练图像增强 Albumentations
def get_train_transform():
    return A.Compose([
        A.Flip(0.5),
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

#验证图像增强
def get_valid_transform():
    return A.Compose([
        ToTensorV2(p=1.0)
    ], bbox_params={
     'format': 'pascal_voc', 'label_fields': ['labels']})

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

#-----------------------------------------------------------------------------
#第二步 定义变量并读取数据
#-----------------------------------------------------------------------------
DIR_INPUT = 'data'
DIR_TRAIN = f'{
       DIR_INPUT}/train'
DIR_TEST = f'{
       DIR_INPUT}/test'
train_df = pd.read_csv(f'{
       DIR_INPUT}/train.csv')
print(train_df.shape)

train_df['x'] = -1
train_df['y'] = -1
train_df['w'] = -1
train_df['h'] = -1

#读取box四个坐标
train_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))
train_df.drop(columns=['bbox'], inplace=True)
train_df['x'] = train_df['x'].astype(np.float)
train_df['y'] = train_df['y'].astype(np.float)
train_df['w'] = train_df['w'].astype(np.float)
train_df['h'] = train_df['h'].astype(np.float)

#获取图像id
image_ids = train_df['image_id'].unique()
valid_ids = image_ids[-665:]
train_ids = image_ids[:-665]
valid_df = train_df[train_df['image_id'].isin(valid_ids)]
train_df = train_df[train_df['image_id'].isin(train_ids)]
print(valid_df.shape, train_df.shape)
print(train_df.head())

#-----------------------------------------------------------------------------
#第三步 加载数据
#-----------------------------------------------------------------------------
train_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())
valid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())

train_data_loader = DataLoader(
    train_dataset,
    batch_size=2,
    shuffle=False,
    num_workers=0,
    collate_fn=collate_fn
)

valid_data_loader = DataLoader(
    valid_dataset,
    batch_size=2,
    shuffle=False,
    num_workers=0,
    collate_fn=collate_fn
)

#-----------------------------------------------------------------------------
#第四步 数据可视化
#-----------------------------------------------------------------------------
#提取训练数据和类别
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
images, targets, image_ids = next(iter(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]
boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)
sample = images[0].permute(1, 2, 0).cpu().numpy()

fig, ax = plt.subplots(1, 1, figsize=(10, 8))

#绘制小麦目标识别box
for box in boxes:
    cv2.rectangle(sample,
                  (box[0], box[1]),
                  (box[2], box[3]),
                  (255, 0, 0), 3)

    ax.text(box[0], 
            box[1] - 2, 
            '{:s}'.format('wheat'), 
            bbox=dict(facecolor='blue', alpha=0.5),
            fontsize=12, 
            color='white')

ax.set_axis_off()
ax.imshow(sample)
plt.show()

#-----------------------------------------------------------------------------
#第五步 模型构建
#-----------------------------------------------------------------------------
num_classes = 2  #1 class (wheat) + background
lr_scheduler = None
num_epochs = 1
itr = 1

class Averager:
    def __init__(self):
        self.current_total = 0.0
        self.iterations = 0.0

    def send(self, value):
        self.current_total += value
        self.iterations += 1

    @property
    def value(self):
        if self.iterations == 0:
            return 0
        else:
            return 1.0 * self.current_total / self.iterations

    def reset(self):
        self.current_total = 0.0
        self.iterations = 0.0

#load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

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

#replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

#参数设置
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
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)

loss_hist = Averager()
print("Start training....")

# 迭代训练
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]
        for t in targets:
            t['boxes'] = t['boxes'].float()
        
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()
        loss_hist.send(loss_value)
        #print("loss is :",loss_value)

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        if itr % 50 == 0:
            print(f"Iteration #{
       itr}/{
       len(train_data_loader)} loss: {
       loss_value}")
        itr += 1

    #更新学习率
    if lr_scheduler is not None:
        lr_scheduler.step()
    print(f"Epoch #{
       epoch} loss: {
       loss_hist.value}")

torch.save(model.state_dict(), 'fasterrcnn_resnet50_fpn.pth')
print("Next Test....")

#-----------------------------------------------------------------------------
#第六步 模型测试
#-----------------------------------------------------------------------------
images, targets, image_ids = next(iter(valid_data_loader))
images = list(img.to(device) for img in images)
targets = [{
     k: v.to(device) for k, v in t.items()} for t in targets]
boxes = targets[0]['boxes'].cpu().numpy().astype(np.int32)
sample = images[0].permute(1, 2, 0).cpu().numpy()

model.eval()
cpu_device = torch.device("cpu")

outputs = model(images)
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)
plt.show()  

模型运行结果如下图所示，可以看到迭代的loss，推荐大家用好的目标检测环境实验。

同时对测试集或验证集的图像进行识别，如下图所示：

---

四.总结

写到这里，这篇文章就介绍结束了，希望对您有所帮助。详细的对比实验和算法评估还请读者自行完成，后续作者的文章也会深入介绍。

• 一.Pytorch安装
• 二.数据集描述
  1.Kaggle赛题
  2.数据集介绍
• 三.代码实现
  1.读取小麦数据
  2.可视化展示
  3.构建Faster-RCNN模型
  4.模型预测

代码和数据集下载地址:

• https://github.com/eastmountyxz/ImageProcessing-Python

《立冬–小珞情》
初冬已至，泛黄的银杏叶随着寒风飘落，观山大道旁不时传来阵阵残香，冷风也裹紧了路人的衣裳。随着时光的流淌，我渡过了生命中的第一年，虽然还不能用言语表达，但我早已熟悉了这美妙的世界，感受到了家人对我的疼爱，当然也偶尔会经历一些烦恼，比这两天的感冒。每当我难受的时候，我就会嘶吼，或烦躁，还好，妈妈和婆婆总能第一时间将我抱起，在怀抱中摇摆着安抚我稳定。摇着摇着，看到妈妈那双疼爱的眼睛，我总会立时扬起嘴角，那是一种只有在妈妈怀抱中才会扬起的傻笑，接着进入甜美的梦境中。虽然我还小，但似乎也能感受到妈妈对这个微笑的喜欢，恰是立冬的第一杯奶茶，暖暖地流进她的心底。

“叮咚叮…”，醒来的我听到了电话里传来了远方的微信视频。我迫不及待地抢走了手机，但又不知道如何接通，此时的妈妈扶着我按下了绿色接听键，看到那似熟非熟的眼镜娃娃胡子脸，我叫了一声“粑粑”（就是这发音）。亲情，那一朵永远微笑着的蔷薇，无论过了多久都让人无法忘却，散发着芬芳。电话里，听爸爸说，他朋友圈里的北方迎来了初雪，EDG夺得了冠军，乔木和街道昨夜已悄然换上了雪白的冬装，让我们记得保暖。妈妈回复到，贵阳的寒风这两天格外的凛冽，小珞珞感冒快好了，也让爸爸记得多穿衣服。听着他俩唠家常，我在旁边不时地翻滚嬉闹，似乎要证明我才是家里最重要的一员。“妈做的酸汤鱼好了，我们准备吃饭去了，你也早点吃饭，珞珞别担心，保重身体 ”，随着家里那口老锅的热气升腾，就此结束通话。

或许，一岁的我还不知道这意味着什么，但我知道在妈妈怀抱中看着视频里远方的爸爸就很快乐，像极了我吃西瓜的甘甜。大一点，我的小学作文里，可能会写上一句：“这就是家的味道，也是人们在生命里追踪的情思，恰似这灶火间酸汤鱼燃烧的人间喜怒酸甜。正是因为这种味道，妈妈才会爱上爸爸，爸爸也才会追求妈妈，他们都才会爱最可爱的我”。

或许，我童真的世界记不了这些，但爸爸和妈妈会永远记住我的每一天。爱，在当下，爱，在顷刻间。小珞珞祝大家立冬快乐。

(By:Eastmount 2021-11-08 夜于武汉 http://blog.csdn.net/eastmount/ )

---

感谢几位大佬的分享，参考文献如下：

• https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train
• https://blog.csdn.net/qq_39071739/article/details/108193935
• https://maoli.blog.csdn.net/article/details/118575041
• https://maoli.blog.csdn.net/article/details/119571695