Cách cô lập các đối tượng phân đoạn với Ultralytics YOLO#

Instance segmentation tạo ra một mặt nạ (mask) có độ chính xác đến từng pixel cho mọi đối tượng được phát hiện, có nghĩa là bạn có thể tách riêng từng đối tượng ra khỏi một bức ảnh. Hướng dẫn này chỉ cho bạn cách biến kết quả segmentation của Ultralytics YOLO thành các đối tượng được cô lập bằng cách sử dụng Predict Mode và OpenCV, với nền đen đặc hoặc nền trong suốt để lưu dưới dạng PNG.

Tại sao cần cô lập đối tượng phân đoạn?#

Việc tách từng đối tượng ra khỏi ảnh mở ra nhiều quy trình xử lý tiếp theo:

• Xóa nền cho ảnh sản phẩm, danh mục hoặc chỉnh sửa sáng tạo.
• Cắt theo đối tượng để xây dựng tập dữ liệu phân loại từ các phát hiện của bạn.
• Xử lý tập trung giúp các bước sau như OCR, phân tích màu sắc hoặc đo lường chỉ nhìn thấy đối tượng, thay vì khung cảnh xung quanh.
• Xuất file PNG trong suốt để ghép các đối tượng lên nền mới.

Quy trình này hoạt động với bất kỳ mô hình segmentation nào của Ultralytics YOLO và tuân theo bốn giai đoạn: chạy inference → trích xuất từng đường viền → cô lập đối tượng → lưu kết quả.

Chạy Segmentation Inference#

Cài đặt các thư viện cần thiết, sau đó tải mô hình segmentation (hậu tố -seg, bắt buộc để tạo mặt nạ) và chạy dự đoán trên ảnh nguồn của bạn:

from ultralytics import YOLO

# Load a segmentation model
model = YOLO("yolo26n-seg.pt")

# Run inference on a source
results = model.predict(source="path/to/image.jpg")

Trích xuất đường bao đối tượng#

Mỗi mục trong results tương ứng với một hình ảnh và việc lặp qua kết quả sẽ trả về từng phát hiện một. Đối với mỗi phát hiện, hãy sao chép hình ảnh gốc, đọc nhãn lớp và vẽ đường viền mặt nạ của đối tượng lên một mặt nạ nhị phân trống. Vùng màu trắng của mặt nạ này đánh dấu chính xác những pixel nào thuộc về đối tượng.

Các đoạn mã trong phần này và phần tiếp theo chạy bên trong vòng lặp phát hiện dưới đây; tập lệnh hoàn chỉnh, có thể sao chép nằm ở phần Full Example.

{ width="240", align="right" }

from pathlib import Path

import cv2
import numpy as np

for r in results:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem  # source image base-name

    # Iterate each detected object in the image
    for ci, c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]  # class name

        # Build a binary mask and draw the object contour onto it
        b_mask = np.zeros(img.shape[:2], np.uint8)
        contour = c.masks.xy[0].astype(np.int32).reshape(-1, 1, 2)
        cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

Cô lập đối tượng#

Với mask nhị phân đã sẵn sàng, hãy kết hợp nó với ảnh gốc. Có hai kiểu phổ biến, tùy thuộc vào nền bạn muốn:

# Isolate object with a black background
mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
isolated = cv2.bitwise_and(mask3ch, img)

# Bounding box coordinates
x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
# Crop the isolated image to the object region
iso_crop = isolated[y1:y2, x1:x2]

Lưu kết quả (tùy chọn)#

Việc bạn làm gì với từng đối tượng đã cô lập là tùy thuộc vào bạn. Một bước tiếp theo phổ biến là ghi nó vào ổ đĩa để sử dụng sau:

# Save the isolated object to file
cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated)

Ở đây img_name là phần gốc tên tệp ảnh nguồn, label là tên lớp và ci là chỉ mục phát hiện, do đó nhiều thể hiện của cùng một lớp sẽ có tên tệp độc nhất. Hãy thay thế isolated bằng iso_crop nếu bạn đã áp dụng tùy chọn cắt ảnh ở trên.

Ví dụ đầy đủ#

Tập lệnh dưới đây kết hợp mọi bước thành một khối duy nhất có thể chạy được. Nó sử dụng nền đen theo mặc định; hãy chuyển dòng được đánh dấu đơn thành np.dstack([img, b_mask]) để có tệp PNG trong suốt thay thế:

from pathlib import Path

import cv2
import numpy as np

from ultralytics import YOLO

model = YOLO("yolo26n-seg.pt")
results = model.predict(source="path/to/image.jpg")

for r in results:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem

    for ci, c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]

        # Build a binary mask from the object contour
        b_mask = np.zeros(img.shape[:2], np.uint8)
        contour = c.masks.xy[0].astype(np.int32).reshape(-1, 1, 2)
        cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

        # Isolate the object (black background)
        mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
        isolated = cv2.bitwise_and(mask3ch, img)  # transparent PNG: isolated = np.dstack([img, b_mask])

        # Save or add your custom post-processing here
        cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated)

        # Optional: crop to the bounding box before saving
        # x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
        # cv2.imwrite(f"{img_name}_{label}-{ci}.png", isolated[y1:y2, x1:x2])

Để sử dụng nhiều lần, hãy gói phần thân vòng lặp trong một hàm để bạn có thể gọi nó trên nhiều ảnh.

Kết luận#

Bây giờ bạn đã có một công thức hoàn chỉnh để cô lập các đối tượng được phân đoạn bằng Ultralytics YOLO: chạy inference, xây dựng mặt nạ nhị phân từ mỗi đường viền, sau đó trích xuất đối tượng trên nền đen hoặc nền trong suốt và tùy chọn cắt nó theo khung bao quanh. Khám phá tài liệu đầy đủ về Segment Task và Predict Mode để tùy chỉnh quy trình làm việc cho các lớp của riêng bạn.

Câu hỏi thường gặp#