kwflyer / crab

Attention-Based Convolution Neural Network for Magnetic Tile Surface Defect Classification and Detection (convolutional retinal attention block:CRAB)

Ju Li，Kai Wang, Mengfan He, Luyao Ke, and Heng Wang. Attention-Based Convolution Neural Network for Magnetic Tile Surface Defect Classification and Detection[J]. Applied Soft Computing, 2024: 111631. (https://doi.org/10.1016/j.asoc.2024.111631)

If you find this code is useful and helpful to your work, please cite our paper in your research work. Thanks.

If there are any questions about source code, please do not hesitate to contact Ju Li([email protected]) and me ([email protected]). Special thanks to Yang Wei for helping to arrange and re-implement the code.

How to use the code

Running environment:

The proposed methods are implemented in Python 3.7 with PyTorch framework on a desktop computer equipped with an Intel i9 7920X CPU and an NVIDIA RTX 3090 GPU.

Dataset used in this paper:

1. CIFAR-10 & CIFAR-100
2. CALTECH 101
3. Magnetic tile
4. AMT

Experimental results:

1. Surface defect classification based on Resnet-18 and 50 with CIFAR-10 dataset.
2. Surface defect classification based on Resnet-18 and 50 with CIFAR-100 dataset.
3. Surface defect detection based on YOLOv4 with AMT dataset.

Source code:

1. Surface defect classification based on Resnet-18 and 50 with AMT dataset.
2. Surface defect detection based on YOLOv4 with AMT dataset.

How to reproduce the experimental results of surface defect classification.

1. find the cal_101_resnet.py file in classification folder:
2. configurable arguments: train.py [--link_place LINK_PLACE] [--re_co RE_CO] [--backbone BACKBONE]
3. train your backbone
4. evaluating results will automatically show after training.

Similarly, How to reproduce the experimental results of surface defect detection.

1. find the train.py file in detection folder:
2. configurable arguments: train.py [--weight_path WEIGHT_PATH] [--resume] [--gpu_id GPU_ID] [--log_path LOG_PATH] [--accumulate ACCUMULATE] [--fp_16 FP_16]
3. train your backbone
4. evaluating results will automatically show after training.

Tips: The experimental results included in the above "Experimental_Results" files are re-calculated when submitting code, which may be a slight deviation from the results reported in our paper due to the effects of randomness of choosing training samples.

Reference：The original Magnetic tile (MT) dataset is available at: https://github.com/abin24/Saliency-detection-toolbox

1. Y. Huang, C. Qiu, K. Yuan, Surface defect saliency of magnetic tile, The Visual Computer 36 (2020) 85–96.