Huấn luyện (fine-tune) một mô hình nhận diện vật thể tùy biến trên Google Colab có thể thực hiện nhanh chóng thông qua họ mô hình YOLO (cụ thể là YOLOv8 của Ultralytics). YOLOv8 hỗ trợ huấn luyện trên dữ liệu riêng (Custom Dataset) và tận dụng GPU miễn phí của Google Colab để hoàn tất quá trình train chỉ trong 5–10 phút.
Bài toán minh họa: Hệ thống đếm sản phẩm trái cây/sản phẩm công nghiệp trên băng chuyền.
[1. Chuẩn bị dữ liệu] ──> [2. Khai báo data.yaml] ──> [3. Train / Fine-tune] ──> [4. Suy luận & Đếm]
(Ảnh + Tệp nhãn txt) (Đường dẫn & Tên lớp) (YOLOv8 Nano + GPU) (cv2 + Bounding Box)
Trên giao diện Google Colab, bạn hãy bật GPU trước:
Vào menu Runtime $\rightarrow$ Change runtime type $\rightarrow$ Chọn T4 GPU $\rightarrow$ Nhấn Save.
Dưới đây là mã nguồn chia làm 4 ô (cells) chạy tuần tự.
(Để bạn không cần tốn thời gian chụp ảnh và gắn nhãn thủ công hàng giờ bằng các công cụ như Roboflow, đoạn mã sau sẽ dùng OpenCV và NumPy để tự tạo 40 ảnh train và 10 ảnh validation mô phỏng sản phẩm tròn trên băng chuyền kèm file nhãn tọa độ chuẩn định dạng YOLO).
# Cài đặt thư viện Ultralytics (YOLO)
!pip install -q ultralytics
import os
import cv2
import numpy as np
# Tạo cấu trúc thư mục chuẩn của YOLO
for split in ['train', 'val']:
os.makedirs(f"dataset/images/{split}", exist_ok=True)
os.makedirs(f"dataset/labels/{split}", exist_ok=True)
def create_synthetic_sample(img_path, label_path):
# Khung hình kích thước 640x640 mô phỏng băng chuyền
img = np.full((640, 640, 3), 40, dtype=np.uint8)
# Vẽ các đường rãnh băng chuyền mờ
for y in range(0, 640, 80):
cv2.line(img, (0, y), (640, y), (55, 55, 55), 2)
labels = []
# Sinh ngẫu nhiên từ 1 đến 3 sản phẩm trên khung hình
num_objects = np.random.randint(1, 4)
for _ in range(num_objects):
r = np.random.randint(35, 60)
cx = np.random.randint(r + 20, 640 - r - 20)
cy = np.random.randint(r + 20, 640 - r - 20)
# Vẽ sản phẩm hình tròn màu cam/xanh
color = (0, 165, 255) # BGR
cv2.circle(img, (cx, cy), r, color, -1)
cv2.circle(img, (cx, cy), r, (0, 200, 255), 3)
# Chuẩn hóa tọa độ theo định dạng chuẩn YOLO: <class_id> <x_center> <y_center> <width> <height> (từ 0.0 đến 1.0)
norm_cx = cx / 640.0
norm_cy = cy / 640.0
norm_w = (r * 2) / 640.0
norm_h = (r * 2) / 640.0
labels.append(f"0 {norm_cx:.6f} {norm_cy:.6f} {norm_w:.6f} {norm_h:.6f}")
# Ghi ảnh và file nhãn txt
cv2.imwrite(img_path, img)
with open(label_path, "w") as f:
f.write("\n".join(labels))
print("⚙️ Đang sinh tập dữ liệu mẫu (40 ảnh train, 10 ảnh validation)...")
for i in range(40):
create_synthetic_sample(f"dataset/images/train/img_{i}.jpg", f"dataset/labels/train/img_{i}.txt")
for i in range(10):
create_synthetic_sample(f"dataset/images/val/img_{i}.jpg", f"dataset/labels/val/img_{i}.txt")
print("✅ Đã tạo xong tập dữ liệu cấu trúc YOLO!")data.yamlYOLO cần một file cấu hình mô tả đường dẫn thư mục và danh sách tên các lớp cần nhận diện.
yaml_content = """
path: /content/dataset # Đường dẫn gốc tới dataset
train: images/train # Thư mục ảnh huấn luyện
val: images/val # Thư mục ảnh đánh giá
# Danh sách tên các nhãn
names:
0: product
"""
with open("data.yaml", "w") as f:
f.write(yaml_content.strip())
print("✅ Đã tạo file data.yaml!")Chúng ta lấy trọng số khởi điểm từ mô hình siêu nhẹ yolov8n.pt (đã được huấn luyện sẵn trên 80 lớp của tập COCO), sau đó fine-tune lại lớp phân loại cuối cùng theo đặc trưng dữ liệu mới trong 15 epochs.
from ultralytics import YOLO
# 1. Nạp mô hình pre-trained nền tảng (Backbone)
model = YOLO("yolov8n.pt")
# 2. Bắt đầu quá trình huấn luyện chuyển giao (Fine-tuning)
# epochs=15: Chạy 15 lượt (mất khoảng 1-2 phút trên GPU T4)
# imgsz=640: Kích thước ảnh đầu vào chuẩn
results = model.train(
data="data.yaml",
epochs=15,
imgsz=640,
batch=8,
device=0 # Sử dụng GPU
)
print("🎉 Huấn luyện thành công!")
print("Trọng số tối ưu đã được lưu tại: runs/detect/train/weights/best.pt")cv2 để đếm sản phẩmSau khi train xong, file trọng số best.pt chứa toàn bộ tri thức nhận diện sản phẩm của bạn. Giờ ta nạp file này vào để thực hiện bài toán đếm:
import cv2
from google.colab.patches import cv2_imshow
from ultralytics import YOLO
# 1. Nạp mô hình tùy biến vừa huấn luyện
custom_model = YOLO("runs/detect/train/weights/best.pt")
# 2. Tạo một ảnh thực tế mới để kiểm thử độ chính xác
test_img = np.full((640, 640, 3), 40, dtype=np.uint8)
cv2.circle(test_img, (150, 300), 45, (0, 165, 255), -1)
cv2.circle(test_img, (400, 200), 55, (0, 165, 255), -1)
cv2.circle(test_img, (500, 480), 40, (0, 165, 255), -1)
cv2.imwrite("test_sample.jpg", test_img)
# 3. Chạy suy luận (Inference)
results = custom_model.predict(source="test_sample.jpg", conf=0.5)
# 4. Trích xuất kết quả nhận diện và đếm số lượng
result = results[0]
boxes = result.boxes
product_count = len(boxes)
# Vẽ bounding box bằng OpenCV lên ảnh kết quả
annotated_frame = test_img.copy()
for box in boxes:
# Lấy tọa độ x1, y1, x2, y2
x1, y1, x2, y2 = map(int, box.xyxy[0])
conf = float(box.conf[0])
cv2.rectangle(annotated_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(annotated_frame, f"Product {conf:.2f}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
# Gắn bảng điều khiển đếm tổng số lượng lên góc ảnh
cv2.rectangle(annotated_frame, (20, 20), (320, 80), (0, 0, 0), -1)
cv2.putText(annotated_frame, f"TOTAL COUNT: {product_count}", (35, 62),
cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2)
print(f"\n📊 KẾT QUẢ: Hệ thống phát hiện và đếm được {product_count} sản phẩm!")
cv2_imshow(annotated_frame)Chuẩn bị dữ liệu thật (Annotation):
Trong thực tế, bạn sẽ quay video hoặc chụp khoảng 100–300 bức ảnh sản phẩm thật ở các góc độ và điều kiện ánh sáng khác nhau.
Dùng công cụ miễn phí như Roboflow hoặc LabelImg để khoanh hộp (vẽ Bounding Box) và xuất định dạng nhãn theo chuẩn YOLOv8 PyTorch (gồm thư mục images, labels và file data.yaml).
Khái niệm Transfer Learning:
Thay vì huấn luyện từ đầu (mất hàng tuần trên hàng triệu ảnh), việc fine-tune từ yolov8n.pt giúp mô hình kế thừa sẵn khả năng nhận biết hình khối cơ bản của CNN, từ đó chỉ cần vài chục đến vài trăm ảnh mẫu để học nhận diện sản phẩm mới.