[人工智能]基于语义分割的钢材表面缺陷检测算法

• 项目名称： 基于语义分割的钢材表面缺陷检测算法

• 学习思路：

• 数据集的整理与分析

• 常用语义分割模型的训练与性能对比

• 模型的改进与性能提升

• 编程环境： Python3.7、Pytorch1.9、Google Colab

• 源码链接： https://colab.research.google.com/drive/1pUzWIJf2Sezm0ziuZF2h4PwU1tPcv3TE?usp=sharing

• 2021年7月26日学习内容：

• 钢材表面缺陷数据集的整理

• 数据集的可视化

1.导入模块

# 获取权限：访问云端硬盘中的所有文件
from os.path import join
from google.colab import drive

ROOT = "/content/drive"
drive.mount(ROOT)

!pip install albumentations==0.4.6
import albumentations 
from albumentations.pytorch import ToTensorV2

import os
import time
from typing import Optional, Tuple
import torch
from torch.utils.data import Dataset, random_split, DataLoader
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import pandas as pd
import numpy as np
import cv2
from albumentations import HorizontalFlip, VerticalFlip, RandomBrightness,  ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise, ElasticTransform
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix
from tqdm.notebook import tqdm
!pip install segmentation_models_pytorch
import segmentation_models_pytorch as smp # https://github.com/qubvel/segmentation_models.pytorch
%matplotlib inline

2.数据集

Kaggle中谢韦尔钢铁公司提供的带钢缺陷数据集，该数据集中提供了四种类型的带钢表面缺陷。训练集共有12568张，测试集5506张。图像尺寸为1600×256。

2.1 加载数据集

在这里插入代码片# 数据集原始地址：https://www.kaggle.com/c/severstal-steel-defect-detection/data
!git clone https://github.com/JasonYangCode/Data_set_04

# loading dataset
DATA_DIR = '/content/Data_set_04'

TRAIN_IMG_DIR = DATA_DIR + '/train_images' # Contains training images
TEST_IMG_DIR = DATA_DIR + '/test_images'  # Contains test images

TRAIN_CSV = DATA_DIR + '/train.csv'  # Contains real labels for training images
TEST_CSV = DATA_DIR + '/sample_submission.csv'   # Contains dummy labels for test image
TRAINED_WEIGHTS = './model.pth'     # Contains trained weights
SUBMISSION_FILE = './submission.csv'     # Contains trined labels for test image
# PRE_TRAINED_WEIGHTS = '../input/asigment5-zero-to-gans-novikov-ms-model/model.pth'

train_df = pd.read_csv(TRAIN_CSV)
test_df = pd.read_csv(TEST_CSV)

# 计算每张图片的masks的数量

def survey(df_in: pd.DataFrame):

    df_maskCnt = pd.DataFrame({'maskCount' : df_in.groupby('ImageId').size()})
    df_out = pd.merge(df_in, df_maskCnt, on='ImageId')
    df_out = df_out.sort_values(by=['maskCount', 'ImageId'], ascending=False)

    df_out['ClassIds'] = pd.Series(dtype=object)

    for i, row_i in df_out.iterrows():
        ClassId_list = []
        for j, row_j in df_out.loc[df_out['ImageId'] == row_i['ImageId']].iterrows():
            ClassId_list.append(row_j['ClassId'])
        df_out.at[i,'ClassIds'] = ClassId_list
    # print(df_out.head())    
    return df_out

df = survey(train_df)

# 将ClassId转换为masks中的索引
def class_id2index(val: int):
    return int(val-1)

# 将masks中的索引转换为ClassId
def index2class_id(val: int):
    """ converts index to ClassId in masks"""
    return int(val+1)
    
# 从 df 返回 num classId 的总数和列表
def counter_func(df_in):
    length = 4
    counter = np.zeros(length, dtype=int)
    total = 0
    for i in range(length):
        try:
            index = class_id2index(df_in.index[i])
            counter[index] = df_in.iloc[i, 0]  
        except:
            continue
        
        
    total = counter.sum()
    return total, counter 

mask_count_df_pivot = pd.DataFrame({'ClassCount' : df.groupby('ImageId').size()})
mask_count_df_pivot = pd.DataFrame({'Num' : mask_count_df_pivot.groupby('ClassCount').size()})
mask_count_df_pivot.sort_values('ClassCount', ascending=True, inplace=True)
ClassId_count_df = df.set_index(["ImageId", "ClassId"]).count(level='ClassId')

total, counter = counter_func(ClassId_count_df)
print('Total strings: {0}, 1 class: {1}, 2 class: {2}, 3 class: {3}, 4 class: {4}'.format(total, *counter))
total, counter = counter_func(mask_count_df_pivot)
print('Total images: {0}, one class: {1}, two classes: {2}, three classes: {3}, four classes: {4}'.format(total, *counter))

# https://www.kaggle.com/paulorzp/rle-functions-run-lenght-encode-decode

# img：numpy 数组，1 - 掩码，0 - 背景，以字符串格式返回运行长度（起始长度）
def mask2rle(img: np.array):

    pixels = img.T.flatten()
    pixels = np.concatenate([[0], pixels, [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    rle = ' '.join(str(x) for x in runs)
    if rle == np.nan:
        return ''
    else:
        return rle #returns a string formated (start length)


def rle2mask(mask_rle: str, input_shape: Tuple[int, int, int]=(256,1600,1)):
    
    """    
    img
    The pixel order of the line is from top to bottom from the left vertical line.
     It must be made and handed over.
     width/height should be made [height, width, ...] to fit the rows and columns.

     example when width=4, height=3
    
    s = [1,2,3,4,5,6,7,8,9,10,11,12]
        => 1,2,3 First row on the left, second row on 4,5,6
    
    mask_rle: run-length as string formated (start length)
    shape: (height,width)!!!  of array to return
    Returns numpy array, 1 - mask, 0 - background
    """
    
    # print('>>> rle2mask rle2mask: ', mask_rle)
    height, width = input_shape[:2]



    mask = np.zeros(width * height, dtype=np.uint8)
    if mask_rle is not np.nan:
        s = mask_rle.split()
        array = np.asarray([int(x) for x in s])
        starts = array[0::2]
        lengths = array[1::2]

        for index, start in enumerate(starts):
            begin = int(start - 1)
            end = int(begin + lengths[index])
            mask[begin : end] = 1

    rle_mask = mask.reshape(width, height).T
    # print('>>> rle2mask mask.shape: ', rle_mask.shape)
    # print('>>> rle2mask mask: ', rle_mask)
    return rle_mask

# Test RLE functions
assert mask2rle(rle2mask(df['EncodedPixels'].iloc[0]))==df['EncodedPixels'].iloc[0]
assert mask2rle(rle2mask('1 1'))=='1 1'


def build_masks(rle_labels: pd.DataFrame, input_shape: Tuple[int, int, int]=(256, 1600, 4)):
 
    masks = np.zeros(input_shape)
    # print('>>> build_masks')
    # print('>>> rle_labels.head()', rle_labels.head())
    for _, val in rle_labels.iterrows():
        # print('>>> build_masks: val[ImageId] :', val['ImageId'])
        # print('>>> build_masks: val[ClassId] :', val['ClassId'])
        masks[:, :, class_id2index(val['ClassId'])] = rle2mask(val['EncodedPixels'], input_shape)
        # print(">>> build_masks: masks ", masks.shape)
        # print('<<< build_masks')
    return masks #(256, 1600, 4)

def make_mask(row_id_in: int, df_in: pd.DataFrame,  input_shape_in: Tuple[int, int, int] = (256, 1600, 4)):

    fname = df_in.iloc[row_id_in].ImageId
    # print('>>>make_mask row_id_in: ', row_id_in, ' fname: ', fname)
    rle_labels = df_in[df_in['ImageId'] == fname][['ImageId', 'ClassId', 'EncodedPixels']]
    # print('>>>make_mask rle_labels', rle_labels)
    masks = build_masks(rle_labels, input_shape=input_shape_in) #(256, 1600, 4)
    # print('>>>make_mask masks: {0}'.format(masks))
    return fname, masks
assert mask2rle(np.zeros((256, 1600), np.float32)) == ''

2.2数据集可视化

# 可视化
def show_images(df_in: pd.DataFrame, img_dir: str, trained_df_in: pd.DataFrame = None):

    local_df = df_in
    local_trained_df = trained_df_in
    columns = 1
    if type(trained_df_in) == pd.DataFrame:
        rows = 15
    else:
        rows = 10
    
    
    fig = plt.figure(figsize=(20,80))
    
    def sorter(local_df):
        local_df = local_df.sort_values(by=['maskCount', 'ImageId'], ascending=False) # To show as many valid masks as possible
        grp = local_df['ImageId'].drop_duplicates()[0:rows]
        return grp

    ax_idx = 1
    for filename in sorter(df_in):
        if ax_idx > rows * columns * 2:
            break

        subdf = local_df[local_df['ImageId'] == filename].reset_index() # chose file with masks

        fig.add_subplot(rows * 2, columns, ax_idx).set_title(filename)
        img = cv2.imread(os.path.join(img_dir, filename ))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_2 = cv2.imread(os.path.join(img_dir, filename ))
        img_2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        plt.imshow(img)

        # 显示来自第一个df的mask  
        ax_idx += 1
        fig.add_subplot(rows * 2, columns, ax_idx).\
            set_title(filename + 'Highlighted defects ClassIds:' + str(subdf['ClassIds'][0]))

        colors = [(255, 51, 51),(255,255,51), (51,255,51), (51,51,255)]
        masks = build_masks(subdf, (256, 1600, 4)) # get masks (256, 1600, 4)
        masks_len = masks.shape[2] # get 4
        
        for i in range(masks_len):
            img[masks[:, :, i] == 1] = colors[i]

        plt.imshow(img)
        ax_idx += 1


        # 显示来自第二个df的mask
        if type(trained_df_in) == pd.DataFrame:
            subdf_trained = local_trained_df[local_trained_df['ImageId'] == filename].reset_index() # chose file with masks
            fig.add_subplot(rows * 2, columns, ax_idx).\
                set_title('Trained  '+ filename + ' Highlighted defects ClassIds: ' + str(subdf['ClassIds'][0]))

            colors = [(204, 51, 51),(204,204,51), (51,204,51), (51,51,204)]
            masks = build_masks(subdf_trained, (256, 1600, 4)) # get masks (256, 1600, 4)
            masks_len = masks.shape[2] # get 4

            for i in range(masks_len):
                img_2[masks[:, :, i] == 1] = colors[i]
                
            plt.imshow(img_2)
            ax_idx += 1
        
    print("Class 1 = Red","Class 2 = Yellow","Class 3 = Green","Class 4 = Blue", sep='\n')
    plt.show()
    
show_images(df, TRAIN_IMG_DIR)