In industrial manufacturing, the detection of casting defects in critical components such as railway wagon couplers is paramount for ensuring structural integrity, safety, and compliance with quality standards. Traditional inspection methods rely heavily on manual visual examination, which is labor-intensive, time-consuming, and prone to human error due to factors like fatigue and subjective judgment. With the increasing demand for production efficiency and automation, there is a pressing need for robust, real-time, and accurate defect detection systems. The advent of deep learning, particularly convolutional neural networks (CNNs), has revolutionized computer vision tasks, including object detection and classification. Among various algorithms, the YOLO (You Only Look Once) series has gained prominence for its balance between speed and accuracy, making it suitable for real-time applications. However, detecting casting defects in coupler images presents unique challenges: defects often exhibit subtle color variations, irregular shapes, and are influenced by environmental factors like varying illumination and surface textures. These challenges can lead to reduced detection precision and increased false positives or negatives in standard models.
This work addresses these issues by proposing an enhanced defect detection method based on the YOLOv8n architecture, specifically tailored for casting defects in railway wagon coupler images. The primary goal is to achieve high detection accuracy while maintaining real-time performance, which is crucial for integration into automated production lines. To this end, we introduce two key modifications: first, we replace the original backbone network of YOLOv8n with the lightweight MobileNetv3 to reduce computational complexity and accelerate inference speed; second, we develop a novel SA-BRA (Spatial Attention-Bi-level Routing Attention) module, incorporated after the backbone, to enhance feature extraction by dynamically focusing on spatially relevant regions associated with casting defects. These improvements aim to mitigate the limitations posed by indistinct defect characteristics and external disturbances. We validate our approach on a custom dataset of coupler casting images, containing various defect types, and demonstrate its superiority over existing methods through comprehensive experiments.
The detection of casting defects is a critical aspect of quality control in metalworking industries. Casting defects can arise from numerous factors, including improper mold design, material impurities, or cooling irregularities, leading to flaws such as porosity, cracks, inclusions, or surface irregularities. In railway wagon couplers, which are subjected to high mechanical stresses, even minor casting defects can compromise performance and safety. Automated visual inspection systems leveraging machine learning have shown promise, but they often struggle with the variability inherent in casting defects. For instance, defects may appear as low-contrast regions, blend with background textures, or vary in size and shape, making them difficult to distinguish from normal surface features. Moreover, industrial environments often have inconsistent lighting conditions, further complicating detection. Therefore, developing a specialized algorithm that can adapt to these challenges is essential for reliable defect identification.
In recent years, deep learning-based object detection algorithms have been widely adopted for industrial inspection tasks. Two-stage detectors like R-CNN and Faster R-CNN offer high accuracy but at the cost of slower inference speeds, making them less suitable for real-time applications. Single-stage detectors, such as the YOLO family and SSD, provide a better speed-accuracy trade-off. YOLOv8, the latest iteration in the series, introduces enhancements like a decoupled head and anchor-free design, improving both performance and flexibility. However, the standard YOLOv8n model, while efficient, may still lack the precision needed for subtle casting defects due to its fixed architecture. To address this, researchers have explored various modifications, including lightweight backbones for speed and attention mechanisms for accuracy. For example, MobileNet variants utilize depthwise separable convolutions to reduce parameters, while attention modules like SENet or CBAM enhance feature representation by emphasizing informative regions. Our work builds upon these ideas by integrating a lightweight backbone with a custom attention mechanism specifically designed for casting defects.
The core of our method lies in the integration of MobileNetv3 and the SA-BRA module into the YOLOv8n framework. MobileNetv3 is chosen for its efficiency, achieved through depthwise separable convolutions and optimized network design. The depthwise separable convolution decomposes a standard convolution into two steps: a depthwise convolution that applies a single filter per input channel, and a pointwise convolution that combines the outputs via 1×1 convolutions. This significantly reduces computational load. Mathematically, for an input feature map of size $P_F \times P_F$ with $Q_{in}$ channels, a convolutional kernel of size $P_K \times P_K$, and $Q_{out}$ output channels, the computational cost of a standard convolution is $P_K^2 \times Q_{in} \times Q_{out} \times P_F^2$, whereas for depthwise separable convolution, it is $P_K^2 \times Q_{in} \times P_F^2 + 1^2 \times Q_{in} \times Q_{out} \times P_F^2$. The ratio of costs is given by:
$$ \frac{P_K^2 \times Q_{in} \times P_F^2 + 1^2 \times Q_{in} \times Q_{out} \times P_F^2}{P_K^2 \times Q_{in} \times Q_{out} \times P_F^2} = \frac{1}{Q_{out}} + \frac{1}{P_K^2} $$
For a typical kernel size $P_K = 3$, this results in approximately 8 to 9 times reduction in computations, making the network lighter and faster. MobileNetv3 further incorporates squeeze-and-excitation (SE) attention and the h-swish activation function to improve accuracy without substantial overhead. By replacing the YOLOv8n backbone with MobileNetv3, we aim to boost inference speed, which is vital for real-time detection of casting defects in production lines.
However, lightweight networks often sacrifice some detection accuracy due to reduced model capacity. To compensate, we propose the SA-BRA module, which combines a bi-level routing attention (BRA) mechanism with spatial attention. The BRA mechanism dynamically selects relevant regions for attention computation, reducing redundancy and focusing on areas likely to contain casting defects. It operates by dividing the input feature map into $s \times s$ non-overlapping regions, generating query ($Q$), key ($K$), and value ($V$) tensors through linear projections:
$$ Q = x_r N_q, \quad K = x_r N_k, \quad V = x_r N_v $$
where $x_r$ represents a sub-region of the feature map, and $N_q$, $N_k$, $N_v$ are weight matrices. Then, region-to-region adjacency matrices $D_r$ are computed to measure inter-region relevance:
$$ D_r = Q_r (K_r)^T $$
Here, $Q_r$ and $K_r$ are average queries and keys within each region. Using a top-$k$ selection algorithm, the most relevant $k$ regions are retained, forming an index matrix $H_r$. The gathered key and value tensors, $K_g$ and $V_g$, are obtained via:
$$ K_g = \text{Gather}(K, H_r), \quad V_g = \text{Gather}(V, H_r) $$
The output feature map $C_{out}$ is then produced through attention computation:
$$ C_{out} = F(Q, K_g, V_g) $$
This bi-level routing allows the model to concentrate on sparse, defect-prone regions, enhancing feature discrimination. To further refine spatial focus, we augment BRA with a spatial attention module (SAM). SAM applies average and max pooling along the channel dimension, concatenates the results, and processes them through a convolutional layer and sigmoid activation to generate a spatial weight map. This map highlights important spatial locations, complementing the channel-wise emphasis of BRA. The combined SA-BRA module is inserted after the MobileNetv3 backbone, enabling the network to adaptively weigh both channel and spatial information for improved detection of casting defects.
The overall architecture of our improved YOLOv8n model is illustrated in the network diagram. The input images are resized to 640×640 pixels and fed into the MobileNetv3 backbone for initial feature extraction. The SA-BRA module then processes these features to emphasize regions containing casting defects. The enhanced features are passed through the neck network, which employs feature pyramid networks (FPN) and path aggregation networks (PAN) for multi-scale fusion, and finally to the decoupled head for classification and bounding box regression. This design ensures that the model remains lightweight yet powerful, capable of handling the variability in casting defects.
To evaluate our method, we constructed a custom dataset of railway wagon coupler casting images, as no public dataset exists for this specific application. The dataset comprises 3,276 images in JPG format, annotated with four types of casting defects: porosity, cracks, inclusions, and surface irregularities. The images were collected under varying lighting conditions to simulate real industrial environments. We split the dataset into training (70%), validation (20%), and test (10%) sets. Data augmentation techniques, including mosaic augmentation during training, were applied to increase diversity and prevent overfitting.
Experimental setup involved hardware with an Intel Core i5-7300HQ CPU and NVIDIA GeForce GTX 1050 GPU, and software with Python 3.8 and PyTorch framework. Training parameters included an input size of 640×640, SGD optimizer with a learning rate of 0.01, batch size of 2 (due to GPU limitations), and 300 epochs. We used standard evaluation metrics: precision (P), recall (R), mean average precision (mAP) at IoU thresholds of 0.5 (mAP@0.5) and 0.5:0.95 (mAP@0.5:0.95), number of parameters (Params), and frames per second (FPS) for speed. Precision and recall are defined as:
$$ P = \frac{TP}{TP + FP}, \quad R = \frac{TP}{TP + FN} $$
where $TP$ is true positive, $FP$ false positive, and $FN$ false negative. The average precision (AP) for each class is computed from the precision-recall curve, and mAP is the mean over all classes. FPS measures real-time capability, with higher values indicating faster inference.
We conducted ablation studies to assess the contribution of each component. The baseline YOLOv8n model was progressively modified by adding MobileNetv3, BRA, and SA-BRA modules. Results are summarized in the table below, showing metrics for each configuration.
| Model | MobileNetv3 | BRA | SA-BRA | mAP@0.5 (%) | mAP@0.5:0.95 (%) | P (%) | R (%) | Params (10^6) | FPS |
|---|---|---|---|---|---|---|---|---|---|
| YOLOv8n | No | No | No | 94.5 | 79.4 | 96.7 | 92.2 | 3.0 | 34.9 |
| Model 1 | Yes | No | No | 90.4 | 74.6 | 92.4 | 89.1 | 1.2 | 65.3 |
| Model 2 | No | Yes | No | 94.7 | 80.9 | 97.1 | 94.5 | 3.0 | 34.5 |
| Model 3 | No | No | Yes | 95.0 | 81.7 | 97.6 | 95.7 | 3.0 | 33.9 |
| Model 4 | Yes | Yes | No | 92.9 | 76.1 | 91.2 | 90.8 | 1.2 | 63.2 |
| Model 5 (Ours) | Yes | No | Yes | 93.7 | 77.3 | 94.3 | 91.5 | 1.2 | 61.8 |
The ablation results reveal several insights. Model 1, with MobileNetv3, shows a significant increase in FPS (65.3 vs. 34.9) due to parameter reduction (1.2M vs. 3.0M), but mAP@0.5 drops by 4.1%, indicating a trade-off between speed and accuracy. Model 2, with BRA, improves mAP@0.5 by 0.2% and recall by 2.3%, demonstrating the benefit of attention for casting defects detection, though with a slight FPS decrease. Model 3, with SA-BRA, achieves further gains in mAP@0.5 (95.0%) and recall (95.7%), highlighting the effectiveness of combined spatial and routing attention. Our final model (Model 5) integrates MobileNetv3 and SA-BRA, achieving a balance: mAP@0.5 is 93.7%, which is only 0.8% lower than the baseline, while FPS reaches 61.8, a 77% improvement in speed. This makes it suitable for real-time detection of casting defects, addressing the core objective.
We also compared our method with other mainstream object detection algorithms on the same dataset. The table below presents the results, including YOLOv5s, YOLOv7, Faster R-CNN, and baseline YOLOv8n.
| Model | mAP@0.5 (%) | mAP@0.5:0.95 (%) | P (%) | R (%) | Params (10^6) | FPS |
|---|---|---|---|---|---|---|
| YOLOv5s | 89.7 | 74.9 | 91.5 | 86.1 | 7.2 | 33.6 |
| YOLOv7 | 92.1 | 75.5 | 92.6 | 89.3 | 36.5 | 27.6 |
| Faster R-CNN | 93.2 | 79.6 | 93.8 | 91.5 | 100.8 | 18.4 |
| YOLOv8n | 94.5 | 79.4 | 96.7 | 92.2 | 3.0 | 34.9 |
| Improved YOLOv8n (Ours) | 93.7 | 77.3 | 94.3 | 91.5 | 1.2 | 61.8 |
Our model outperforms YOLOv5s and YOLOv7 in both accuracy and speed, with higher mAP@0.5 and FPS. Compared to Faster R-CNN, which has high accuracy but low FPS (18.4), our method offers a better speed-accuracy trade-off, crucial for industrial deployment. Although baseline YOLOv8n has slightly higher mAP@0.5 (94.5% vs. 93.7%), our improved version is significantly faster (61.8 FPS vs. 34.9 FPS), making it more practical for real-time casting defects detection. The reduction in parameters (1.2M vs. 3.0M) also means lower computational resource requirements, which is beneficial for edge devices in manufacturing settings.
To assess practical performance, we tested our model under different lighting conditions, simulating real-world scenarios. Under strong illumination, the model achieved high precision (e.g., 89% for porosity, 94% for cracks) with clear bounding boxes around casting defects. In low-light conditions, where defect features are less distinct, it maintained robust detection (e.g., 90% for inclusions, 84% for surface irregularities), demonstrating adaptability to environmental variations. Qualitative examples show that the model effectively localizes casting defects without significant false positives, even when defects blend with background textures. This robustness stems from the SA-BRA module’s ability to focus on relevant spatial regions and the lightweight backbone’s efficient feature extraction.
The success of our method can be attributed to several factors. First, MobileNetv3’s depthwise separable convolutions reduce computational load without severely compromising feature quality, enabling faster processing of casting defects images. Second, the SA-BRA module enhances feature representation by selectively attending to defect-prone areas, which is particularly useful for subtle casting defects that may be overlooked by standard convolutions. The bi-level routing mechanism efficiently prunes irrelevant regions, reducing noise and improving accuracy. Spatial attention further refines this by emphasizing key locations within those regions. Together, these components allow the model to learn discriminative features for casting defects, even under challenging conditions.
Nevertheless, there are limitations. In extremely low-light scenarios, detection performance may degrade, as the model relies on visual features that become less discernible. Future work could incorporate infrared or multi-spectral imaging to address this. Additionally, the dataset, while comprehensive, could be expanded to include more defect varieties and environmental conditions to improve generalization. Techniques like generative adversarial networks (GANs) could synthesize additional training data for rare casting defects. Another direction is to explore transformer-based architectures, which have shown promise in vision tasks, but may require optimization for real-time performance.
In conclusion, we have presented an enhanced defect detection method for casting defects in railway wagon coupler images, based on YOLOv8n. By integrating MobileNetv3 for lightweight efficiency and a novel SA-BRA module for improved feature attention, our model achieves a balance between accuracy and speed. Experimental results on a custom dataset show that it maintains high detection precision (93.7% mAP@0.5) while operating at 61.8 FPS, making it suitable for real-time industrial inspection. The method demonstrates robustness to varying lighting conditions and effectively identifies multiple types of casting defects. Compared to existing algorithms, it offers a superior trade-off, addressing the needs of automated production lines. This work contributes to the advancement of intelligent manufacturing systems, where reliable and fast detection of casting defects is essential for quality assurance and safety. Future efforts will focus on further optimizing the model for edge deployment and extending it to other casting components in the transportation industry.
The implications of this research extend beyond railway couplers to various metal casting applications, where similar challenges with defect detection persist. By leveraging deep learning innovations, we can move towards fully automated inspection pipelines, reducing human effort and increasing consistency. The proposed SA-BRA module, with its dynamic routing and spatial attention, could be adapted to other vision tasks involving complex textures or subtle anomalies. As industries continue to embrace Industry 4.0, such algorithms will play a pivotal role in enhancing productivity and ensuring product reliability. Our method represents a step forward in this direction, providing a practical solution for casting defects detection that balances computational efficiency with detection accuracy.
From a broader perspective, the integration of lightweight networks and attention mechanisms reflects a trend in computer vision towards more efficient and interpretable models. The ability to process images in real-time while maintaining high accuracy is critical for applications like autonomous systems, medical imaging, and security surveillance. In the context of casting defects, where defects can have severe consequences if undetected, the stakes are particularly high. Therefore, continuous improvement of detection algorithms, through techniques like those proposed here, is essential. We hope that this work inspires further research into specialized architectures for industrial vision tasks, ultimately contributing to safer and more efficient manufacturing processes worldwide.
To summarize, our improved YOLOv8n model effectively addresses the challenges of casting defects detection in railway wagon couplers by combining speed and precision. The use of MobileNetv3 and SA-BRA modules demonstrates how targeted modifications can enhance performance for specific applications. We believe that this approach can be generalized to other defect detection scenarios, offering a blueprint for developing custom deep learning solutions in industrial settings. As data availability and computational resources grow, we expect such methods to become increasingly prevalent, driving the next wave of automation in quality control.