工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测

基于FPN_Tensorflow的PCB缺陷检测

这一篇实战,参考于:
基于深度学习的印刷电路板瑕疵识别.
因为介绍的数据集的预处理不是很详细,所以在这里再整理一遍,顺便学习新的深度学习网络模型。
一. 原理
Faster R-CNN.
深度学习网络 | FPN和fast rcnn及RPN网络是如何结合的细节.
二. 训练过程
1.数据准备
你可以从这里下载数据集:http://robotics.pkusz.edu.cn/resources/dataset/
我的已经将数据的训练数据和验证数据分配好了,这里按照从官网的数据集下载的来说。

下载好原始数据集的后,你会发现有这么多文件夹:
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第1张图片
Annotations文件夹:
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第2张图片

需要把各个文件夹的文件全部移动到Annotations/,想这样:
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第3张图片
同样,对images/,就这样:
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第4张图片
6个文件夹的图片全部移动出来,然后把空的文件夹删掉。

2.下载源代码

环境要求:注意版本匹配
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第5张图片

3.数据拆分
将自己的数据集拆分为训练集、测试集合:

python /data/io/divide_data.py

将会生成两个文件夹:
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第6张图片

4.修改配置

在/libs/label_name_dict/label_dict.py中添加自己的数据集标签类别。具体如下:

# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import

from libs.configs import cfgs

if cfgs.DATASET_NAME == 'ship':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'ship': 1
    }
elif cfgs.DATASET_NAME == 'aeroplane':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'aeroplane': 1
    }
elif cfgs.DATASET_NAME == 'WIDER':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'face': 1
    }
elif cfgs.DATASET_NAME == 'jyzdata':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'classone': 1,
        'boli': 2,
        'dangeboli': 3,
        'taoci': 4,
        'taoci2': 4
    }
elif cfgs.DATASET_NAME == 'icdar':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'text': 1
    }
elif cfgs.DATASET_NAME.startswith('DOTA'):
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'roundabout': 1,
        'tennis-court': 2,
        'swimming-pool': 3,
        'storage-tank': 4,
        'soccer-ball-field': 5,
        'small-vehicle': 6,
        'ship': 7,
        'plane': 8,
        'large-vehicle': 9,
        'helicopter': 10,
        'harbor': 11,
        'ground-track-field': 12,
        'bridge': 13,
        'basketball-court': 14,
        'baseball-diamond': 15
    }
elif cfgs.DATASET_NAME == 'pascal':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'aeroplane': 1,
        'bicycle': 2,
        'bird': 3,
        'boat': 4,
        'bottle': 5,
        'bus': 6,
        'car': 7,
        'cat': 8,
        'chair': 9,
        'cow': 10,
        'diningtable': 11,
        'dog': 12,
        'horse': 13,
        'motorbike': 14,
        'person': 15,
        'pottedplant': 16,
        'sheep': 17,
        'sofa': 18,
        'train': 19,
        'tvmonitor': 20
    }
elif cfgs.DATASET_NAME == 'pcb':
    NAME_LABEL_MAP = {
        'back_ground': 0,
        'missing_hole': 1,
        'mouse_bite': 2,
        'open_circuit': 3,
        'short': 4,
        'spur': 5,
        'spurious_copper': 6
    }
else:
    assert 'please set label dict!'


def get_label_name_map():
    reverse_dict = {}
    for name, label in NAME_LABEL_MAP.items():
        reverse_dict[label] = name
    return reverse_dict

LABEl_NAME_MAP = get_label_name_map()

然后就可以运行convert_data_to_tfrecord.py了。别忘了填上四个参数。

python convert_data_to_tfrecord.py --VOC_dir='***/VOCdevkit/VOCdevkit_train/' --save_name='train' --img_format='.jpg' --dataset='ship'

5.数据转换
下一步,需要分别将训练和测试数据转换成 TFrecord 格式。该格式是TensorFlow的高效存储格式,连续的内存二进制存储能够有效的加快数据的读取和写入。而像原始文件夹那样每个数据独立存储。

python convert_data_to_tfrecord.py --VOC_dir='***/data/pcb_test/' --save_name='test' --img_format='.jpg' --dataset='pcb'
python convert_data_to_tfrecord.py --VOC_dir='***/data/pcb_train/' --save_name='train' --img_format='.jpg' --dataset='pcb'
# 上一步训练、测试数据分离后的文件夹路径
--VOC_dir='你的数据路径' 

#tfrecord名字。训练数据用'train' ,测试数据用'test'
--save_name='train' 

#你的图片格式。有Png\tif\jpeg
--img_format='.jpg' 

#你的数据集名字。这个需要在/libs/label_name_dict/label_dict.py中添加自己的数据集
--dataset='mydataset

6.参数设置

对应里面的model,应先下载好预训练模型,放在 /data/pretrained_weights/ 中。

预训练模型的下载地址:
Res101: resnet_v1_101_2016_08_28.tar.gz.
Res50: resnet_v1_50_2016_08_28.tar.gz.

数据准备好了,那么训练就要调参和调一些路径啦。打开/libs/configs/cfgs.py。

# -*- coding: utf-8 -*-
from __future__ import division, print_function, absolute_import
import os
import tensorflow as tf
# ------------------------------------------------
VERSION = 'FPN_Res101_0117_OHEM'#-----------------NET_NAME = 'resnet_v1_101'#-----------------ADD_BOX_IN_TENSORBOARD = True

# ---------------------------------------- System_config
ROOT_PATH = os.path.abspath('/home/cg/机器视觉与机器学习/工业缺陷检测项目/基于FPN的PCB缺陷检测/Xia-Detection-for-PCB')#工程文件路径-----------------print (20*"++--")
print (ROOT_PATH)
GPU_GROUP = "2"
SHOW_TRAIN_INFO_INTE = 10
SMRY_ITER = 100
SAVE_WEIGHTS_INTE = 10000

SUMMARY_PATH = ROOT_PATH + '/output/summary'
TEST_SAVE_PATH = ROOT_PATH + '/tools/test_result'
INFERENCE_IMAGE_PATH = ROOT_PATH + '/tools/inference_image'
INFERENCE_SAVE_PATH = ROOT_PATH + '/tools/inference_results'

if NET_NAME.startswith("resnet"):
    weights_name = NET_NAME
elif NET_NAME.startswith("MobilenetV2"):
    weights_name = "mobilenet/mobilenet_v2_1.0_224"
else:
    raise NotImplementedError

PRETRAINED_CKPT = ROOT_PATH + '/data/pretrained_weights/' + weights_name + '.ckpt'
TRAINED_CKPT = os.path.join(ROOT_PATH, 'output/trained_weights')

EVALUATE_DIR = ROOT_PATH + '/output/evaluate_result_pickle/'
#test_annotate_path = '/home/yjr/DataSet/VOC/VOC_test/VOC2007/Annotations'
test_annotate_path = '/home/cg/机器视觉与机器学习/工业缺陷检测项目/基于FPN的PCB缺陷检测/pcb数据集/PCB_DATASET/Annotations/'#-----------------改

# ------------------------------------------ Train config
RESTORE_FROM_RPN = False
IS_FILTER_OUTSIDE_BOXES = False
FIXED_BLOCKS = 0  # allow 0~3
USE_07_METRIC = False

RPN_LOCATION_LOSS_WEIGHT = 1.
RPN_CLASSIFICATION_LOSS_WEIGHT = 1.0

FAST_RCNN_LOCATION_LOSS_WEIGHT = 1.0
FAST_RCNN_CLASSIFICATION_LOSS_WEIGHT = 1.0
RPN_SIGMA = 3.0
FASTRCNN_SIGMA = 1.0

MUTILPY_BIAS_GRADIENT = None   # 2.0  # if None, will not multipy
GRADIENT_CLIPPING_BY_NORM = None   # 10.0  if None, will not clip

EPSILON = 1e-5
MOMENTUM = 0.9
LR = 0.001  # 0.001  # 0.0003
#DECAY_STEP = [60000, 80000]  # 50000, 70000
DECAY_STEP = [10000, 20000]  # 50000, 70000
#MAX_ITERATION = 150000
MAX_ITERATION = 30000

# ------------------------------------------- Data_preprocess_config
DATASET_NAME = 'pcb'  # 'ship', 'spacenet', 'pascal', 'coco' #-----------------改
# PIXEL_MEAN = [123.68, 116.779, 103.939]  # R, G, B. In tf, channel is RGB. In openCV, channel is BGR
PIXEL_MEAN = [21.25, 85.936, 28.729]  #-----------------IMG_SHORT_SIDE_LEN =  600  # 600
IMG_MAX_LENGTH =  3000  # 1000
CLASS_NUM = 6 #-----------------改

# --------------------------------------------- Network_config
BATCH_SIZE = 1
INITIALIZER = tf.random_normal_initializer(mean=0.0, stddev=0.01)
BBOX_INITIALIZER = tf.random_normal_initializer(mean=0.0, stddev=0.001)
WEIGHT_DECAY = 0.00004 if NET_NAME.startswith('Mobilenet') else 0.0001

# ---------------------------------------------Anchor config
USE_CENTER_OFFSET = False

LEVLES = ['P2', 'P3', 'P4', 'P5', 'P6']
# BASE_ANCHOR_SIZE_LIST = [32, 64, 128, 256, 512]  # addjust the base anchor size for voc.
BASE_ANCHOR_SIZE_LIST = [15, 25, 40, 60, 80]  # addjust the base anchor size for voc.
#BASE_ANCHOR_SIZE_LIST = [8, 15, 25, 40, 60]
ANCHOR_STRIDE_LIST = [4, 8, 16, 32, 64]
ANCHOR_SCALES = [2., 3., 4.]
ANCHOR_RATIOS = [2., 3., 4., 5.]
# ANCHOR_SCALES = [1.0]
# ANCHOR_RATIOS = [0.5, 1., 2.0]
ROI_SCALE_FACTORS = [10., 10., 5.0, 5.0]
ANCHOR_SCALE_FACTORS = None

# --------------------------------------------FPN config
SHARE_HEADS = True
KERNEL_SIZE = 3
RPN_IOU_POSITIVE_THRESHOLD = 0.7
RPN_IOU_NEGATIVE_THRESHOLD = 0.3
TRAIN_RPN_CLOOBER_POSITIVES = False

RPN_MINIBATCH_SIZE = 256
RPN_POSITIVE_RATE = 0.5
RPN_NMS_IOU_THRESHOLD = 0.7
RPN_TOP_K_NMS_TRAIN = 12000
#RPN_MAXIMUM_PROPOSAL_TARIN = 2000
RPN_MAXIMUM_PROPOSAL_TARIN = 2000

RPN_TOP_K_NMS_TEST = 6000
RPN_MAXIMUM_PROPOSAL_TEST = 1000

# specific settings for FPN
# FPN_TOP_K_PER_LEVEL_TRAIN = 2000
# FPN_TOP_K_PER_LEVEL_TEST = 1000

# -------------------------------------------Fast-RCNN config
ROI_SIZE = 14
ROI_POOL_KERNEL_SIZE = 2
#USE_DROPOUT = False
USE_DROPOUT = True
KEEP_PROB = 1.0
SHOW_SCORE_THRSHOLD = 0.6  # only show in tensorboard

#FAST_RCNN_NMS_IOU_THRESHOLD = 0.3  # 0.6
FAST_RCNN_NMS_IOU_THRESHOLD = 0.3
FAST_RCNN_NMS_MAX_BOXES_PER_CLASS = 100
FAST_RCNN_IOU_POSITIVE_THRESHOLD = 0.5
FAST_RCNN_IOU_NEGATIVE_THRESHOLD = 0.0   # 0.1 < IOU < 0.5 is negative
FAST_RCNN_MINIBATCH_SIZE = 256  # if is -1, that is train with OHEM
# FAST_RCNN_MINIBATCH_SIZE = -1
FAST_RCNN_POSITIVE_RATE = 0.25

#ADD_GTBOXES_TO_TRAIN = False
ADD_GTBOXES_TO_TRAIN = True

注意配置好路径,和参数。
7.训练

python /tools/train.py 

等着训练完成吧。

8.验证测试

cd $PATH_ROOT/tools
python inference.py --data_dir='/PATH/TO/IMAGES/' 
                    --save_dir='/PATH/TO/SAVE/RESULTS/' 
                    --GPU='0'

训练完之后会在output/summary文件夹里,生成
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第7张图片

9.效果

发现效果根训练的迭代次数有很大关系,迭代迭代次数越多,训练时间越长,一开始训练了两天时间,30000的迭代值。
工业缺陷检测项目实战(三)——基于FPN_Tensorflow的PCB缺陷检测_第8张图片

需要源代码包的私信我就可啦。期待下一次的实战项目。

你可能感兴趣的:(机器学习与深度学习,ubuntu,算法,计算机视觉)