• Pytorch R-CNN目标检测-汽车car

    概述

    目标检测(Object Detection)就是一种基于目标几何和统计特征的图像分割,它将目标的分割和识别合二为一,通俗点说就是给定一张图片要精确的定位到物体所在位置,并完成对物体类别的识别。其准确性和实时性是整个系统的一项重要能力。

    R-CNN的全称是Region-CNN(区域卷积神经网络),是第一个成功将深度学习应用到目标检测上的算法。R-CNN基于卷积神经网络(CNN),线性回归,和支持向量机(SVM)等算法,实现目标检测技术。

    以下的代码和项目工程是引用他人的,此文只对其做一个简单的流程梳理。

    这里先贴出工具脚本util.py的代码如下

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/2/29 下午7:31
    5. @file: util.py
    6. @author: zj
    7. @description: 
    8. """
    9.  
    10. import os
    11. import numpy as np
    12. import xmltodict
    13. import torch
    14. import matplotlib.pyplot as plt
    15.  
    16.  
    17. def check_dir(data_dir):
    18.     if not os.path.exists(data_dir):
    19.         os.mkdir(data_dir)
    20.  
    21.  
    22. def parse_car_csv(csv_dir):
    23.     csv_path = os.path.join(csv_dir, 'car.csv')
    24.     samples = np.loadtxt(csv_path, dtype='str')
    25.     return samples
    26.  
    27.  
    28. def parse_xml(xml_path):
    29.     """
    30.     解析xml文件，返回标注边界框坐标
    31.     """
    32.     # print(xml_path)
    33.     with open(xml_path, 'rb') as f:
    34.         xml_dict = xmltodict.parse(f)
    35.         # print(xml_dict)
    36.  
    37.         bndboxs = list()
    38.         objects = xml_dict['annotation']['object']
    39.         if isinstance(objects, list):
    40.             for obj in objects:
    41.                 obj_name = obj['name']
    42.                 difficult = int(obj['difficult'])
    43.                 if 'car'.__eq__(obj_name) and difficult != 1:
    44.                     bndbox = obj['bndbox']
    45.                     bndboxs.append((int(bndbox['xmin']), int(bndbox['ymin']), int(bndbox['xmax']), int(bndbox['ymax'])))
    46.         elif isinstance(objects, dict):
    47.             obj_name = objects['name']
    48.             difficult = int(objects['difficult'])
    49.             if 'car'.__eq__(obj_name) and difficult != 1:
    50.                 bndbox = objects['bndbox']
    51.                 bndboxs.append((int(bndbox['xmin']), int(bndbox['ymin']), int(bndbox['xmax']), int(bndbox['ymax'])))
    52.         else:
    53.             pass
    54.  
    55.         return np.array(bndboxs)
    56.  
    57.  
    58. def iou(pred_box, target_box):
    59.     """
    60.     计算候选建议和标注边界框的IoU
    61.     :param pred_box: 大小为[4]
    62.     :param target_box: 大小为[N, 4]
    63.     :return: [N]
    64.     """
    65.     if len(target_box.shape) == 1:
    66.         target_box = target_box[np.newaxis, :]
    67.  
    68.     xA = np.maximum(pred_box[0], target_box[:, 0])
    69.     yA = np.maximum(pred_box[1], target_box[:, 1])
    70.     xB = np.minimum(pred_box[2], target_box[:, 2])
    71.     yB = np.minimum(pred_box[3], target_box[:, 3])
    72.     # 计算交集面积
    73.     intersection = np.maximum(0.0, xB - xA) * np.maximum(0.0, yB - yA)
    74.     # 计算两个边界框面积
    75.     boxAArea = (pred_box[2] - pred_box[0]) * (pred_box[3] - pred_box[1])
    76.     boxBArea = (target_box[:, 2] - target_box[:, 0]) * (target_box[:, 3] - target_box[:, 1])
    77.  
    78.     scores = intersection / (boxAArea + boxBArea - intersection)
    79.     return scores
    80.  
    81.  
    82. def compute_ious(rects, bndboxs):
    83.     iou_list = list()
    84.     for rect in rects:
    85.         scores = iou(rect, bndboxs)
    86.         iou_list.append(max(scores))
    87.     return iou_list
    88.  
    89.  
    90. def save_model(model, model_save_path):
    91.     # 保存最好的模型参数
    92.     check_dir('./models')
    93.     torch.save(model.state_dict(), model_save_path)
    94.  
    95.  
    96. def plot_loss(loss_list):
    97.     x = list(range(len(loss_list)))
    98.     fg = plt.figure()
    99.  
    100.     plt.plot(x, loss_list)
    101.     plt.title('loss')
    102.     plt.savefig('./loss.png')

    数据集准备

    数据集下载

    运行项目中的pascal_voc.py脚本,这个脚本是下载数据集。

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/2/29 下午2:51
    5. @file: pascal_voc.py
    6. @author: zj
    7. @description: 加载PASCAL VOC 2007数据集
    8. """
    9.  
    10. import cv2
    11. import numpy as np
    12. from torchvision.datasets import VOCDetection
    13.  
    14. if __name__ == '__main__':
    15.     """
    16.     下载PASCAL VOC数据集
    17.     """
    18.     dataset = VOCDetection('../../data', year='2007', image_set='trainval', download=True)
    19.  
    20.     # img, target = dataset.__getitem__(1000)
    21.     # img = np.array(img)
    22.     # 
    23.     # print(target)
    24.     # print(img.shape)
    25.     # 
    26.     # cv2.imshow('img', img)
    27.     # cv2.waitKey(0)

    从数据集中提取出car相关的数据

    由于本文只针对汽车car进行目标检测,所以只需要car相关的数据。

    执行pascal_voc_car.py脚本,脚本依次做了以下事:

    ①读取'../../data/VOCdevkit/VOC2007/ImageSets/Main/car_train.txt'文件和'../../data/VOCdevkit/VOC2007/ImageSets/Main/car_val.txt'文件

    car_train.txt和car_val.txt文件的内容格式如下

    ②然后将以上文件内容分别保存到'../../data/voc_car/train/car.csv'和'../../data/voc_car/val/car.csv'中

    car.csv的内容格式如下

    ③最后根据筛选出来的car的相关数据,从'../../data/VOCdevkit/VOC2007/Annotations/'中复制相关.xml文件到'../../data/voc_car/train/Annotations/'和'../../data/voc_car/val/Annotations/',以及从'../../data/VOCdevkit/VOC2007/JPEGImages/'中复制相关.jpg文件到'../../data/voc_car/train/JPEGImages/'和'../../data/voc_car/val/JPEGImages/'

    以下是pascal_voc_car.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/2/29 下午2:43
    5. @file: pascal_voc_car.py
    6. @author: zj
    7. @description: 从PASCAL VOC 2007数据集中抽取类别Car。保留1/10的数目
    8. """
    9.  
    10. import os
    11. import shutil
    12. import random
    13. import numpy as np
    14. import xmltodict
    15. from utils.util import check_dir
    16.  
    17. suffix_xml = '.xml'
    18. suffix_jpeg = '.jpg'
    19.  
    20. car_train_path = '../../data/VOCdevkit/VOC2007/ImageSets/Main/car_train.txt'
    21. car_val_path = '../../data/VOCdevkit/VOC2007/ImageSets/Main/car_val.txt'
    22.  
    23. voc_annotation_dir = '../../data/VOCdevkit/VOC2007/Annotations/'
    24. voc_jpeg_dir = '../../data/VOCdevkit/VOC2007/JPEGImages/'
    25.  
    26. car_root_dir = '../../data/voc_car/'
    27.  
    28.  
    29. def parse_train_val(data_path):
    30.     """
    31.     提取指定类别图像
    32.     """
    33.     samples = []
    34.  
    35.     with open(data_path, 'r') as file:
    36.         lines = file.readlines()
    37.         for line in lines:
    38.             res = line.strip().split(' ')
    39.             if len(res) == 3 and int(res[2]) == 1:
    40.                 samples.append(res[0])
    41.  
    42.     return np.array(samples)
    43.  
    44.  
    45. def sample_train_val(samples):
    46.     """
    47.     随机采样样本，减少数据集个数（留下1/10）
    48.     """
    49.     for name in ['train', 'val']:
    50.         dataset = samples[name]
    51.         length = len(dataset)
    52.  
    53.         random_samples = random.sample(range(length), int(length / 10))
    54.         # print(random_samples)
    55.         new_dataset = dataset[random_samples]
    56.         samples[name] = new_dataset
    57.  
    58.     return samples
    59.  
    60. def save_car(car_samples, data_root_dir, data_annotation_dir, data_jpeg_dir):
    61.     """
    62.     保存类别Car的样本图片和标注文件
    63.     """
    64.     for sample_name in car_samples:
    65.         src_annotation_path = os.path.join(voc_annotation_dir, sample_name + suffix_xml)
    66.         dst_annotation_path = os.path.join(data_annotation_dir, sample_name + suffix_xml)
    67.         shutil.copyfile(src_annotation_path, dst_annotation_path)
    68.  
    69.         src_jpeg_path = os.path.join(voc_jpeg_dir, sample_name + suffix_jpeg)
    70.         dst_jpeg_path = os.path.join(data_jpeg_dir, sample_name + suffix_jpeg)
    71.         shutil.copyfile(src_jpeg_path, dst_jpeg_path)
    72.  
    73.     csv_path = os.path.join(data_root_dir, 'car.csv')
    74.     np.savetxt(csv_path, np.array(car_samples), fmt='%s')
    75.  
    76.  
    77. if __name__ == '__main__':
    78.     samples = {'train': parse_train_val(car_train_path), 'val': parse_train_val(car_val_path)}
    79.     print(samples)
    80.     # samples = sample_train_val(samples)
    81.     # print(samples)
    82.  
    83.     check_dir(car_root_dir)
    84.     for name in ['train', 'val']:
    85.         data_root_dir = os.path.join(car_root_dir, name)
    86.         data_annotation_dir = os.path.join(data_root_dir, 'Annotations')
    87.         data_jpeg_dir = os.path.join(data_root_dir, 'JPEGImages')
    88.  
    89.         check_dir(data_root_dir)
    90.         check_dir(data_annotation_dir)
    91.         check_dir(data_jpeg_dir)
    92.         save_car(samples[name], data_root_dir, data_annotation_dir, data_jpeg_dir)
    93.  
    94.     print('done')

    卷积神经网络微调模型

    准备微调数据正负样本集

    执行create_finetune_data.py脚本,这个脚本主要做了以下事

    ①把'../../data/voc_car/train/JPEGImages/'和'../../data/voc_car/val/JPEGImages/'中的.jpg文件复制到'../../data/finetune_car/train/JPEGImages/'和'../../data/finetune_car/val/JPEGImages/',然后又把'../../data/voc_car/train/car.csv'和'../../data/voc_car/val/car.csv'分别复制到'../../data/finetune_car/train/car.csv'和'../../data/finetune_car/val/car.csv'

    ②根据'../../data/finetune_car/train/car.csv'和'../../data/finetune_car/val/car.csv'文件内容分别读取'../../data/finetune_car/train/JPEGImages/'和'../../data/finetune_car/val/JPEGImages/'中的图片,并传入parse_annotation_jpeg方法

    ③parse_annotation_jpeg方法中,先获取候选框rects,然后从.xml文件中获取标注框bndboxs,接着计算候选框和标注框的IoU得到iou_list,遍历iou_list,选出IoU≥0.5的作为正样本,0

    其文件内容格式如下

    以下是create_finetune_data.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/2/29 下午7:22
    5. @file: create_finetune_data.py
    6. @author: zj
    7. @description: 创建微调数据集
    8. """
    9.  
    10. import time
    11. import shutil
    12. import numpy as np
    13. import cv2
    14. import os
    15. import selectivesearch
    16. from utils.util import check_dir
    17. from utils.util import parse_car_csv
    18. from utils.util import parse_xml
    19. from utils.util import compute_ious
    20.  
    21.  
    22. # train
    23. # positive num: 66517
    24. # negatie num: 464340
    25. # val
    26. # positive num: 64712
    27. # negative num: 415134
    28.  
    29.  
    30. def parse_annotation_jpeg(annotation_path, jpeg_path, gs):
    31.     """
    32.     获取正负样本（注：忽略属性difficult为True的标注边界框）
    33.     正样本：候选建议与标注边界框IoU大于等于0.5
    34.     负样本：IoU大于0,小于0.5。为了进一步限制负样本数目，其大小必须大于标注框的1/5
    35.     """
    36.     img = cv2.imread(jpeg_path)
    37.  
    38.     selectivesearch.config(gs, img, strategy='q')
    39.     # 计算候选建议
    40.     rects = selectivesearch.get_rects(gs)
    41.     # 获取标注边界框
    42.     bndboxs = parse_xml(annotation_path)
    43.  
    44.     # 标注框大小
    45.     maximum_bndbox_size = 0
    46.     for bndbox in bndboxs:
    47.         xmin, ymin, xmax, ymax = bndbox
    48.         bndbox_size = (ymax - ymin) * (xmax - xmin)
    49.         if bndbox_size > maximum_bndbox_size:
    50.             maximum_bndbox_size = bndbox_size
    51.  
    52.     # 获取候选建议和标注边界框的IoU
    53.     iou_list = compute_ious(rects, bndboxs)
    54.  
    55.     positive_list = list()
    56.     negative_list = list()
    57.     for i in range(len(iou_list)):
    58.         xmin, ymin, xmax, ymax = rects[i]
    59.         rect_size = (ymax - ymin) * (xmax - xmin)
    60.  
    61.         iou_score = iou_list[i]
    62.         if iou_list[i] >= 0.5:
    63.             # 正样本
    64.             positive_list.append(rects[i])
    65.         if 0 < iou_list[i] < 0.5 and rect_size > maximum_bndbox_size / 5.0:
    66.             # 负样本
    67.             negative_list.append(rects[i])
    68.         else:
    69.             pass
    70.  
    71.     return positive_list, negative_list
    72.  
    73.  
    74. if __name__ == '__main__':
    75.     car_root_dir = '../../data/voc_car/'
    76.     finetune_root_dir = '../../data/finetune_car/'
    77.     check_dir(finetune_root_dir)
    78.  
    79.     gs = selectivesearch.get_selective_search()
    80.     for name in ['train', 'val']:
    81.         src_root_dir = os.path.join(car_root_dir, name)
    82.         src_annotation_dir = os.path.join(src_root_dir, 'Annotations')
    83.         src_jpeg_dir = os.path.join(src_root_dir, 'JPEGImages')
    84.  
    85.         dst_root_dir = os.path.join(finetune_root_dir, name)
    86.         dst_annotation_dir = os.path.join(dst_root_dir, 'Annotations')
    87.         dst_jpeg_dir = os.path.join(dst_root_dir, 'JPEGImages')
    88.         check_dir(dst_root_dir)
    89.         check_dir(dst_annotation_dir)
    90.         check_dir(dst_jpeg_dir)
    91.  
    92.         total_num_positive = 0
    93.         total_num_negative = 0
    94.  
    95.         samples = parse_car_csv(src_root_dir)
    96.         # 复制csv文件
    97.         src_csv_path = os.path.join(src_root_dir, 'car.csv')
    98.         dst_csv_path = os.path.join(dst_root_dir, 'car.csv')
    99.         shutil.copyfile(src_csv_path, dst_csv_path)
    100.         for sample_name in samples:
    101.             since = time.time()
    102.  
    103.             src_annotation_path = os.path.join(src_annotation_dir, sample_name + '.xml')
    104.             src_jpeg_path = os.path.join(src_jpeg_dir, sample_name + '.jpg')
    105.             # 获取正负样本
    106.             positive_list, negative_list = parse_annotation_jpeg(src_annotation_path, src_jpeg_path, gs)
    107.             total_num_positive += len(positive_list)
    108.             total_num_negative += len(negative_list)
    109.  
    110.             dst_annotation_positive_path = os.path.join(dst_annotation_dir, sample_name + '_1' + '.csv')
    111.             dst_annotation_negative_path = os.path.join(dst_annotation_dir, sample_name + '_0' + '.csv')
    112.             dst_jpeg_path = os.path.join(dst_jpeg_dir, sample_name + '.jpg')
    113.             # 保存图片
    114.             shutil.copyfile(src_jpeg_path, dst_jpeg_path)
    115.             # 保存正负样本标注
    116.             np.savetxt(dst_annotation_positive_path, np.array(positive_list), fmt='%d', delimiter=' ')
    117.             np.savetxt(dst_annotation_negative_path, np.array(negative_list), fmt='%d', delimiter=' ')
    118.  
    119.             time_elapsed = time.time() - since
    120.             print('parse {}.png in {:.0f}m {:.0f}s'.format(sample_name, time_elapsed // 60, time_elapsed % 60))
    121.         print('%s positive num: %d' % (name, total_num_positive))
    122.         print('%s negative num: %d' % (name, total_num_negative))
    123.     print('done')

    自定义微调数据集类

    custom_finetune_dataset.py,该脚本不用主动执行,在训练微调模型的时候,自然会调用到,以下只说这个脚本做了什么事

    ①CustomFinetuneDataset类继承自Dataset

    ②__init__时读取'../../data/finetune_car/train/JPEGImages/'或'../../data/finetune_car/val/JPEGImages/'文件夹中的图片,读取'../../data/finetune_car/train/Annotations/'或'../../data/finetune_car/val/Annotations/'中的正负样本集,记录正样本总数self.total_positive_num,负样本总数self.total_negative_num,正样本候选框positive_rects,负样本候选框negative_rects

    ③__getitem__方法传入index,如果index小于正样本总数self.total_positive_num,那么返回对应正样本的图片和标签(1),否则返回对应负样本的图片和标签(0)。

    以下是custom_finetune_dataset.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/3 下午7:06
    5. @file: custom_finetune_dataset.py
    6. @author: zj
    7. @description: 自定义微调数据类
    8. """
    9.  
    10. import numpy  as np
    11. import os
    12. import cv2
    13. from PIL import Image
    14. from torch.utils.data import Dataset
    15. from torch.utils.data import DataLoader
    16. import torchvision.transforms as transforms
    17.  
    18. from utils.util import parse_car_csv
    19.  
    20.  
    21. class CustomFinetuneDataset(Dataset):
    22.  
    23.     def __init__(self, root_dir, transform=None):
    24.         samples = parse_car_csv(root_dir)
    25.  
    26.         jpeg_images = [cv2.imread(os.path.join(root_dir, 'JPEGImages', sample_name + ".jpg"))
    27.                        for sample_name in samples]
    28.  
    29.         positive_annotations = [os.path.join(root_dir, 'Annotations', sample_name + '_1.csv')
    30.                                 for sample_name in samples]
    31.         negative_annotations = [os.path.join(root_dir, 'Annotations', sample_name + '_0.csv')
    32.                                 for sample_name in samples]
    33.  
    34.         # 边界框大小
    35.         positive_sizes = list()
    36.         negative_sizes = list()
    37.         # 边界框坐标
    38.         positive_rects = list()
    39.         negative_rects = list()
    40.  
    41.         for annotation_path in positive_annotations:
    42.             rects = np.loadtxt(annotation_path, dtype=int, delimiter=' ')
    43.             # 存在文件为空或者文件中仅有单行数据
    44.             if len(rects.shape) == 1:
    45.                 # 是否为单行
    46.                 if rects.shape[0] == 4:
    47.                     positive_rects.append(rects)
    48.                     positive_sizes.append(1)
    49.                 else:
    50.                     positive_sizes.append(0)
    51.             else:
    52.                 positive_rects.extend(rects)
    53.                 positive_sizes.append(len(rects))
    54.         for annotation_path in negative_annotations:
    55.             rects = np.loadtxt(annotation_path, dtype=int, delimiter=' ')
    56.             # 和正样本规则一样
    57.             if len(rects.shape) == 1:
    58.                 if rects.shape[0] == 4:
    59.                     negative_rects.append(rects)
    60.                     negative_sizes.append(1)
    61.                 else:
    62.                     positive_sizes.append(0)
    63.             else:
    64.                 negative_rects.extend(rects)
    65.                 negative_sizes.append(len(rects))
    66.         print(positive_rects)
    67.         self.transform = transform
    68.         self.jpeg_images = jpeg_images
    69.         self.positive_sizes = positive_sizes
    70.         self.negative_sizes = negative_sizes
    71.         self.positive_rects = positive_rects
    72.         self.negative_rects = negative_rects
    73.         self.total_positive_num = int(np.sum(positive_sizes))
    74.         self.total_negative_num = int(np.sum(negative_sizes))
    75.  
    76.     def __getitem__(self, index: int):
    77.         # 定位下标所属图像
    78.         image_id = len(self.jpeg_images) - 1
    79.         if index < self.total_positive_num:
    80.             # 正样本
    81.             target = 1
    82.             xmin, ymin, xmax, ymax = self.positive_rects[index]
    83.             # 寻找所属图像
    84.             for i in range(len(self.positive_sizes) - 1):
    85.                 if np.sum(self.positive_sizes[:i]) <= index < np.sum(self.positive_sizes[:(i + 1)]):
    86.                     image_id = i
    87.                     break
    88.             image = self.jpeg_images[image_id][ymin:ymax, xmin:xmax]
    89.         else:
    90.             # 负样本
    91.             target = 0
    92.             idx = index - self.total_positive_num
    93.             xmin, ymin, xmax, ymax = self.negative_rects[idx]
    94.             # 寻找所属图像
    95.             for i in range(len(self.negative_sizes) - 1):
    96.                 if np.sum(self.negative_sizes[:i]) <= idx < np.sum(self.negative_sizes[:(i + 1)]):
    97.                     image_id = i
    98.                     break
    99.             image = self.jpeg_images[image_id][ymin:ymax, xmin:xmax]
    100.  
    101.         # print('index: %d image_id: %d target: %d image.shape: %s [xmin, ymin, xmax, ymax]: [%d, %d, %d, %d]' %
    102.         #       (index, image_id, target, str(image.shape), xmin, ymin, xmax, ymax))
    103.         if self.transform:
    104.             image = self.transform(image)
    105.  
    106.         return image, target
    107.  
    108.     def __len__(self) -> int:
    109.         return self.total_positive_num + self.total_negative_num
    110.  
    111.     def get_positive_num(self) -> int:
    112.         return self.total_positive_num
    113.  
    114.     def get_negative_num(self) -> int:
    115.         return self.total_negative_num
    116.  
    117.  
    118. def test(idx):
    119.     root_dir = '../../data/finetune_car/train'
    120.     train_data_set = CustomFinetuneDataset(root_dir)
    121.  
    122.     print('positive num: %d' % train_data_set.get_positive_num())
    123.     print('negative num: %d' % train_data_set.get_negative_num())
    124.     print('total num: %d' % train_data_set.__len__())
    125.  
    126.     # 测试id=3/66516/66517/530856
    127.     image, target = train_data_set.__getitem__(idx)
    128.     print('target: %d' % target)
    129.  
    130.     image = Image.fromarray(image)
    131.     print(image)
    132.     print(type(image))
    133.  
    134.     # cv2.imshow('image', image)
    135.     # cv2.waitKey(0)
    136.  
    137.  
    138. def test2():
    139.     root_dir = '../../data/finetune_car/train'
    140.     transform = transforms.Compose([
    141.         transforms.ToPILImage(),
    142.         transforms.Resize((227, 227)),
    143.         transforms.ToTensor(),
    144.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    145.     ])
    146.  
    147.     train_data_set = CustomFinetuneDataset(root_dir, transform=transform)
    148.     image, target = train_data_set.__getitem__(530856)
    149.     print('target: %d' % target)
    150.     print('image.shape: ' + str(image.shape))
    151.  
    152.  
    153. def test3():
    154.     root_dir = '../../data/finetune_car/train'
    155.     transform = transforms.Compose([
    156.         transforms.ToPILImage(),
    157.         transforms.Resize((227, 227)),
    158.         transforms.ToTensor(),
    159.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    160.     ])
    161.  
    162.     train_data_set = CustomFinetuneDataset(root_dir, transform=transform)
    163.     data_loader = DataLoader(train_data_set, batch_size=128, num_workers=8, drop_last=True)
    164.  
    165.     inputs, targets = next(data_loader.__iter__())
    166.     print(targets)
    167.     print(inputs.shape)
    168.  
    169.  
    170. if __name__ == '__main__':
    171.     # test(159622)
    172.     # test(4051)
    173.     test3()

    自定义批量采样器类

    custom_batch_sampler.py,该脚本也不用主动执行,在训练微调模型的时候,自然会调用到,以下只说这个脚本做了什么事

    ①CustomBatchSampler类继承自(Sampler)

    ②__init__时通过传入的正样本总数num_positive和负样本总数num_negative得出一个列表self.idx_list,并结合传入的单次正样本数batch_positive和单次负样本数batch_negative算出可迭代次数self.num_iter

    ③__iter__方法中循环self.num_iter次,每次循环中会对正样本随机采集self.batch_positive次index,以及对负样本随机采集self.batch_negative次index,然后打乱存入sampler_list,最后返回一个迭代器iter(sampler)

    以下是custom_batch_sampler.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/3 下午7:38
    5. @file: custom_batch_sampler.py
    6. @author: zj
    7. @description: 自定义采样器
    8. """
    9.  
    10. import numpy  as np
    11. import random
    12. from torch.utils.data import Sampler
    13. from torch.utils.data import DataLoader
    14. import torchvision.transforms as transforms
    15. from utils.data.custom_finetune_dataset import CustomFinetuneDataset
    16.  
    17.  
    18. class CustomBatchSampler(Sampler):
    19.  
    20.     def __init__(self, num_positive, num_negative, batch_positive, batch_negative) -> None:
    21.         """
    22.         2分类数据集
    23.         每次批量处理，其中batch_positive个正样本，batch_negative个负样本
    24.         @param num_positive: 正样本数目
    25.         @param num_negative: 负样本数目
    26.         @param batch_positive: 单次正样本数
    27.         @param batch_negative: 单次负样本数
    28.         """
    29.         self.num_positive = num_positive
    30.         self.num_negative = num_negative
    31.         self.batch_positive = batch_positive
    32.         self.batch_negative = batch_negative
    33.  
    34.         length = num_positive + num_negative
    35.         self.idx_list = list(range(length))
    36.  
    37.         self.batch = batch_negative + batch_positive
    38.         self.num_iter = length // self.batch
    39.  
    40.     def __iter__(self):
    41.         sampler_list = list()
    42.         for i in range(self.num_iter):
    43.             tmp = np.concatenate(
    44.                 (random.sample(self.idx_list[:self.num_positive], self.batch_positive),
    45.                  random.sample(self.idx_list[self.num_positive:], self.batch_negative))
    46.             )
    47.             random.shuffle(tmp)
    48.             sampler_list.extend(tmp)
    49.         return iter(sampler_list)
    50.  
    51.     def __len__(self) -> int:
    52.         return self.num_iter * self.batch
    53.  
    54.     def get_num_batch(self) -> int:
    55.         return self.num_iter
    56.  
    57.  
    58. def test():
    59.     root_dir = '../../data/finetune_car/train'
    60.     train_data_set = CustomFinetuneDataset(root_dir)
    61.     train_sampler = CustomBatchSampler(train_data_set.get_positive_num(), train_data_set.get_negative_num(), 32, 96)
    62.  
    63.     print('sampler len: %d' % train_sampler.__len__())
    64.     print('sampler batch num: %d' % train_sampler.get_num_batch())
    65.  
    66.     first_idx_list = list(train_sampler.__iter__())[:128]
    67.     print(first_idx_list)
    68.     # 单次批量中正样本个数
    69.     print('positive batch: %d' % np.sum(np.array(first_idx_list) < 66517))
    70.  
    71.  
    72. def test2():
    73.     root_dir = '../../data/finetune_car/train'
    74.     transform = transforms.Compose([
    75.         transforms.ToPILImage(),
    76.         transforms.Resize((227, 227)),
    77.         transforms.ToTensor(),
    78.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    79.     ])
    80.  
    81.     train_data_set = CustomFinetuneDataset(root_dir, transform=transform)
    82.     train_sampler = CustomBatchSampler(train_data_set.get_positive_num(), train_data_set.get_negative_num(), 32, 96)
    83.     data_loader = DataLoader(train_data_set, batch_size=128, sampler=train_sampler, num_workers=8, drop_last=True)
    84.  
    85.     inputs, targets = next(data_loader.__iter__())
    86.     print(targets)
    87.     print(inputs.shape)
    88.  
    89.  
    90. if __name__ == '__main__':
    91.     test()

    训练微调模型

    执行finetune.py脚本

    ①调用custom_finetune_dataset.py脚本和custom_batch_sampler.py脚本,得到训练数据data_loaders

    ②使用预训练模型AlexNet网络模型,修改分类器对象classifier的输出为2类(1类是car,一类是背景),然后定义损失函数为交叉熵损失函数,优化函数为SGD,学习率更新器为StepLR,然后开始训练,保存准确率最高的权重到'models/alexnet_car.pth'

    以下是finetune.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/1 上午9:54
    5. @file: finetune.py
    6. @author: zj
    7. @description: 
    8. """
    9.  
    10. import os
    11. import copy
    12. import time
    13. import torch
    14. import torch.nn as nn
    15. import torch.optim as optim
    16. from torch.utils.data import DataLoader
    17. import torchvision.transforms as transforms
    18. import torchvision.models as models
    19. from torchvision.models import AlexNet_Weights
    20.  
    21. from utils.data.custom_finetune_dataset import CustomFinetuneDataset
    22. from utils.data.custom_batch_sampler import CustomBatchSampler
    23. from utils.util import check_dir
    24.  
    25.  
    26. def load_data(data_root_dir):
    27.     transform = transforms.Compose([
    28.         transforms.ToPILImage(),
    29.         transforms.Resize((227, 227)),
    30.         transforms.RandomHorizontalFlip(),
    31.         transforms.ToTensor(),
    32.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    33.     ])
    34.  
    35.     data_loaders = {}
    36.     data_sizes = {}
    37.     for name in ['train', 'val']:
    38.         data_dir = os.path.join(data_root_dir, name)
    39.         data_set = CustomFinetuneDataset(data_dir, transform=transform)
    40.         data_sampler = CustomBatchSampler(data_set.get_positive_num(), data_set.get_negative_num(), 32, 96)
    41.         data_loader = DataLoader(data_set, batch_size=128, sampler=data_sampler, num_workers=8, drop_last=True)
    42.  
    43.         data_loaders[name] = data_loader
    44.         data_sizes[name] = data_sampler.__len__()
    45.  
    46.     return data_loaders, data_sizes
    47.  
    48.  
    49. def train_model(data_loaders, model, criterion, optimizer, lr_scheduler, num_epochs=25, device=None):
    50.     since = time.time()
    51.  
    52.     best_model_weights = copy.deepcopy(model.state_dict())
    53.     best_acc = 0.0
    54.  
    55.     for epoch in range(num_epochs):
    56.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    57.         print('-' * 10)
    58.  
    59.         # Each epoch has a training and validation phase
    60.         for phase in ['train', 'val']:
    61.             if phase == 'train':
    62.                 model.train()  # Set model to training mode
    63.             else:
    64.                 model.eval()  # Set model to evaluate mode
    65.  
    66.             running_loss = 0.0
    67.             running_corrects = 0
    68.  
    69.             # Iterate over data.
    70.             for inputs, labels in data_loaders[phase]:
    71.                 inputs = inputs.to(device)
    72.                 labels = labels.to(device)
    73.  
    74.                 # zero the parameter gradients
    75.                 optimizer.zero_grad()
    76.  
    77.                 # forward
    78.                 # track history if only in train
    79.                 with torch.set_grad_enabled(phase == 'train'):
    80.                     outputs = model(inputs)
    81.                     _, preds = torch.max(outputs, 1)
    82.                     loss = criterion(outputs, labels)
    83.  
    84.                     # backward + optimize only if in training phase
    85.                     if phase == 'train':
    86.                         loss.backward()
    87.                         optimizer.step()
    88.  
    89.                 # statistics
    90.                 running_loss += loss.item() * inputs.size(0)
    91.                 running_corrects += torch.sum(preds == labels.data)
    92.             if phase == 'train':
    93.                 lr_scheduler.step()
    94.  
    95.             epoch_loss = running_loss / data_sizes[phase]
    96.             epoch_acc = running_corrects.double() / data_sizes[phase]
    97.  
    98.             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
    99.                 phase, epoch_loss, epoch_acc))
    100.  
    101.             # deep copy the model
    102.             if phase == 'val' and epoch_acc > best_acc:
    103.                 best_acc = epoch_acc
    104.                 best_model_weights = copy.deepcopy(model.state_dict())
    105.  
    106.         print()
    107.  
    108.     time_elapsed = time.time() - since
    109.     print('Training complete in {:.0f}m {:.0f}s'.format(
    110.         time_elapsed // 60, time_elapsed % 60))
    111.     print('Best val Acc: {:4f}'.format(best_acc))
    112.  
    113.     # load best model weights
    114.     model.load_state_dict(best_model_weights)
    115.     return model
    116.  
    117.  
    118. if __name__ == '__main__':
    119.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    120.  
    121.     data_loaders, data_sizes = load_data('./data/finetune_car')
    122.  
    123.     model = models.alexnet(weights=AlexNet_Weights.IMAGENET1K_V1)
    124.     # print(model)
    125.     num_features = model.classifier[6].in_features
    126.     model.classifier[6] = nn.Linear(num_features, 2)
    127.     # print(model)
    128.     model = model.to(device)
    129.  
    130.     criterion = nn.CrossEntropyLoss()
    131.     optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
    132.     lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
    133.  
    134.     best_model = train_model(data_loaders, model, criterion, optimizer, lr_scheduler, device=device, num_epochs=25)
    135.     # 保存最好的模型参数
    136.     check_dir('./models')
    137.     torch.save(best_model.state_dict(), 'models/alexnet_car.pth')

    分类器训练

    准备分类器数据集

    执行create_classifier_data.py脚本

    ①把'../../data/finetune_car/train/JPEGImages/'和'../../data/finetune_car/val/JPEGImages/'中的.jpg文件复制到'../../data/classifier_car/train/JPEGImages/'和'../../data/classifier_car/val/JPEGImages/',然后又把'../../data/finetune_car/train/car.csv'和'../../data/finetune_car/val/car.csv'分别复制到'../../data/classifier_car/train/car.csv'和'../../data/classifier_car/val/car.csv'

    ②根据'../../data/classifier_car/train/car.csv'和'../../data/classifier_car/val/car.csv'文件内容分别读取'../../data/classifier_car/train/JPEGImages/'和'../../data/classifier_car/val/JPEGImages/'中的图片,并传入parse_annotation_jpeg方法

    ③parse_annotation_jpeg方法中,先获取候选框rects,然后从.xml文件中获取标注框bndboxs,接着计算候选框和标注框的IoU得到iou_list,遍历iou_list,选出0

    以下是create_classifier_data.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/1 下午7:17
    5. @file: create_classifier_data.py
    6. @author: zj
    7. @description: 创建分类器数据集
    8. """
    9.  
    10. import random
    11. import numpy as np
    12. import shutil
    13. import time
    14. import cv2
    15. import os
    16. import xmltodict
    17. import selectivesearch
    18. from utils.util import check_dir
    19. from utils.util import parse_car_csv
    20. from utils.util import parse_xml
    21. from utils.util import iou
    22. from utils.util import compute_ious
    23.  
    24. # train
    25. # positive num: 625
    26. # negative num: 366028
    27. # val
    28. # positive num: 625
    29. # negative num: 321474
    30.  
    31. def parse_annotation_jpeg(annotation_path, jpeg_path, gs):
    32.     """
    33.     获取正负样本（注：忽略属性difficult为True的标注边界框）
    34.     正样本：标注边界框
    35.     负样本：IoU大于0，小于等于0.3。为了进一步限制负样本数目，其大小必须大于标注框的1/5
    36.     """
    37.     img = cv2.imread(jpeg_path)
    38.  
    39.     selectivesearch.config(gs, img, strategy='q')
    40.     # 计算候选建议
    41.     rects = selectivesearch.get_rects(gs)
    42.     # 获取标注边界框
    43.     bndboxs = parse_xml(annotation_path)
    44.  
    45.     # 标注框大小
    46.     maximum_bndbox_size = 0
    47.     for bndbox in bndboxs:
    48.         xmin, ymin, xmax, ymax = bndbox
    49.         bndbox_size = (ymax - ymin) * (xmax - xmin)
    50.         if bndbox_size > maximum_bndbox_size:
    51.             maximum_bndbox_size = bndbox_size
    52.  
    53.     # 获取候选建议和标注边界框的IoU
    54.     iou_list = compute_ious(rects, bndboxs)
    55.  
    56.     positive_list = list()
    57.     negative_list = list()
    58.     for i in range(len(iou_list)):
    59.         xmin, ymin, xmax, ymax = rects[i]
    60.         rect_size = (ymax - ymin) * (xmax - xmin)
    61.  
    62.         iou_score = iou_list[i]
    63.         if 0 < iou_score <= 0.3 and rect_size > maximum_bndbox_size / 5.0:
    64.             # 负样本
    65.             negative_list.append(rects[i])
    66.         else:
    67.             pass
    68.  
    69.     return bndboxs, negative_list
    70.  
    71.  
    72. if __name__ == '__main__':
    73.     car_root_dir = '../../data/voc_car/'
    74.     classifier_root_dir = '../../data/classifier_car/'
    75.     check_dir(classifier_root_dir)
    76.  
    77.     gs = selectivesearch.get_selective_search()
    78.     for name in ['train', 'val']:
    79.         src_root_dir = os.path.join(car_root_dir, name)
    80.         src_annotation_dir = os.path.join(src_root_dir, 'Annotations')
    81.         src_jpeg_dir = os.path.join(src_root_dir, 'JPEGImages')
    82.  
    83.         dst_root_dir = os.path.join(classifier_root_dir, name)
    84.         dst_annotation_dir = os.path.join(dst_root_dir, 'Annotations')
    85.         dst_jpeg_dir = os.path.join(dst_root_dir, 'JPEGImages')
    86.         check_dir(dst_root_dir)
    87.         check_dir(dst_annotation_dir)
    88.         check_dir(dst_jpeg_dir)
    89.  
    90.         total_num_positive = 0
    91.         total_num_negative = 0
    92.  
    93.         samples = parse_car_csv(src_root_dir)
    94.         # 复制csv文件
    95.         src_csv_path = os.path.join(src_root_dir, 'car.csv')
    96.         dst_csv_path = os.path.join(dst_root_dir, 'car.csv')
    97.         shutil.copyfile(src_csv_path, dst_csv_path)
    98.         for sample_name in samples:
    99.             since = time.time()
    100.  
    101.             src_annotation_path = os.path.join(src_annotation_dir, sample_name + '.xml')
    102.             src_jpeg_path = os.path.join(src_jpeg_dir, sample_name + '.jpg')
    103.             # 获取正负样本
    104.             positive_list, negative_list = parse_annotation_jpeg(src_annotation_path, src_jpeg_path, gs)
    105.             total_num_positive += len(positive_list)
    106.             total_num_negative += len(negative_list)
    107.  
    108.             dst_annotation_positive_path = os.path.join(dst_annotation_dir, sample_name + '_1' + '.csv')
    109.             dst_annotation_negative_path = os.path.join(dst_annotation_dir, sample_name + '_0' + '.csv')
    110.             dst_jpeg_path = os.path.join(dst_jpeg_dir, sample_name + '.jpg')
    111.             # 保存图片
    112.             shutil.copyfile(src_jpeg_path, dst_jpeg_path)
    113.             # 保存正负样本标注
    114.             np.savetxt(dst_annotation_positive_path, np.array(positive_list), fmt='%d', delimiter=' ')
    115.             np.savetxt(dst_annotation_negative_path, np.array(negative_list), fmt='%d', delimiter=' ')
    116.  
    117.             time_elapsed = time.time() - since
    118.             print('parse {}.png in {:.0f}m {:.0f}s'.format(sample_name, time_elapsed // 60, time_elapsed % 60))
    119.         print('%s positive num: %d' % (name, total_num_positive))
    120.         print('%s negative num: %d' % (name, total_num_negative))
    121.     print('done')

    自定义分类器数据集类

    custom_classifier_dataset.py,该脚本不用主动执行,在训练分类器模型的时候,自然会调用到,以下只说这个脚本做了什么事

    ①CustomClassifierDataset类继承自Dataset

    ②__init__时读取'../../data/classifier_car/train/JPEGImages/'或'../../data/classifier_car/val/JPEGImages/'文件夹中的图片,读取'../../data/classifier_car/train/Annotations/'或'../../data/classifier_car/val/Annotations/'中的正负样本集,记录正样本列表self.positive_list,负样本总数self.negative_list,正样本候选框positive_rects,负样本候选框negative_rects

    ③__getitem__方法传入index,如果index小于正样本总数len(self.positive_list),那么返回对应正样本的图片和标签(1),否则返回对应负样本的图片和标签(0)。

    以下是custom_classifier_dataset.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/4 下午4:00
    5. @file: custom_classifier_dataset.py
    6. @author: zj
    7. @description: 分类器数据集类，可进行正负样本集替换，适用于hard negative mining操作
    8. """
    9.  
    10. import numpy  as np
    11. import os
    12. import cv2
    13. from PIL import Image
    14. from torch.utils.data import Dataset
    15. from torch.utils.data import DataLoader
    16. import torchvision.transforms as transforms
    17.  
    18. from utils.util import parse_car_csv
    19.  
    20.  
    21. class CustomClassifierDataset(Dataset):
    22.  
    23.     def __init__(self, root_dir, transform=None):
    24.         samples = parse_car_csv(root_dir)
    25.  
    26.         jpeg_images = list()
    27.         positive_list = list()
    28.         negative_list = list()
    29.         for idx in range(len(samples)):
    30.             sample_name = samples[idx]
    31.             jpeg_images.append(cv2.imread(os.path.join(root_dir, 'JPEGImages', sample_name + ".jpg")))
    32.  
    33.             positive_annotation_path = os.path.join(root_dir, 'Annotations', sample_name + '_1.csv')
    34.             positive_annotations = np.loadtxt(positive_annotation_path, dtype=int, delimiter=' ')
    35.             # 考虑csv文件为空或者仅包含单个标注框
    36.             if len(positive_annotations.shape) == 1:
    37.                 # 单个标注框坐标
    38.                 if positive_annotations.shape[0] == 4:
    39.                     positive_dict = dict()
    40.  
    41.                     positive_dict['rect'] = positive_annotations
    42.                     positive_dict['image_id'] = idx
    43.                     # positive_dict['image_name'] = sample_name
    44.  
    45.                     positive_list.append(positive_dict)
    46.             else:
    47.                 for positive_annotation in positive_annotations:
    48.                     positive_dict = dict()
    49.  
    50.                     positive_dict['rect'] = positive_annotation
    51.                     positive_dict['image_id'] = idx
    52.                     # positive_dict['image_name'] = sample_name
    53.  
    54.                     positive_list.append(positive_dict)
    55.  
    56.             negative_annotation_path = os.path.join(root_dir, 'Annotations', sample_name + '_0.csv')
    57.             negative_annotations = np.loadtxt(negative_annotation_path, dtype=int, delimiter=' ')
    58.             # 考虑csv文件为空或者仅包含单个标注框
    59.             if len(negative_annotations.shape) == 1:
    60.                 # 单个标注框坐标
    61.                 if negative_annotations.shape[0] == 4:
    62.                     negative_dict = dict()
    63.  
    64.                     negative_dict['rect'] = negative_annotations
    65.                     negative_dict['image_id'] = idx
    66.                     # negative_dict['image_name'] = sample_name
    67.  
    68.                     negative_list.append(negative_dict)
    69.             else:
    70.                 for negative_annotation in negative_annotations:
    71.                     negative_dict = dict()
    72.  
    73.                     negative_dict['rect'] = negative_annotation
    74.                     negative_dict['image_id'] = idx
    75.                     # negative_dict['image_name'] = sample_name
    76.  
    77.                     negative_list.append(negative_dict)
    78.  
    79.         self.transform = transform
    80.         self.jpeg_images = jpeg_images
    81.         self.positive_list = positive_list
    82.         self.negative_list = negative_list
    83.  
    84.     def __getitem__(self, index: int):
    85.         # 定位下标所属图像
    86.         if index < len(self.positive_list):
    87.             # 正样本
    88.             target = 1
    89.             positive_dict = self.positive_list[index]
    90.  
    91.             xmin, ymin, xmax, ymax = positive_dict['rect']
    92.             image_id = positive_dict['image_id']
    93.  
    94.             image = self.jpeg_images[image_id][ymin:ymax, xmin:xmax]
    95.             cache_dict = positive_dict
    96.         else:
    97.             # 负样本
    98.             target = 0
    99.             idx = index - len(self.positive_list)
    100.             negative_dict = self.negative_list[idx]
    101.  
    102.             xmin, ymin, xmax, ymax = negative_dict['rect']
    103.             image_id = negative_dict['image_id']
    104.  
    105.             image = self.jpeg_images[image_id][ymin:ymax, xmin:xmax]
    106.             cache_dict = negative_dict
    107.  
    108.         # print('index: %d image_id: %d target: %d image.shape: %s [xmin, ymin, xmax, ymax]: [%d, %d, %d, %d]' %
    109.         #       (index, image_id, target, str(image.shape), xmin, ymin, xmax, ymax))
    110.         if self.transform:
    111.             image = self.transform(image)
    112.  
    113.         return image, target, cache_dict
    114.  
    115.     def __len__(self) -> int:
    116.         return len(self.positive_list) + len(self.negative_list)
    117.  
    118.     def get_transform(self):
    119.         return self.transform
    120.  
    121.     def get_jpeg_images(self) -> list:
    122.         return self.jpeg_images
    123.  
    124.     def get_positive_num(self) -> int:
    125.         return len(self.positive_list)
    126.  
    127.     def get_negative_num(self) -> int:
    128.         return len(self.negative_list)
    129.  
    130.     def get_positives(self) -> list:
    131.         return self.positive_list
    132.  
    133.     def get_negatives(self) -> list:
    134.         return self.negative_list
    135.  
    136.     # 用于hard negative mining
    137.     # 替换负样本
    138.     def set_negative_list(self, negative_list):
    139.         self.negative_list = negative_list
    140.  
    141.  
    142. def test(idx):
    143.     root_dir = '../../data/classifier_car/val'
    144.     train_data_set = CustomClassifierDataset(root_dir)
    145.  
    146.     print('positive num: %d' % train_data_set.get_positive_num())
    147.     print('negative num: %d' % train_data_set.get_negative_num())
    148.     print('total num: %d' % train_data_set.__len__())
    149.  
    150.     # 测试id=3/66516/66517/530856
    151.     image, target, cache_dict = train_data_set.__getitem__(idx)
    152.     print('target: %d' % target)
    153.     print('dict: ' + str(cache_dict))
    154.  
    155.     image = Image.fromarray(image)
    156.     print(image)
    157.     print(type(image))
    158.  
    159.     # cv2.imshow('image', image)
    160.     # cv2.waitKey(0)
    161.  
    162.  
    163. def test2():
    164.     root_dir = '../../data/classifier_car/train'
    165.     transform = transforms.Compose([
    166.         transforms.ToPILImage(),
    167.         transforms.Resize((227, 227)),
    168.         transforms.ToTensor(),
    169.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    170.     ])
    171.  
    172.     train_data_set = CustomClassifierDataset(root_dir, transform=transform)
    173.     image, target, cache_dict = train_data_set.__getitem__(230856)
    174.     print('target: %d' % target)
    175.     print('dict: ' + str(cache_dict))
    176.     print('image.shape: ' + str(image.shape))
    177.  
    178.  
    179. def test3():
    180.     root_dir = '../../data/classifier_car/train'
    181.     transform = transforms.Compose([
    182.         transforms.ToPILImage(),
    183.         transforms.Resize((227, 227)),
    184.         transforms.ToTensor(),
    185.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    186.     ])
    187.  
    188.     train_data_set = CustomClassifierDataset(root_dir, transform=transform)
    189.     data_loader = DataLoader(train_data_set, batch_size=128, num_workers=8, drop_last=True)
    190.  
    191.     inputs, targets, cache_dicts = next(data_loader.__iter__())
    192.     print(targets)
    193.     print(inputs.shape)
    194.  
    195.  
    196. if __name__ == '__main__':
    197.     # test(159622)
    198.     # test(4051)
    199.     test(24768)
    200.     # test2()
    201.     # test3()

    自定义批量采样器类

    同"卷积神经网络微调模型"中的"自定义批量采样器类",在训练分类器模型的时候,自然会调用到

    训练分类器

    执行linear_svm.py脚本

    ①调用custom_classifier_dataset.py脚本和custom_batch_sampler.py脚本,得到训练数据data_loaders

    ②使用AlexNet网络模型,修改分类器对象classifier的输出为2类(1类是car,一类是背景),加载之前微调训练的权重alexnet_car.pth,并设置参数冻结,然后再添加一个全连接层作为svm分类器,定义损失函数为折页损失函数,优化函数为SGD,学习率更新器为StepLR,然后开始训练,保存准确率最高的权重到'models/best_linear_svm_alexnet_car.pth'

    以下是linear_svm.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/3/1 下午2:38
    5. @file: linear_svm.py
    6. @author: zj
    7. @description: 
    8. """
    9.  
    10. import time
    11. import copy
    12. import os
    13. import random
    14. import numpy as np
    15. import torch
    16. import torch.nn as nn
    17. import torch.optim as optim
    18. from torch.utils.data import DataLoader
    19. import torchvision.transforms as transforms
    20. from torchvision.models import alexnet
    21.  
    22. from utils.data.custom_classifier_dataset import CustomClassifierDataset
    23. from utils.data.custom_hard_negative_mining_dataset import CustomHardNegativeMiningDataset
    24. from utils.data.custom_batch_sampler import CustomBatchSampler
    25. from utils.util import check_dir
    26. from utils.util import save_model
    27.  
    28. batch_positive = 32
    29. batch_negative = 96
    30. batch_total = 128
    31.  
    32.  
    33. def load_data(data_root_dir):
    34.     transform = transforms.Compose([
    35.         transforms.ToPILImage(),
    36.         transforms.Resize((227, 227)),
    37.         transforms.RandomHorizontalFlip(),
    38.         transforms.ToTensor(),
    39.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    40.     ])
    41.  
    42.     data_loaders = {}
    43.     data_sizes = {}
    44.     remain_negative_list = list()
    45.     for name in ['train', 'val']:
    46.         data_dir = os.path.join(data_root_dir, name)
    47.  
    48.         data_set = CustomClassifierDataset(data_dir, transform=transform)
    49.         if name == 'train':
    50.             """
    51.             使用hard negative mining方式
    52.             初始正负样本比例为1:1。由于正样本数远小于负样本，所以以正样本数为基准，在负样本集中随机提取同样数目负样本作为初始负样本集
    53.             """
    54.             positive_list = data_set.get_positives()
    55.             negative_list = data_set.get_negatives()
    56.  
    57.             init_negative_idxs = random.sample(range(len(negative_list)), len(positive_list))
    58.             init_negative_list = [negative_list[idx] for idx in range(len(negative_list)) if idx in init_negative_idxs]
    59.             remain_negative_list = [negative_list[idx] for idx in range(len(negative_list))
    60.                                     if idx not in init_negative_idxs]
    61.  
    62.             data_set.set_negative_list(init_negative_list)
    63.             data_loaders['remain'] = remain_negative_list
    64.  
    65.         sampler = CustomBatchSampler(data_set.get_positive_num(), data_set.get_negative_num(),
    66.                                      batch_positive, batch_negative)
    67.  
    68.         data_loader = DataLoader(data_set, batch_size=batch_total, sampler=sampler, num_workers=8, drop_last=True)
    69.         data_loaders[name] = data_loader
    70.         data_sizes[name] = len(sampler)
    71.     return data_loaders, data_sizes
    72.  
    73.  
    74. def hinge_loss(outputs, labels):
    75.     """
    76.     折页损失计算
    77.     :param outputs: 大小为(N, num_classes)
    78.     :param labels: 大小为(N)
    79.     :return: 损失值
    80.     """
    81.     num_labels = len(labels)
    82.     corrects = outputs[range(num_labels), labels].unsqueeze(0).T
    83.  
    84.     # 最大间隔
    85.     margin = 1.0
    86.     margins = outputs - corrects + margin
    87.     loss = torch.sum(torch.max(margins, 1)[0]) / len(labels)
    88.  
    89.     # # 正则化强度
    90.     # reg = 1e-3
    91.     # loss += reg * torch.sum(weight ** 2)
    92.  
    93.     return loss
    94.  
    95.  
    96. def add_hard_negatives(hard_negative_list, negative_list, add_negative_list):
    97.     for item in hard_negative_list:
    98.         if len(add_negative_list) == 0:
    99.             # 第一次添加负样本
    100.             negative_list.append(item)
    101.             add_negative_list.append(list(item['rect']))
    102.         if list(item['rect']) not in add_negative_list:
    103.             negative_list.append(item)
    104.             add_negative_list.append(list(item['rect']))
    105.  
    106.  
    107. def get_hard_negatives(preds, cache_dicts):
    108.     fp_mask = preds == 1
    109.     tn_mask = preds == 0
    110.  
    111.     fp_rects = cache_dicts['rect'][fp_mask].numpy()
    112.     fp_image_ids = cache_dicts['image_id'][fp_mask].numpy()
    113.  
    114.     tn_rects = cache_dicts['rect'][tn_mask].numpy()
    115.     tn_image_ids = cache_dicts['image_id'][tn_mask].numpy()
    116.  
    117.     hard_negative_list = [{'rect': fp_rects[idx], 'image_id': fp_image_ids[idx]} for idx in range(len(fp_rects))]
    118.     easy_negatie_list = [{'rect': tn_rects[idx], 'image_id': tn_image_ids[idx]} for idx in range(len(tn_rects))]
    119.  
    120.     return hard_negative_list, easy_negatie_list
    121.  
    122.  
    123. def train_model(data_loaders, model, criterion, optimizer, lr_scheduler, num_epochs=25, device=None):
    124.     since = time.time()
    125.  
    126.     best_model_weights = copy.deepcopy(model.state_dict())
    127.     best_acc = 0.0
    128.  
    129.     for epoch in range(num_epochs):
    130.  
    131.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    132.         print('-' * 10)
    133.  
    134.         # Each epoch has a training and validation phase
    135.         for phase in ['train', 'val']:
    136.             if phase == 'train':
    137.                 model.train()  # Set model to training mode
    138.             else:
    139.                 model.eval()  # Set model to evaluate mode
    140.  
    141.             running_loss = 0.0
    142.             running_corrects = 0
    143.  
    144.             # 输出正负样本数
    145.             data_set = data_loaders[phase].dataset
    146.             print('{} - positive_num: {} - negative_num: {} - data size: {}'.format(
    147.                 phase, data_set.get_positive_num(), data_set.get_negative_num(), data_sizes[phase]))
    148.  
    149.             # Iterate over data.
    150.             for inputs, labels, cache_dicts in data_loaders[phase]:
    151.                 inputs = inputs.to(device)
    152.                 labels = labels.to(device)
    153.  
    154.                 # zero the parameter gradients
    155.                 optimizer.zero_grad()
    156.  
    157.                 # forward
    158.                 # track history if only in train
    159.                 with torch.set_grad_enabled(phase == 'train'):
    160.                     outputs = model(inputs)
    161.                     # print(outputs.shape)
    162.                     _, preds = torch.max(outputs, 1)
    163.                     loss = criterion(outputs, labels)
    164.  
    165.                     # backward + optimize only if in training phase
    166.                     if phase == 'train':
    167.                         loss.backward()
    168.                         optimizer.step()
    169.  
    170.                 # statistics
    171.                 running_loss += loss.item() * inputs.size(0)
    172.                 running_corrects += torch.sum(preds == labels.data)
    173.             if phase == 'train':
    174.                 lr_scheduler.step()
    175.  
    176.             epoch_loss = running_loss / data_sizes[phase]
    177.             epoch_acc = running_corrects.double() / data_sizes[phase]
    178.  
    179.             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
    180.                 phase, epoch_loss, epoch_acc))
    181.  
    182.             # deep copy the model
    183.             if phase == 'val' and epoch_acc > best_acc:
    184.                 best_acc = epoch_acc
    185.                 best_model_weights = copy.deepcopy(model.state_dict())
    186.  
    187.         # 每一轮训练完成后，测试剩余负样本集，进行hard negative mining
    188.         train_dataset = data_loaders['train'].dataset
    189.         remain_negative_list = data_loaders['remain']
    190.         jpeg_images = train_dataset.get_jpeg_images()
    191.         transform = train_dataset.get_transform()
    192.  
    193.         with torch.set_grad_enabled(False):
    194.             remain_dataset = CustomHardNegativeMiningDataset(remain_negative_list, jpeg_images, transform=transform)
    195.             remain_data_loader = DataLoader(remain_dataset, batch_size=batch_total, num_workers=8, drop_last=True)
    196.  
    197.             # 获取训练数据集的负样本集
    198.             negative_list = train_dataset.get_negatives()
    199.             # 记录后续增加的负样本
    200.             add_negative_list = data_loaders.get('add_negative', [])
    201.  
    202.             running_corrects = 0
    203.             # Iterate over data.
    204.             for inputs, labels, cache_dicts in remain_data_loader:
    205.                 inputs = inputs.to(device)
    206.                 labels = labels.to(device)
    207.  
    208.                 # zero the parameter gradients
    209.                 optimizer.zero_grad()
    210.  
    211.                 outputs = model(inputs)
    212.                 # print(outputs.shape)
    213.                 _, preds = torch.max(outputs, 1)
    214.  
    215.                 running_corrects += torch.sum(preds == labels.data)
    216.  
    217.                 hard_negative_list, easy_neagtive_list = get_hard_negatives(preds.cpu().numpy(), cache_dicts)
    218.                 add_hard_negatives(hard_negative_list, negative_list, add_negative_list)
    219.  
    220.             remain_acc = running_corrects.double() / len(remain_negative_list)
    221.             print('remain negative size: {}, acc: {:.4f}'.format(len(remain_negative_list), remain_acc))
    222.  
    223.             # 训练完成后，重置负样本，进行hard negatives mining
    224.             train_dataset.set_negative_list(negative_list)
    225.             tmp_sampler = CustomBatchSampler(train_dataset.get_positive_num(), train_dataset.get_negative_num(),
    226.                                              batch_positive, batch_negative)
    227.             data_loaders['train'] = DataLoader(train_dataset, batch_size=batch_total, sampler=tmp_sampler,
    228.                                                num_workers=8, drop_last=True)
    229.             data_loaders['add_negative'] = add_negative_list
    230.             # 重置数据集大小
    231.             data_sizes['train'] = len(tmp_sampler)
    232.  
    233.         # 每训练一轮就保存
    234.         save_model(model, 'models/linear_svm_alexnet_car_%d.pth' % epoch)
    235.  
    236.     time_elapsed = time.time() - since
    237.     print('Training complete in {:.0f}m {:.0f}s'.format(
    238.         time_elapsed // 60, time_elapsed % 60))
    239.     print('Best val Acc: {:4f}'.format(best_acc))
    240.  
    241.     # load best model weights
    242.     model.load_state_dict(best_model_weights)
    243.     return model
    244.  
    245.  
    246. if __name__ == '__main__':
    247.     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    248.     # device = 'cpu'
    249.  
    250.     data_loaders, data_sizes = load_data('./data/classifier_car')
    251.  
    252.     # 加载CNN模型
    253.     model_path = './models/alexnet_car.pth'
    254.     model = alexnet()
    255.     num_classes = 2
    256.     num_features = model.classifier[6].in_features
    257.     model.classifier[6] = nn.Linear(num_features, num_classes)
    258.     model.load_state_dict(torch.load(model_path))
    259.     model.eval()
    260.     # 固定特征提取
    261.     for param in model.parameters():
    262.         param.requires_grad = False
    263.     # 创建SVM分类器
    264.     model.classifier[6] = nn.Linear(num_features, num_classes)
    265.     # print(model)
    266.     model = model.to(device)
    267.  
    268.     criterion = hinge_loss
    269.     # 由于初始训练集数量很少，所以降低学习率
    270.     optimizer = optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
    271.     # 共训练10轮，每隔4论减少一次学习率
    272.     lr_schduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)
    273.  
    274.     best_model = train_model(data_loaders, model, criterion, optimizer, lr_schduler, num_epochs=10, device=device)
    275.     # 保存最好的模型参数
    276.     save_model(best_model, 'models/best_linear_svm_alexnet_car.pth')

    边界框回归训练

    准备边界框回归数据集

    执行create_bbox_regression_data.py脚本

    ①读取'../../data/voc_car/train/Annotations/'中的标注框信息存入bndboxs和'../../data/finetune_car/train/Annotations/'中的正样本数据存入positive_bndboxes,计算标注框和正样本数据的IoU,针对IoU＞0.6的正样本数据,保存其到'../../data/bbox_regression/positive/'中,并保存对应的图片到'../../data/bbox_regression/JPEGImages/'中,保存对应的标注框信息到'../../data/bbox_regression/bndboxs/'中,保存对应的图片名到'../../data/bbox_regression/car.csv'中

    以下是create_bbox_regression_data.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/4/3 下午7:19
    5. @file: create_bbox_regression_data.py
    6. @author: zj
    7. @description: 创建边界框回归数据集
    8. """
    9.  
    10. import os
    11. import shutil
    12. import numpy as np
    13. import utils.util as util
    14.  
    15. # 正样本边界框数目：37222
    16.  
    17. if __name__ == '__main__':
    18.     """
    19.     从voc_car/train目录中提取标注边界框坐标
    20.     从finetune_car/train目录中提取训练集正样本坐标（IoU>=0.5），进一步提取IoU>0.6的边界框
    21.     数据集保存在bbox_car目录下
    22.     """
    23.     voc_car_train_dir = '../../data/voc_car/train'
    24.     # ground truth
    25.     gt_annotation_dir = os.path.join(voc_car_train_dir, 'Annotations')
    26.     jpeg_dir = os.path.join(voc_car_train_dir, 'JPEGImages')
    27.  
    28.     classifier_car_train_dir = '../../data/finetune_car/train'
    29.     # positive
    30.     positive_annotation_dir = os.path.join(classifier_car_train_dir, 'Annotations')
    31.  
    32.     dst_root_dir = '../../data/bbox_regression/'
    33.     dst_jpeg_dir = os.path.join(dst_root_dir, 'JPEGImages')
    34.     dst_bndbox_dir = os.path.join(dst_root_dir, 'bndboxs')
    35.     dst_positive_dir = os.path.join(dst_root_dir, 'positive')
    36.  
    37.     util.check_dir(dst_root_dir)
    38.     util.check_dir(dst_jpeg_dir)
    39.     util.check_dir(dst_bndbox_dir)
    40.     util.check_dir(dst_positive_dir)
    41.  
    42.     samples = util.parse_car_csv(voc_car_train_dir)
    43.     res_samples = list()
    44.     total_positive_num = 0
    45.     for sample_name in samples:
    46.         # 提取正样本边界框坐标（IoU>=0.5）
    47.         positive_annotation_path = os.path.join(positive_annotation_dir, sample_name + '_1.csv')
    48.         positive_bndboxes = np.loadtxt(positive_annotation_path, dtype=int, delimiter=' ')
    49.         # 提取标注边界框
    50.         gt_annotation_path = os.path.join(gt_annotation_dir, sample_name + '.xml')
    51.         bndboxs = util.parse_xml(gt_annotation_path)
    52.         # 计算符合条件（IoU>0.6）的候选建议
    53.         positive_list = list()
    54.         if len(positive_bndboxes.shape) == 1 and len(positive_bndboxes) != 0:
    55.             scores = util.iou(positive_bndboxes, bndboxs)
    56.             if np.max(scores) > 0.6:
    57.                 positive_list.append(positive_bndboxes)
    58.         elif len(positive_bndboxes.shape) == 2:
    59.             for positive_bndboxe in positive_bndboxes:
    60.                 scores = util.iou(positive_bndboxe, bndboxs)
    61.                 if np.max(scores) > 0.6:
    62.                     positive_list.append(positive_bndboxe)
    63.         else:
    64.             pass
    65.  
    66.         # 如果存在正样本边界框（IoU>0.6），那么保存相应的图片以及标注边界框
    67.         if len(positive_list) > 0:
    68.             # 保存图片
    69.             jpeg_path = os.path.join(jpeg_dir, sample_name + ".jpg")
    70.             dst_jpeg_path = os.path.join(dst_jpeg_dir, sample_name + ".jpg")
    71.             shutil.copyfile(jpeg_path, dst_jpeg_path)
    72.             # 保存标注边界框
    73.             dst_bndbox_path = os.path.join(dst_bndbox_dir, sample_name + ".csv")
    74.             np.savetxt(dst_bndbox_path, bndboxs, fmt='%s', delimiter=' ')
    75.             # 保存正样本边界框
    76.             dst_positive_path = os.path.join(dst_positive_dir, sample_name + ".csv")
    77.             np.savetxt(dst_positive_path, np.array(positive_list), fmt='%s', delimiter=' ')
    78.  
    79.             total_positive_num += len(positive_list)
    80.             res_samples.append(sample_name)
    81.             print('save {} done'.format(sample_name))
    82.         else:
    83.             print('-------- {} 不符合条件'.format(sample_name))
    84.  
    85.     dst_csv_path = os.path.join(dst_root_dir, 'car.csv')
    86.     np.savetxt(dst_csv_path, res_samples, fmt='%s', delimiter=' ')
    87.     print('total positive num: {}'.format(total_positive_num))
    88.     print('done')

    自定义边界框回归训练数据集类

    custom_bbox_regression_dataset.py,该脚本不用主动执行,在训练分类器模型的时候,自然会调用到,以下只说这个脚本做了什么事

    ①BBoxRegressionDataset类继承自Dataset

    ②__init__时读取'../../data/bbox_regression/JPEGImages/'文件夹中的图片,存入self.jpeg_list,又读取'../../data/bbox_regression/bndboxs/'中的标注框信息和'../../data/bbox_regression/positive/'中的正样本数据并都存入self.box_list

    ③__getitem__方法计算并返回图片和相对坐标差

    以下是custom_bbox_regression_dataset.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/4/3 下午8:07
    5. @file: custom_bbox_regression_dataset.py
    6. @author: zj
    7. @description:
    8. """
    9.  
    10. import os
    11. import cv2
    12. import numpy as np
    13. import torch
    14. import torchvision.transforms as transforms
    15. from torch.utils.data import Dataset
    16. from torch.utils.data import DataLoader
    17.  
    18. import utils.util as util
    19.  
    20.  
    21. class BBoxRegressionDataset(Dataset):
    22.  
    23.     def __init__(self, root_dir, transform=None):
    24.         super(BBoxRegressionDataset, self).__init__()
    25.         self.transform = transform
    26.  
    27.         samples = util.parse_car_csv(root_dir)
    28.         jpeg_list = list()
    29.         # 保存{'image_id': ?, 'positive': ?, 'bndbox': ?}
    30.         box_list = list()
    31.         for i in range(len(samples)):
    32.             sample_name = samples[i]
    33.  
    34.             jpeg_path = os.path.join(root_dir, 'JPEGImages', sample_name + '.jpg')
    35.             bndbox_path = os.path.join(root_dir, 'bndboxs', sample_name + '.csv')
    36.             positive_path = os.path.join(root_dir, 'positive', sample_name + '.csv')
    37.  
    38.             jpeg_list.append(cv2.imread(jpeg_path))
    39.             bndboxes = np.loadtxt(bndbox_path, dtype=int, delimiter=' ')
    40.             positives = np.loadtxt(positive_path, dtype=int, delimiter=' ')
    41.  
    42.             if len(positives.shape) == 1:
    43.                 bndbox = self.get_bndbox(bndboxes, positives)
    44.                 box_list.append({'image_id': i, 'positive': positives, 'bndbox': bndbox})
    45.             else:
    46.                 for positive in positives:
    47.                     bndbox = self.get_bndbox(bndboxes, positive)
    48.                     box_list.append({'image_id': i, 'positive': positive, 'bndbox': bndbox})
    49.  
    50.         self.jpeg_list = jpeg_list
    51.         self.box_list = box_list
    52.  
    53.     def __getitem__(self, index: int):
    54.         assert index < self.__len__(), '数据集大小为%d，当前输入下标为%d' % (self.__len__(), index)
    55.  
    56.         box_dict = self.box_list[index]
    57.         image_id = box_dict['image_id']
    58.         positive = box_dict['positive']
    59.         bndbox = box_dict['bndbox']
    60.  
    61.         # 获取预测图像
    62.         jpeg_img = self.jpeg_list[image_id]
    63.         xmin, ymin, xmax, ymax = positive
    64.         image = jpeg_img[ymin:ymax, xmin:xmax]
    65.  
    66.         if self.transform:
    67.             image = self.transform(image)
    68.  
    69.         # 计算P/G的x/y/w/h
    70.         target = dict()
    71.         p_w = xmax - xmin
    72.         p_h = ymax - ymin
    73.         p_x = xmin + p_w / 2
    74.         p_y = ymin + p_h / 2
    75.  
    76.         xmin, ymin, xmax, ymax = bndbox
    77.         g_w = xmax - xmin
    78.         g_h = ymax - ymin
    79.         g_x = xmin + g_w / 2
    80.         g_y = ymin + g_h / 2
    81.  
    82.         # 计算t
    83.         t_x = (g_x - p_x) / p_w
    84.         t_y = (g_y - p_y) / p_h
    85.         t_w = np.log(g_w / p_w)
    86.         t_h = np.log(g_h / p_h)
    87.  
    88.         return image, np.array((t_x, t_y, t_w, t_h))
    89.  
    90.     def __len__(self):
    91.         return len(self.box_list)
    92.  
    93.     def get_bndbox(self, bndboxes, positive):
    94.         """
    95.         返回和positive的IoU最大的标注边界框
    96.         :param bndboxes: 大小为[N, 4]或者[4]
    97.         :param positive: 大小为[4]
    98.         :return: [4]
    99.         """
    100.  
    101.         if len(bndboxes.shape) == 1:
    102.             # 只有一个标注边界框，直接返回即可
    103.             return bndboxes
    104.         else:
    105.             scores = util.iou(positive, bndboxes)
    106.             return bndboxes[np.argmax(scores)]
    107.  
    108.  
    109. def test():
    110.     """
    111.     创建数据集类实例
    112.     """
    113.     transform = transforms.Compose([
    114.         transforms.ToPILImage(),
    115.         transforms.Resize((227, 227)),
    116.         transforms.RandomHorizontalFlip(),
    117.         transforms.ToTensor(),
    118.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    119.     ])
    120.  
    121.     data_root_dir = '../../data/bbox_regression'
    122.     data_set = BBoxRegressionDataset(data_root_dir, transform=transform)
    123.  
    124.     print(data_set.__len__())
    125.     image, target = data_set.__getitem__(10)
    126.     print(image.shape)
    127.     print(target)
    128.     print(target.dtype)
    129.  
    130.  
    131. def test2():
    132.     """
    133.     测试DataLoader使用
    134.     """
    135.     transform = transforms.Compose([
    136.         transforms.ToPILImage(),
    137.         transforms.Resize((227, 227)),
    138.         transforms.RandomHorizontalFlip(),
    139.         transforms.ToTensor(),
    140.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    141.     ])
    142.  
    143.     data_root_dir = '../../data/bbox_regression'
    144.     data_set = BBoxRegressionDataset(data_root_dir, transform=transform)
    145.     data_loader = DataLoader(data_set, batch_size=128, shuffle=True, num_workers=8)
    146.  
    147.     items = next(data_loader.__iter__())
    148.     datas, targets = items
    149.     print(datas.shape)
    150.     print(targets.shape)
    151.     print(targets.dtype)
    152.  
    153.  
    154. if __name__ == '__main__':
    155.     test()
    156.     # test2()

    训练边界框回归

    执行bbox_regression.py脚本

    ①调用custom_bbox_regression_dataset.py脚本,得到训练数据data_loader

    ②使用AlexNet网络模型,修改分类器对象classifier的输出为2类(1类是car,一类是背景),加载权重best_linear_svm_alexnet_car.pth,并设置参数冻结,然后再添加一个线性层作为全连接层,定义损失函数为均方误差损失函数,优化函数为SGD,学习率更新器为StepLR,然后开始训练,保存模型到'models/bbox_regression_%d.pth'

    以下是bbox_regression.py脚本代码

    1. # -*- coding: utf-8 -*-
    2.  
    3. """
    4. @date: 2020/4/3 下午6:55
    5. @file: bbox_regression.py
    6. @author: zj
    7. @description: 边界框回归训练
    8. """
    9.  
    10. import os
    11. import copy
    12. import time
    13. import torch
    14. import torch.nn as nn
    15. import torch.optim as optim
    16. from torch.utils.data import DataLoader
    17. import torchvision.transforms as transforms
    18. from torchvision.models import AlexNet
    19.  
    20. from utils.data.custom_bbox_regression_dataset import BBoxRegressionDataset
    21. import utils.util as util
    22.  
    23.  
    24. def load_data(data_root_dir):
    25.     transform = transforms.Compose([
    26.         transforms.ToPILImage(),
    27.         transforms.Resize((227, 227)),
    28.         transforms.RandomHorizontalFlip(),
    29.         transforms.ToTensor(),
    30.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    31.     ])
    32.  
    33.     data_set = BBoxRegressionDataset(data_root_dir, transform=transform)
    34.     data_loader = DataLoader(data_set, batch_size=128, shuffle=True, num_workers=8)
    35.  
    36.     return data_loader
    37.  
    38.  
    39. def train_model(data_loader, feature_model, model, criterion, optimizer, lr_scheduler, num_epochs=25, device=None):
    40.     since = time.time()
    41.  
    42.     model.train()  # Set model to training mode
    43.     loss_list = list()
    44.     for epoch in range(num_epochs):
    45.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    46.         print('-' * 10)
    47.  
    48.         running_loss = 0.0
    49.  
    50.         # Iterate over data.
    51.         for inputs, targets in data_loader:
    52.             inputs = inputs.to(device)
    53.             targets = targets.float().to(device)
    54.  
    55.             features = feature_model.features(inputs)
    56.             features = torch.flatten(features, 1)
    57.  
    58.             # zero the parameter gradients
    59.             optimizer.zero_grad()
    60.  
    61.             # forward
    62.             outputs = model(features)
    63.             loss = criterion(outputs, targets)
    64.  
    65.             loss.backward()
    66.             optimizer.step()
    67.  
    68.             # statistics
    69.             running_loss += loss.item() * inputs.size(0)
    70.             lr_scheduler.step()
    71.  
    72.         epoch_loss = running_loss / data_loader.dataset.__len__()
    73.         loss_list.append(epoch_loss)
    74.  
    75.         print('{} Loss: {:.4f}'.format(epoch, epoch_loss))
    76.  
    77.         # 每训练一轮就保存
    78.         util.save_model(model, './models/bbox_regression_%d.pth' % epoch)
    79.  
    80.     print()
    81.  
    82.     time_elapsed = time.time() - since
    83.     print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    84.  
    85.     return loss_list
    86.  
    87.  
    88. def get_model(device=None):
    89.     # 加载CNN模型
    90.     model = AlexNet(num_classes=2)
    91.     model.load_state_dict(torch.load('./models/best_linear_svm_alexnet_car.pth'))
    92.     model.eval()
    93.  
    94.     # 取消梯度追踪
    95.     for param in model.parameters():
    96.         param.requires_grad = False
    97.     if device:
    98.         model = model.to(device)
    99.  
    100.     return model
    101.  
    102.  
    103. if __name__ == '__main__':
    104.     data_loader = load_data('./data/bbox_regression')
    105.  
    106.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    107.     feature_model = get_model(device)
    108.  
    109.     # AlexNet最后一个池化层计算得到256*6*6输出
    110.     in_features = 256 * 6 * 6
    111.     out_features = 4
    112.     model = nn.Linear(in_features, out_features)
    113.     model.to(device)
    114.  
    115.     criterion = nn.MSELoss()
    116.     optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4)
    117.     lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
    118.  
    119.     loss_list = train_model(data_loader, feature_model, model, criterion, optimizer, lr_scheduler, device=device,
    120.                             num_epochs=12)
    121.     util.plot_loss(loss_list)

    汽车car目标检测器实现

    读取图片,先检测图片中是否有汽车,然后使用非极大值抑制(NMS)算法消除冗余边界框,最后输出目标检测结果,如下图

    工程下载

    pytorch-r-cnn工程文件

  • 相关阅读:
    国内的软件测试真的这么不受待见吗？
    C#Lambda让代码变得更加简洁而优雅
    【Linux API 揭秘】container_of函数详解
    怎样在Windows10系统中安装配置PL/SQL
    数据分析常见的业务面试题
    阿里巴巴API接口解析，实现获得商品详情
    上6休3上3休2……这烧脑的调休安排，国内外AI都算不明白，集体大“翻车”！...
    光电柴微电网日前调度报告
    【电梯控制系统】基于VHDL语言和状态机实现的电梯控制系统的设计，使用了状态机
    Java的垃圾回收机制详解——从入门到出土，学不会接着来砍我！
  • 原文地址：https://blog.csdn.net/Victor_Li_/article/details/134274482