Entrenamiento de un Modelo YOLO para Detectar Defectos en Soldaduras Industriales - Parte 1

Artificial intelligence-based weld flaw detection represents a significant breakthrough for the manufacturing industry. Using computer vision models such as YOLO, it is possible to automate quality control processes that traditionally required manual inspection, saving time and resources while improving accuracy. This customized approach demonstrates how AI can be tailored to specific industry needs.

How to create a customized weld flaw detection model with YOLO?

When we are faced with very specific problems such as weld flaw detection on mechanical parts, generic pre-trained models often fall short. In these cases, we need to create and train a customized model that exactly fits the customer's needs.

To begin this process, it is essential to have a suitable GPU, as computer vision model training is computationally expensive. In this case, we use a T4 GPU that allows us to perform the training efficiently.

What data do we need to train our model?

The starting point for any machine learning model is a quality data set. For our welding fault detector, we need:

• Dataset provided by the customer: Contains real business images.
• Dataset structure: Organized into test, training and validation folders.
• Labeled images: Each image has a corresponding label file indicating the position of the faults.
• Configuration file: A YAML file that specifies the folder paths and classes to detect.

In this specific case, our model must identify three different classes:

• Bad weld
• Good weld
• Defect

How to configure the model training?

To train our custom model with YOLOv11, we follow these steps:

1. Dependencies installation:

# We install Ultralytics to access YOLO!pip install ultralytics.

2. Load the pre-trained model:

# We load YOLOv11 as a base for our modelfrom ultralytics import YOLOmodel = YOLO('yolov11.pt').

3. Configuration file definition:

# We define the path to the configuration YAML fileyaml_path = 'data.yaml'

4. Training configuration:

# default configurationresults = model.train(  data=yaml_path, epochs=10, imgsz=640, augment=True)

In this first approach, we use a default configuration with 10 epochs (full iterations of the dataset) and an image size of 640 pixels. We also activated data augmentation to make the model more robust through image transformations.

How to evaluate and visualize the training results?

Once the training is completed, it is crucial to evaluate the performance of the model to determine its effectiveness. YOLO provides detailed metrics and visualization tools that facilitate this analysis.

Key metrics to evaluate the model

After training, we obtain several important metrics:

• Accuracy: 0.41 (indicates what proportion of positive detections were correct).
• Recall: 0.51 (indicates what proportion of the real objects were detected)
• F1-Score: Ideally, we look for values above 0.6-0.7.

Graphical display of results:

• F1-Score curves for each class.
• Confusion matrix showing hits and misses between classes.

In the confusion matrix, we were able to observe that for the "bad weld" class, the model got it right 40 times, but there was confusion with the "good weld" class in 53% of the cases. This information is valuable to identify areas of improvement in the model.

How to save and load the trained model?

Once training is complete, it is important to save the model for later use:

# We save the trained modelmodel.save('my_model.pt')
 # To load the model latermodel_loaded = YOLO('my_model.pt')

How to make predictions with the trained model?

The ultimate goal of this whole process is to use the model to detect faults in new images. YOLO offers several ways to perform predictions:

Basic visualization of results.

# We create a folder to save the results!mkdir testing
 # We perform the prediction and visualize the resultsresults = model_loaded('path_to_the_image.jpg')results[0].show()

Filtering results by confidence

In many cases, we want to filter detections according to their confidence level:

# We only show detections with confidence higher than 0.3results = loaded_model('path_to_the_image.jpg',  conf=0.3)results[0].show()

Filtering by specific class

If we are only interested in detecting a particular class:

# we only detect the "good weld" class (class 1)results = loaded_model('path_to_the_image.jpg', classes=1)results[0].show()

Implementation for production use

For production implementations, such as APIs or GUIs, we can create a function that processes the images and returns structured information:

def process_image(image_path, model): # We perform inference results = model(image_path) result = results[0]    
 # We extract relevant information boxes = result.boxes.xyxy.tolist() confidences = result.boxes.conf.tolist() class_names = [result.names.get(int(c), str(int(c))) for c in result.boxes.cls.tolist()]    
 # Get the annotated image annotated_image = result.plot()    
 return boxes, class_names, confidences, annotated_image

This function returns the coordinates of the bounding boxes, the names of the detected classes, the confidence levels and the image with the visual annotations, which facilitates its integration in more complex systems.

Artificial intelligence-based weld failure detection represents a significant breakthrough for the manufacturing industry. With the customized approach we have explored, it is possible to adapt computer vision models to specific needs, improving the quality and efficiency of control processes. Have you implemented similar solutions in your industry? Share your experience in the comments.