AnyLabeling - gán nhãn ảnh thông minh với Segment Anything và YOLO

Gần đây tôi ra mắt AnyLabeling, một công cụ gán nhãn ảnh có tích hợp Segment Anything và các mô hình YOLO. AnyLabeling được phát triển từ Labelme, với mục tiêu giúp quá trình gán nhãn nhanh và chính xác hơn. Bài viết này trình bày cách tôi xây dựng công cụ đó.

1. Vì sao lại làm AnyLabeling?

Khi Meta công bố Segment Anything Model (SAM), mô hình AI có thể tách một vật thể bất kỳ trong ảnh chỉ bằng một cú nhấp chuột, cộng đồng AI rất hào hứng. Sau khi xem bản demo, tôi nhận ra SAM hoàn toàn có thể trở thành nền tảng cho một công cụ gán nhãn thông minh. Đó là điểm khởi đầu của AnyLabeling.

Kết quả từ Segment Anything Model - Nguồn: Meta.

2. Lấy Labelme làm nền

Labelme là công cụ gán nhãn phổ biến với hơn 10.000 sao trên GitHub. Dự án được viết bằng Python và Qt, hỗ trợ nhiều định dạng ảnh, đồng thời có kiến trúc tương đối dễ mở rộng. Vì vậy, tôi chọn Labelme làm nền cho AnyLabeling và thay đổi một số phần so với dự án gốc:

• Thiết kế lại UI cho hiện đại và dễ dùng hơn.
• Bổ sung kiến trúc phục vụ suy luận mô hình; hiện đã tích hợp Segment Anything và các mô hình YOLO.
• Bổ sung tính năng gán nhãn OCR, cho phép nhập nội dung văn bản trong ảnh và gom các vùng văn bản thành nhóm.
• Chuẩn hoá mã Python theo PEP 8. Đây chỉ là lựa chọn cá nhân, không ảnh hưởng đến chức năng, nhưng cũng tiêu tốn khá nhiều thời gian vì quy ước mã nguồn ban đầu của Labelme khác PEP 8.

UI và bộ icon đã được thiết kế lại trong AnyLabeling cho hiện đại hơn.

Gán nhãn dữ liệu OCR trong AnyLabeling.

3. Suy luận mô hình

Khả năng chạy suy luận ngay trong công cụ là phần cốt lõi để tự động hoá quy trình gán nhãn. Phiên bản đầu tiên của AnyLabeling hỗ trợ Segment Anything và YOLO. Tôi tích hợp các mô hình Segment Anything cùng một mô hình YOLOv5; những mô hình YOLOv5 và YOLOv8 khác do Henry đóng góp. Kiến trúc suy luận được minh hoạ dưới đây:

Trong kiến trúc của AnyLabeling, LabelingWidget là widget chính, còn lớp Canvas phụ trách vùng vẽ. Tôi bổ sung AutoLabelingWidget làm giao diện chính cho tính năng gán nhãn tự động và ModelManager để quản lý, thực thi các mô hình AI.

a. Tích hợp Segment Anything

Segment Anything gồm hai phần: encoder tương đối nặng và decoder nhẹ hơn. Encoder trích xuất embedding từ ảnh đầu vào. Từ embedding này và prompt của người dùng (điểm, bounding box hoặc mask), decoder sinh ra một hay nhiều mask kết quả.

Segment Anything trong AnyLabeling

Trong bản demo web, Meta chạy encoder trên máy chủ, còn decoder chạy theo thời gian thực ngay trong trình duyệt: người dùng chọn điểm hoặc vẽ bounding box và nhận kết quả tức thì. AnyLabeling cũng chỉ chạy encoder một lần cho mỗi ảnh. Sau đó, mỗi khi prompt thay đổi, hệ thống chỉ cần chạy lại decoder để tạo mask mới. Một bước hậu xử lý tiếp tục tìm đường biên và chuyển mask thành các đối tượng gán nhãn như đa giác hoặc hình chữ nhật.

Để giảm số lượng dependency, tôi không dùng trực tiếp gói segment_anything của Meta mà viết lại phần suy luận chỉ với ONNX Runtime và NumPy. Mã thực thi mô hình ONNX nằm tại đây; các tệp mô hình được phát hành trong AnyLabeling Assets.

Segment Anything trong AnyLabeling

Người dùng tương tác với mô hình qua nhóm công cụ Auto segmentation marking tools. Bản demo và hướng dẫn chi tiết có trong tài liệu AnyLabeling.

Video dưới đây minh hoạ cách tính năng này hoạt động:

Watch

Tăng tốc:

Do encoder tốn khá nhiều thời gian, tôi lưu kết quả vào bộ nhớ đệm và tính trước embedding cho các ảnh kế tiếp. Nhờ vậy, người dùng không phải chờ encoder sau mỗi lần chuyển ảnh.

• Bộ nhớ đệm: một LRU cache lưu kết quả của encoder, với đường dẫn tệp nhãn làm khoá. Nếu embedding đã có trong cache, hệ thống không chạy lại encoder. Mặc định cache chứa kết quả của 10 ảnh.
• Tính toán trước: một worker thread chạy encoder cho các ảnh kế tiếp. Khi người dùng mở ảnh mới, ảnh đó cùng những ảnh phía sau được đưa vào hàng đợi xử lý. Embedding sau khi tính xong được lưu vào LRU cache; worker bỏ qua những ảnh đã có kết quả.

b. Tích hợp YOLO

Tích hợp YOLOv5 và YOLOv8 đơn giản hơn Segment Anything. Các mô hình YOLO phát hiện đối tượng trong ảnh và trả về danh sách bounding box hoặc mask phân đoạn, kèm nhãn lớp và độ tin cậy. Những kết quả này có thể chuyển thẳng thành dữ liệu gán nhãn.

Các mô hình YOLO trong AnyLabeling

Tôi đã tích hợp một mô hình YOLOv5 trước khi bắt tay vào Segment Anything. Sau đó, Henry bổ sung các biến thể YOLOv5 và YOLOv8, bao gồm cả mô hình phân đoạn.

Để áp dụng cho các bài toán cụ thể, chúng tôi sẽ bổ sung hướng dẫn xây dựng và nạp mô hình tuỳ chỉnh với cấu hình, tập nhãn riêng.

4. Gán nhãn dữ liệu OCR

Gán nhãn dữ liệu OCR là một tính năng mới của AnyLabeling. Đây là nhu cầu phổ biến trong nhiều dự án, nhưng cả Labelme lẫn LabelImg đều chưa hỗ trợ tốt, vì vậy tôi quyết định bổ sung vào AnyLabeling.

Gán nhãn dữ liệu OCR trong AnyLabeling.

Bản đầu tiên hỗ trợ các kiểu gán nhãn sau:

• Gán nhãn văn bản cho toàn ảnh: chuyển sang chế độ Edit để nhập tên hoặc mô tả cho ảnh.
• Gán nhãn phát hiện văn bản với mọi loại hình học (hình chữ nhật, đa giác...): sau khi tạo một đối tượng, chuyển sang chế độ Edit để nhập nội dung văn bản tương ứng.
• Gom nhóm văn bản: với bài toán KIE (Key Information Extraction), người dùng thường cần liên kết các vùng văn bản, chẳng hạn một tên trường với giá trị của nó. Chọn các đối tượng rồi nhấn G để gom nhóm; các đối tượng cùng nhóm được hiển thị bằng một màu. Nhấn U để tách nhóm.

Nội dung văn bản và thông tin nhóm được lưu trong cùng tệp JSON với các annotation khác. Nội dung nằm ở trường text, còn mã nhóm nằm ở trường group_id của đối tượng.

5. Đôi lời cuối

Ngay trong tuần đầu ra mắt, AnyLabeling đã nhận được nhiều phản hồi từ cộng đồng. Tôi rất vui khi thấy công cụ có ích cho nhiều người; hơn 200 sao GitHub trong tuần đầu tiên là một cột mốc đáng nhớ.

Trong thời gian tới, tôi sẽ tiếp tục cải thiện AnyLabeling. Một số hạng mục đang được cân nhắc gồm:

• Tài liệu: bổ sung nội dung và hướng dẫn thực hành trên anylabeling.nrl.ai để người dùng dễ bắt đầu hơn.
• Gán nhãn tự động: thêm các công cụ cho OCR, phát hiện và phân loại văn bản.
• Giao diện: cho phép tuỳ chỉnh màu đối tượng, kích thước điểm và nhiều thuộc tính hiển thị khác.
• Hiệu năng: tối ưu cả quá trình render Canvas lẫn tốc độ của các công cụ gán nhãn tự động.

Nếu có ý tưởng, góp ý hoặc muốn đóng góp cho AnyLabeling, bạn có thể liên hệ với tôi qua vietanhdev [atttt] gmail dot com. Tiến độ của dự án được cập nhật trên GitHub repository và GitHub Project.