镜头缺陷检测
1 项目说明
摄像头模组是智能手机最为重要的组成部分之一。随着智能手机行业的快速发展,摄像头模组的需求量增加。高像素摄像头的出现,对模组检测精度要求提出了新的挑战。
项目中以手机镜头为例,向大家介绍如何快速使用实例分割方式进行缺陷检测。
2 数据准备
数据集中包含了992张已经标注好的数据,标注形式为MSCOCO的实例分割格式。点击此处下载数据集
- 数据切分将训练集、验证集和测试集按照7:2:1的比例划分。
paddlex --split_dataset --format COCO --dataset_dir dataset --val_value 0.2 --test_value 0.1
数据文件夹切分前后的状态如下:
dataset/ dataset/ ├── JPEGImages/ --> ├── JPEGImages/ ├── annotations.json ├── annotations.json ├── test.json ├── train.json ├── val.json
3 模型选择
PaddleX提供了丰富的视觉模型,在实例分割中提供了MaskRCNN系列模型.在本项目中采用Mask-RCNN算法
4 模型训练
在项目中,我们采用Mask-RCNN作为镜头缺陷检测的模型。 运行如下代码开始训练模型:
python code/train.py
若输入如下代码,则可在log文件中查看训练日志
python code/train.py > log
- 训练过程说明
# 定义训练和验证时的transforms # API说明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/transforms/operators.py train_transforms = T.Compose([ T.RandomResizeByShort( short_sizes=[640, 672, 704, 736, 768, 800], max_size=1333, interp='CUBIC'), T.RandomHorizontalFlip(), T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) eval_transforms = T.Compose([ T.ResizeByShort( short_size=800, max_size=1333, interp='CUBIC'), T.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])
# 定义训练和验证所用的数据集 # API说明:https://github.com/PaddlePaddle/PaddleX/blob/develop/paddlex/cv/datasets/coco.py#L26 train_dataset = pdx.datasets.CocoDetection( data_dir='dataset/JPEGImages', ann_file='dataset/train.json', # num_workers=0, # 注意:若运行时报错则添加该句 transforms=train_transforms, shuffle=True) eval_dataset = pdx.datasets.CocoDetection( data_dir='dataset/JPEGImages', ann_file='dataset/val.json', # num_workers=0, # 注意:若运行时报错则添加该句 transforms=eval_transforms)
# 初始化模型,并进行训练 # 可使用VisualDL查看训练指标,参考https://github.com/PaddlePaddle/PaddleX/tree/release/2.0-rc/tutorials/train#visualdl可视化训练指标 num_classes = len(train_dataset.labels) model = pdx.det.MaskRCNN( num_classes=num_classes, backbone='ResNet50', with_fpn=True)
# API说明:https://github.com/PaddlePaddle/PaddleX/blob/release/2.0-rc/paddlex/cv/models/detector.py#L155 # 各参数介绍与调整说明:https://paddlex.readthedocs.io/zh_CN/develop/appendix/parameters.html model.train( num_epochs=12, train_dataset=train_dataset, train_batch_size=1, eval_dataset=eval_dataset, learning_rate=0.00125, lr_decay_epochs=[8, 11], warmup_steps=10, warmup_start_lr=0.0, save_dir='output/mask_rcnn_r50_fpn', use_vdl=True)
5 训练可视化
在模型训练过程,在train
函数中,将use_vdl
设为True,则训练过程会自动将训练日志以VisualDL的格式打点在save_dir
(用户自己指定的路径)下的vdl_log
目录。
用户可以使用如下命令启动VisualDL服务,查看可视化指标
visualdl --logdir output/mask_rcnn_r50_fpn/vdl_log --port 8001
服务启动后,按照命令行提示,使用浏览器打开 http://localhost:8001/
6 模型导出
模型训练处理被保存在了output文件夹,此时模型文件还是动态图文档,需要导出成静态图的模型才可以进一步部署预测,运行如下命令,会自动在output文件夹下创建一个inference_model
的文件夹,用来存放预测好的模型。
paddlex --export_inference --model_dir=output/mask_rcnn_r50_fpn/best_model --save_dir=output/inference_model
7 模型预测
运行如下代码:
python code/infer.py
文件内容如下:
import glob import numpy as np import threading import time import random import os import base64 import cv2 import json import paddlex as pdx image_name = 'dataset/JPEGImages/Image_370.jpg' model = pdx.load_model('output/mask_rcnn_r50_fpn/best_model') img = cv2.imread(image_name) result = model.predict(img) keep_results = [] areas = [] f = open('result.txt','a') count = 0 for dt in np.array(result): cname, bbox, score = dt['category'], dt['bbox'], dt['score'] if score < 0.5: continue keep_results.append(dt) count+=1 f.write(str(dt)+'\n') f.write('\n') areas.append(bbox[2] * bbox[3]) areas = np.asarray(areas) sorted_idxs = np.argsort(-areas).tolist() keep_results = [keep_results[k] for k in sorted_idxs] if len(keep_results) > 0 else [] print(keep_results) print(count) f.write("the total number is :"+str(int(count))) f.close() pdx.det.visualize(image_name, result, threshold=0.5, save_dir='./output/mask_rcnn_r50_fpn')
则可生成result.txt文件并显示预测结果图片,result.txt文件中会显示图片中每个检测框的位置、类别及置信度,并给出检测框的总个数.
预测结果如下: