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

人工智能等各类编程培训资料整理，所有资源无秘无压缩-购买会员

介绍

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

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

使用 Yolv5 进行预测

这是模型的最终结果。

学习目标

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

收集数据集

使用 Mask RCNN 进行实例分割

2.1 导入库

2.2 划分数据集

2.3 创建一个 Scratch 类

2.4 数据增强

2.5 创建模型

2.6 训练模型

通过 Yolov5 进行目标检测

3.1 数据标注

3.2 训练

3.3 结果

结论和要点

收集数据集

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

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

#makearequesttotheurl

r=requests.get(url)

#createoursoup

soup=BeautifulSoup(r.text,html.parser)

print(soup.title.text)

images=soup.find_all(img)

forimageinimages[-1]:

name=image[alt]

link=image[src]

withopen(name.replace(,-).replace(/,)+.jpg,wb)asf:

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 算法所需的所有库。

#importinglibraries

importpandasaspd

importnumpyasnp

importcv2

importos

importre

fromPILimportImage

importalbumentationsasA

fromalbumentations.pytorch.transformsimportToTensorV2

importtorch

importtorchvision

fromtorchvision.models.detection.faster_rcnnimportFastRCNNPredictor

fromtorchvision.models.detectionimportFasterRCNN

fromtorchvision.models.detection.rpnimportAnchorGenerator

fromtorch.utils.dataimportDataLoader,Dataset

fromtorch.utils.data.samplerimportSequentialSampler

frommatplotlibimportpyplotasplt

划分数据集

此处使用的数据采用 .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]

#validandtraindf

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 类，它转换我们的数据集并返回所需的参数。

classScratchDataset(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)

#thereisonlyoneclass

labels=torch.ones((records.shape[0],),dtype=torch.int64)

#supposeallinstancesarenotcrowd

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

ifself.transforms:

sample={

image:image,

bboxes:target[boxes],

labels:labels

}

sample=self.transforms(**sample)

image=sample[image]

target[boxes]=torch.tensor(sample[bboxes])

returnimage,target,image_id

def__len__(self)->int:

returnself.image_ids.shape[0]

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

数据扩充

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

#Albumenations

defget_train_transform():

returnA.Compose([

A.Flip(0.5),

ToTensorV2(p=1.0)

],bbox_params={format:pascal_voc,label_fields:[labels]})

defget_valid_transform():

returnA.Compose([

ToTensorV2(p=1.0)

],bbox_params={format:pascal_voc,label_fields:[labels]})

创建模型

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

#loadamodelpre-trainedonCOCO

model=torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

num_classes=2#1classscratch+background

#getnumberofinputfeaturesfortheclassifier

in_features=model.roi_heads.box_predictor.cls_score.in_features

#replacethepre-trainedheadwithanewone

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

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

classAverager:

def__init__(self):

self.current_total=0.0

self.iterations=0.0

defsend(self,value):

self.current_total+=value

self.iterations+=1

@property

defvalue(self):

ifself.iterations==0:

return0

else:

return1.0*self.current_total/self.iterations

defreset(self):

self.current_total=0.0

self.iterations=0.0

defcollate_fn(batch):

returntuple(zip(*batch))

train_dataset=WheatDataset(train_df,DIR_TRAIN,get_train_transform())

valid_dataset=WheatDataset(valid_df,DIR_TRAIN,get_valid_transform())

#splitthedatasetintrainandtestset

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)iftorch.cuda.is_available()elsetorch.device(cpu)

images,targets,image_ids=next(iter(train_data_loader))

model.to(device)

params=[pforpinmodel.parameters()ifp.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

forepochinrange(num_epochs):

loss_hist.reset()

forimages,targets,image_ids,intrain_data_loader:

images=list(image.to(device)forimageinimages)

targets=[{k:v.to(device)fork,vint.items()}fortintargets]

loss_dict=model(images,targets)

losses=sum(lossforlossinloss_dict.values())

loss_value=losses.item()

loss_hist.send(loss_value)

optimizer.zero_grad()

losses.backward()

optimizer.step()

ifitr%50==0:

print(fIteration#{itr}loss:{loss_value})

itr+=1

#updatethelearningrate

iflr_schedulerisnotNone:

lr_scheduler.step()

print(fEpoch#{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。

!cdyolov5&&pythontrain.py–img320–batch16–epochs50–datacarScr_up.yaml–weightslast.pt

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

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

!cdyolov5&&pythontrain.py–img320–batch32–epochs10–datacarScr_up.yaml–weightsyolov5s.pt–cache–evolve

结果

ExpprecisionrecallmAP_0.51个0.0030.5110.0012个0.6590.3110.3633个0.6240.5360.5124个0.5720.6100.519

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

使用 Yolov5 进行预测

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

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

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

结论

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

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

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

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