使用 Mask RCNN 和 Yolov5 进行划痕检测

2023-04-20 10:15 779 阅读 ID：985
磐创AI
磐创AI
主页

本文的重点是检测汽车划痕，与针对不同类型产品的自主质量检测系统的开发同步。在停车场，这种检测为客户提供了他们的汽车安全的保证；此外，如果发生某些事情，检测系统将有助于对这种情况进行谨慎处理。

我将这个问题作为一个单类分类问题来解决，将凹痕、损坏和划痕视为划痕，并进一步在 Flask 的帮助下制作了一个基本应用程序。我将向你介绍我在做这个项目时获得的所有想法、代码、算法和知识，我将通过Mask RCNN和Yolov5实现这个项目。

使用 Yolv5 进行预测

这是模型的最终结果。

学习目标

  1. 了解如何使用Mask RCNN和Yolov5 执行 自定义对象检测。
  2. 在使用在 coco 数据集和 Resnet50 上训练的模型时利用迁移学习。
  3. 了解质量数据收集和数据注释的重要性是任何项目中不可或缺且最耗时的部分。

目录

  1. 收集数据集
  2. 使用 Mask RCNN 进行实例分割 2.1 导入库2.2 划分数据集2.3 创建一个 Scratch 类2.4 数据增强2.5 创建模型2.6 训练模型
  3. 通过 Yolov5 进行目标检测3.1 数据标注3.2 训练3.3 结果
  4. 结论和要点

收集数据集

为了收集数据，我制作了一个数据抓取器，使用 Beautiful Soup 从adobe、Istock photo等在线网站抓取数据。

url = 'https://stock.adobe.com/in/search/images?k=car%20scratch'

# make a request to the url
r = requests.get(url)

# create our soup
soup = BeautifulSoup(r.text, 'html.parser')

print(soup.title.text)
images = soup.find_all('img')

for image in images[-1]:
name = image['alt']
link = image['src']

with open(name.replace(' ', '-').replace('/', '') + '.jpg', 'wb') as f:
    im = requests.get(link)
    f.write(im.content)

由于网站关于抓取的隐私政策，大部分图像都没有被抓取。由于隐私问题，我直接从 Istock photo、Shutter photo 和 Adobe 下载了图像。

我从大约 80 张图像开始，增加到 350 张图像，并进一步增加到大约 900 张图像以进行最终注释。

使用 Mask RCNN 进行实例分割

图像分割是基于像素将图像分割成不同的区域。Mask RCNN 是一种用于实例分割的模型，它是图像分割的一种子类型，可在对象边界中分离实例。它是在 Faster RCNN 的基础上进一步构建的。Faster RCNN 对每个对象都有两个输出，分别是类标签和边界框偏移，Mask RCNN 是第三个输出的附加，即对象的掩码。

Mask RCNN 架构

Mask RCNN 的架构由以下部分组成：

  • 骨干网络
  • 区域提案网络
  • 掩码表示
  • 感兴趣区域RoI

使用 Mask RCNN 检测汽车划痕的优势在于，我们可以使用多边形而不仅仅是边界框，并在我们的目标上创建一个掩码，使我们能够以更准确和简洁的方式获得和可视化结果。

让我们开始使用 Mask RCNN 来实现我们的问题。

导入库

导入实施我们的 Mask RCNN 算法所需的所有库。

# importing libraries
import pandas as pd
import numpy as np
import cv2
import os
import re
from PIL import Image
import albumentations as A

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 matplotlib import pyplot as plt

划分数据集

此处使用的数据采用 .csv 格式，其中包含边界框的 x、y、w 和 h 坐标，而数据使用数据注释器 make-sense 进行注释。

make-sense：https://www.makesense.ai/

image_ids = train_df['image_id'].unique()
print(len(image_ids))
valid_ids = image_ids[-10:]
train_ids = image_ids[:-10]

# valid and train df
valid_df = train_df[train_df['image_id'].isin(valid_ids)]
train_df = train_df[train_df['image_id'].isin(train_ids)]

创建一个 Scratch 类

创建我们的 Scratch Dataset 类，它转换我们的数据集并返回所需的参数。

class ScratchDataset(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['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.tensor(sample['bboxes'])
        return image, target, image_id
    def __len__(self) -> int:
        return self.image_ids.shape[0]

这里'img_dir'，是保存图像的目录路径。

数据扩充

在这里，我们使用 Albumentations 进行数据扩充。

# Albumenations
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']})

创建模型

我们将使用 Resnet50 模型和 Mask RCNN。

# load a model pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2 # 1 class scratch+ background

# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features

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

让我们开始创建一个 Averager 类以及训练和验证数据加载器，这将成为训练我们模型时的关键组件。

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
def collate_fn(batch):
    return tuple(zip(*batch))
train_dataset = WheatDataset(train_df, DIR_TRAIN, get_train_transform())
valid_dataset = WheatDataset(valid_df, DIR_TRAIN, get_valid_transform())
# split the dataset in train and test set
indices = torch.randperm(len(train_dataset)).tolist()
train_data_loader = DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=False,
    num_workers=4,
    collate_fn=collate_fn
)
valid_data_loader = DataLoader(
    valid_dataset,
    batch_size=8,
    shuffle=False,
    num_workers=4,
    collate_fn=collate_fn
)

训练模型

我们正在激活“cuda”并访问可用的 GPU。weight_decay=0.0005，momentum=0.9，动态学习率从 0.05 开始。

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
images, targets, image_ids = next(iter(train_data_loader))
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)
lr_scheduler = None
num_epochs = 2
loss_hist = Averager()
itr=1
for epoch in range(num_epochs):
    loss_hist.reset()
    for images, targets, image_ids, in train_data_loader:
        images = list(image.to(device) for image in images)
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        loss_value = losses.item()
        loss_hist.send(loss_value)
        optimizer.zero_grad()
        losses.backward()
        optimizer.step()
        if itr % 50 == 0:
            print(f'Iteration #{itr} loss: {loss_value}')
        itr += 1
    # update the learning rate
    if lr_scheduler is not None:
        lr_scheduler.step()
    print(f'Epoch #{epoch} loss: {loss_hist.value}')

但我无法实现这一点，因为仅仅 80 张图片就花了 10 个小时甚至更多时间。

使用 Mask RCNN 进行自定义训练的时间复杂度是巨大的，而且你需要大量的计算能力，这对我来说是没有的。

我希望你有好的计算资源并且你可以实现它。

通过 Yolov5 进行目标检测

Yolo 主要用于对象检测，由 Ultralytics 发布，已成为基准算法，例如视觉数据分割。Yolov5 比 Yolov4 更快、更高效，并且可以很好地泛化到新图像。

Yolov5: https://github.com/ultralytics/yolov5/releases

Yolov5架构

该算法的工作原理如下：

  • 残差块
  • 边界框回归
  • 交并比（IOU）
  • 非极大值抑制 a

Yolov5 更快、更小，与以前的版本大致一样准确。在 coco 数据集上训练，它与边界框配合得很好。

让我们从问题案例中 Yolov5 的实现开始；我使用 google collab 来运行其中的代码。

数据标注

  • 我使用了一个有意义的数据注释器来注释数据集。
  • 当数据被精确注释时，即小而切中要点时，Yolo 不能很好地工作，因为它不能很好地泛化到小边界框。
  • 因此数据标注有点棘手，需要对区域进行统一标注。

训练

加载模型后，

model = torch.hub.load('ultralytics/yolov5','yolov5s')

我们添加了 Yolo 工作所需的 yaml 文件和数据（图像在一个文件夹中，而注释作为文本文件在另一个文件夹中）我们训练了我们的模型，批量大小为 16，图像大小为 320*320。

!cd yolov5 && python train.py --img 320 --batch 16 --epochs 50 --data carScr_up.yaml --weights last.pt

虽然在 Yolo 文档中，它声称要运行 300 个 epochs 才能获得良好的结果，但我们将其降低到 50 个 epochs，并且在超参数调整之后，我们的模型甚至在 30 个 epochs 内就开始表现不错。

对于超参数调优，我们使用 Yolo 提供的 evolution，其中数据经过 10 个 epoch 的训练，进行 300 次进化。

!cd yolov5 && python train.py --img 320 --batch 32 --epochs 10 --data carScr_up.yaml --weights yolov5s.pt --cache --evolve

结果

下图代表实验 4，每个实验都在不同数量的图像和注释上进行训练。汽车有划痕的预测如下：

使用 Yolov5 进行预测

在这种情况下，精度和召回率很小，因为使用 Yolo，我们正在处理边界框，这些指标取决于实际框和预测框的联合交集 (IOU)。

让我们看看用 Yolov5 训练我们的数据集 50 个时期后获得的指标

我们可以看到，在 20 个 epoch 之后，指标趋于稳定，可以看出 Yolo 可以非常快地学习这种关系并很好地泛化到我们的问题陈述，即使我们拥有的数据少于 1000 张图像。

结论

我们可以看到 Yolov5 和 Mask RCNN 对于我们的问题陈述非常有效，尽管我无法实现后者的代码。在使用 Yolov5 进行自定义训练时，排除指标，它能够非常好地预测，检测样本图像中的所有划痕和损坏。

我们在本文中学习了如何收集、注释和训练不同的模型，以及训练不同模型需要什么。

  • 在上面，我将损坏和划痕视为一个单独的类别。
  • 数据注释和收集是该解决方案不可或缺的部分。
  • 如果我们使用多边形并增加我们的数据集大小，我们肯定可以做得更好。

PS：这不适用于没有损坏的汽车。因为我们只在包含有划痕和损坏的汽车的数据上训练它。

 收起内容
免责声明：作者保留权利，不代表本站立场。如想了解更多和作者有关的信息可以查看页面右侧作者信息卡片。
    快抢沙发！
      加载更多评论
    反馈
    to-top--btn