钢铁表面缺陷检测

Halcom

钢铁表面缺陷检测:
数据：https://www.kaggle.com/c/severstal-steel-defect-detection/data

1. from pathlib import Path
   
2. import pandas as pd
   
3. import numpy as np
   
4. import matplotlib.pyplot as plt
   
5. import cv2
   
6. 
   
7. class RLE:
   
8.     """
   
9.     Encapsulates run-length-encoding functionality.
   
10.     """
    
11. 
    
12.     MASK_H = 256
    
13.     MASK_W = 1600
    
14. 
    
15.     @classmethod
    
16.     def from_str(cls, s):
    
17.         if s != s:
    
18.             return cls()
    
19.         list_ = [int(n) for n in s.split(' ')]
    
20.         return cls.from_list(list_)
    
21. 
    
22.     @classmethod
    
23.     def from_mask(cls, mask):
    
24.         pixels = mask.T.flatten()
    
25.         pixels = np.concatenate([[0], pixels, [0]])
    
26.         runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    
27.         runs[1::2] -= runs[::2]
    
28.         return cls.from_list(runs)
    
29. 
    
30.     @classmethod
    
31.     def from_list(cls, list_):
    
32.         n_items = int(len(list_) / 2)
    
33.         items = np.zeros((n_items, 2), dtype=np.uint64)
    
34.         for i in range(n_items):
    
35.             items[i, 0] = list_[i * 2]
    
36.             items[i, 1] = list_[i * 2 + 1]
    
37.         return cls(items)
    
38. 
    
39.     def __init__(self, items=np.zeros((0, 0))):
    
40.         self._items = items
    
41. 
    
42.     @property
    
43.     def items(self):
    
44.         return self._items
    
45. 
    
46.     def __iter__(self):
    
47.         for idx, item in enumerate(self.items):
    
48.             yield (item[0], item[1])  # run, length
    
49. 
    
50.     def __len__(self):
    
51.         return self.items.shape[0]
    
52. 
    
53.     def to_mask(self):
    
54.         mask = np.zeros(self.MASK_H * self.MASK_W, dtype=np.uint8)
    
55.         for run, length in self:
    
56.             run = int(run - 1)
    
57.             end = int(run + length)
    
58.             mask[run:end] = 1
    
59.         return mask.reshape(self.MASK_H, self.MASK_W, order='F')
    
60. 
    
61.     def to_str_list(self):
    
62.         list_ = []
    
63.         for run, length in self:
    
64.             list_.append(str(run))
    
65.             list_.append(str(length))
    
66.         return list_
    
67. 
    
68.     def __str__(self):
    
69.         if len(self) == 0:
    
70.             return ''
    
71.         return ' '.join(self.to_str_list())
    
72. 
    
73.     def __repr__(self):
    
74.         return self.__str__()
    
75.    
    
76. def make_mask(labels):
    
77.     masks = np.zeros((256, 1600, 4), dtype=np.float32)
    
78.     for c, label in enumerate(labels.values):
    
79.         rle = RLE.from_str(label)
    
80.         masks[:, :, c] = rle.to_mask()
    
81.     return masks
    
82. 
    
83. img_size = 256 * 1600
    
84. 
    
85. train_csv = 'F:/sck/steel/train.csv'
    
86. df = pd.read_csv(train_csv)
    
87. df.head()
    
88. N_CLASSES = 4
    
89. rle_cols = [f'rle{i}' for i in range(N_CLASSES)]
    
90. df['ImageId'], df['ClassId'] = zip(*df['ImageId_ClassId'].str.split('_'))
    
91. df['ClassId'] = df['ClassId'].astype(int)
    
92. df = df.pivot(index='ImageId', columns='ClassId', values='EncodedPixels')
    
93. df.columns = rle_cols
    
94. df['defects'] = df.count(axis=1)
    
95. df = df.loc[df.defects > 0, :]
    
96. df.head()
    
97. 
    
98. r1 = df.loc['0007a71bf.jpg'][rle_cols]
    
99. r1
    
100. 
     
101. m1 = make_mask(r1)[:, :, 2]
     
102. plt.imshow(m1)
     
103. 
     
104. for i in range(len(df)):
     
105.     name=df.index[i]
     
106.     ri = df.loc[name][rle_cols]
     
107.     mi = make_mask(ri)
     
108.     cv2.imwrite('F:/sck/steel/fourclass/1/'+name[0:-4]+'.png', mi[:,:,0])
     
109.     cv2.imwrite('F:/sck/steel/fourclass/2/'+name[0:-4]+'.png', mi[:,:,1])
     
110.     cv2.imwrite('F:/sck/steel/fourclass/3/'+name[0:-4]+'.png', mi[:,:,2])
     
111.     cv2.imwrite('F:/sck/steel/fourclass/4/'+name[0:-4]+'.png', mi[:,:,3])
     
112.     if(mi[:,:,0].max()>1):
     
113.         print('11')
     
114.     if(mi[:,:,1].max()>1):
     
115.         print('11')
     
116.     if(mi[:,:,2].max()>1):
     
117.         print('11')

复制代码

参考：【1】https://www.kaggle.com/lightforever/severstal-mlcomp-catalyst-infer-0-90672
【2】https://blog.csdn.net/HRLSW/article/details/102779214
【3】AI Studio优质数据集
【4】https://aistudio.baidu.com/aistudio/datasetdetail/10941