Trong lĩnh vực thị giác máy tính, việc tự động hóa quá trình kiểm tra chất lượng bề mặt vật liệu là một bài toán quan trọng. Bài viết này trình bày chi tiết quy trình xây dựng một hệ thống phát hiện khuyết tật trên bề mặt nhôm, bao gồm các bước xử lý tập dữ liệu, huấn luyện mô hình YOLOv8 và thiết kế giao diện người dùng đồ họa (GUI) bằng PyQt5.
Chuẩn bị và Phân chia Tập dữ liệu
Tập dữ liệu được sử dụng bao gồm 1400 hình ảnh bề mặt nhôm, chứa bốn loại khuyết tật chính: zhen_kong (lỗ rỗng), ca_shang (trầy xước), zang_wu (vết bẩn) và zhe_zhou (nếp gấp). Dữ liệu đã được gán nhãn sẵn dưới định dạng XML và TXT. Đối với YOLOv8, định dạng TXT (YOLO format) là lựa chọn tối ưu, giúp bỏ qua bước chuyển đổi phức tạp.
Để đánh giá khách quan, tập dữ liệu cần được phân chia thành tập huấn luyện (training set) và tập kiểm định (validation set) theo tỷ lệ 8:2. Dưới đây là đoạn mã Python sử dụng thư viện pathlib và shutil để thực hiện việc phân chia này một cách ngẫu nhiên và cấu trúc lại thư mục theo chuẩn của YOLO:
import shutil
import random
from pathlib import Path
def split_dataset(source_dir: str, dest_dir: str, val_ratio: float = 0.2):
src_path = Path(source_dir)
dest_path = Path(dest_dir)
# Tạo các thư mục đích theo cấu trúc chuẩn
train_img_dir = dest_path / 'train' / 'images'
train_lbl_dir = dest_path / 'train' / 'labels'
val_img_dir = dest_path / 'val' / 'images'
val_lbl_dir = dest_path / 'val' / 'labels'
for d in [train_img_dir, train_lbl_dir, val_img_dir, val_lbl_dir]:
d.mkdir(parents=True, exist_ok=True)
# Lấy danh sách tất cả file ảnh
img_extensions = {'.jpg', '.jpeg', '.png', '.bmp'}
image_files = [f for f in src_path.iterdir() if f.suffix.lower() in img_extensions]
# Xáo trộn ngẫu nhiên với seed cố định để tái lập kết quả
random.seed(42)
random.shuffle(image_files)
# Tính toán điểm cắt
split_idx = int(len(image_files) * (1 - val_ratio))
train_files = image_files[:split_idx]
val_files = image_files[split_idx:]
def copy_files(file_list, img_dest, lbl_dest):
for img_file in file_list:
shutil.copy(img_file, img_dest / img_file.name)
lbl_file = img_file.with_suffix('.txt')
if lbl_file.exists():
shutil.copy(lbl_file, lbl_dest / lbl_file.name)
copy_files(train_files, train_img_dir, train_lbl_dir)
copy_files(val_files, val_img_dir, val_lbl_dir)
print(f"Hoàn tất: {len(train_files)} ảnh huấn luyện, {len(val_files)} ảnh kiểm định.")
split_dataset('./raw_aluminum_data', './yolo_dataset')
Huấn luyện Mô hình YOLOv8
YOLOv8 mang lại sự cân bằng tuyệt vời giữa tốc độ suy luận và độ chính xác. Sau khi cấu hình file data.yaml trỏ đến các thư mục ảnh và nhãn đã tạo ở trên, chúng ta có thể khởi động quá trình huấn luyện. Đoạn mã dưới đây minh họa cách thiết lập các siêu tham số (hyperparameters) thông qua dictionary và bắt đầu training:
from ultralytics import YOLO
def start_training():
# Khởi tạo mô hình từ trọng số pre-trained
detector = YOLO('yolov8s.pt')
# Cấu hình các tham số huấn luyện
train_config = {
'data': './yolo_dataset/data.yaml',
'epochs': 150,
'imgsz': 640,
'batch': 16,
'name': 'aluminum_defect_run',
'patience': 20, # Dừng sớm nếu không cải thiện sau 20 epochs
'optimizer': 'AdamW'
}
# Tiến hành huấn luyện
training_results = detector.train(**train_config)
return training_results
if __name__ == "__main__":
start_training()
Xây dựng Giao diện Người dùng với PyQt5
Để hệ thống thân thiện hơn với người dùng cuối, một giao diện đồ họa được phát triển bằng PyQt5. Giao diện này cho phép người dùng tải ảnh từ máy tính thông qua hộp thoại, chạy mô hình dự đoán và hiển thị trực quan các bounding box của khuyết tật.
import sys
import cv2
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel,
QPushButton, QVBoxLayout, QWidget, QFileDialog)
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt
from ultralytics import YOLO
class InspectionUI(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Hệ Thống Kiểm Tra Khuyết Tật Nhôm")
self.resize(900, 700)
# Tải mô hình đã huấn luyện
self.model = YOLO('runs/detect/aluminum_defect_run/weights/best.pt')
self.setup_ui()
def setup_ui(self):
central_widget = QWidget()
self.setCentralWidget(central_widget)
self.display_label = QLabel("Chưa có ảnh nào được tải", self)
self.display_label.setAlignment(Qt.AlignCenter)
self.display_label.setStyleSheet("background-color: #2b2b2b; color: white; border: 2px solid #555;")
self.load_btn = QPushButton("Tải Ảnh và Phát Hiện", self)
self.load_btn.setStyleSheet("padding: 10px; font-size: 16px; background-color: #4CAF50; color: white;")
self.load_btn.clicked.connect(self.process_image)
layout = QVBoxLayout()
layout.addWidget(self.display_label, stretch=8)
layout.addWidget(self.load_btn, stretch=1)
central_widget.setLayout(layout)
def process_image(self):
file_path, _ = QFileDialog.getOpenFileName(self, "Chọn ảnh nhôm", "", "Images (*.png *.jpg *.jpeg *.bmp)")
if not file_path:
return
frame = cv2.imread(file_path)
if frame is None:
return
# Dự đoán với ngưỡng confidence 0.4
predictions = self.model(frame, conf=0.4)
annotated_frame = predictions[0].plot()
# Chuyển đổi từ BGR (OpenCV) sang RGB (Qt)
rgb_image = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
h, w, ch = rgb_image.shape
bytes_per_line = ch * w
qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
# Hiển thị lên QLabel với tỷ lệ co giãn
self.display_label.setPixmap(pixmap.scaled(self.display_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
if __name__ == '__main__':
app = QApplication(sys.argv)
window = InspectionUI()
window.show()
sys.exit(app.exec_())
Trong đoạn mã trên, QFileDialog được tích hợp để người dùng có thể linh hoạt chọn tệp ảnh đầu vào thay vì hardcode đường dẫn. Hàm cv2.cvtColor được sử dụng để chuyển đổi không gian màu từ BGR sang RGB, đảm bảo màu sắc hiển thị chính xác trên giao diện Qt. Cuối cùng, phương thức scaled kết hợp với Qt.KeepAspectRatio giúp hình ảnh tự động điều chỉnh kích thước phù hợp với khung hiển thị mà không làm méo tỷ lệ gốc của ảnh.