Actual combat of industrial defect detection project (IV) -- ceramic defect detection based on hrnet

be based on HRNet Ceramic defect detection

1. principle ：
Refer to the articles of the big guys
HRNet: HRNet principle .

2. Data set preparation and code
Data download link ：https://aistudio.baidu.com/aistudio/datasetdetail/32615
Code download link ：https://gitee.com/wxyfmq123456/HRNet-Image-Classification?_from=gitee_search
3. Binarization of original graph
Here, the data set has provided a binary image png, We need to use png Image training . Because the original image is particularly inconspicuous .

in total 6 Categories .
4. Parameter configuration
(1) Data storage location ：

(2) Data storage method ：

Each folder represents a type , All the pictures inside are binary pictures （.png）, Original picture （.jpg） It can be deleted or backed up elsewhere .

(3) Modify the code

open cls_hrnet.py, modify

self.classifier = nn.Linear(2048, 1000)

by

self.classifier = nn.Linear(2048, 6)

That is, the following parameters are Number of categories .

(4) Select profile

Let's pick the first one , namely ：
cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml

Modify the parameters inside , take

DATASET:
  DATASET: 'imagenet'
  DATA_FORMAT: 'jpg'
  ROOT: 'data/imagenet/'
  TEST_SET: 'val'
  TRAIN_SET: 'train'

It is amended as follows

DATASET:
  DATASET: 'data'
  DATA_FORMAT: 'png'
  ROOT: 'imagenet'
  TEST_SET: 'val'
  TRAIN_SET: 'train'

The reason is that the path we set is different from the code .
The other parameters , Like the number of iterations epoch,bath_size etc. , You can adjust your parameters .

(5) stay train.py Add the following code to

    train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])
    )
 
 
    # The following is the added code , The above lines are the original code 
    #print(train_dataset.classes)  # The category is determined according to the name of the sub folder 
    with open("class.txt","w") as f1:
        for classname in train_dataset.classes:
            f1.write(classname + "\n")
 
    #print(train_dataset.class_to_idx) # Define indexes for these categories in order as 0,1...
    with open("classToIndex.txt", "w") as f2:
        for key, value in train_dataset.class_to_idx.items():
            f2.write(str(key) + " " + str(value) + '\n')
 
    #print(train_dataset.imgs) # Returns the path and category of pictures from all folders 

You can keep the corresponding index.

5. Training

python  tools/train.py --cfg  experiments/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml

After training output In the folder
Weight offset file ：final_state.pth.tar

6. test
Pictures of the test , Note that the verification set is still read here vaild, So to test a picture , We can turn the picture of the verification set into a , Put in, for example, a place called test In the folder of , The path is shown in the figure ：

stay HRNet-Image-Classification-master\lib\core\function.py Inside def validate function , add to

print('class:{}'.format(output.argmax(1)))

Print the identified category index.
function ：

python  tools/vaild.py --cfg  experiments/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml

The of the category will be printed index.