VietOCR - Nhận Dạng Tiếng Việt Sử Dụng Mô Hình Transformer và AttentionOCR

AttentionOCR

AttentionOCR là sự kết hợp giữa mô hình CNN và mô hình Attention Seq2Seq. Cách hoạt động của mô hình này tương tự như kiến trúc của mô hình seq2seq trong bài toán dịch máy. Với bài toán dịch máy từ tiếng việt sang anh, chúng ta cần encode một chuỗi tiếng việt thành một vector đặc trưng, còn trong mô hình AttentionOCR, thì dữ liệu đầu vào này là một ảnh.

Một ảnh qua mô hình CNN, sẽ cho một feature maps có kích thước channelxheightxwidth, feature maps này sẽ trở thành đầu vào cho mô hình LSTM, tuy nhiên, mô hình LSTM chỉ nhận chỉ nhận đầu vào có kích thước là hiddenxtime_step. Một cách đơn giản và hợp lý là 2 chiều cuối cùng heightxwidth của feature maps sẽ được duổi thẳng. Feature maps lúc này sẽ có kích thước phù hợp với yêu cầu của mô hình LSTM.

Feature maps của mô hình CNN sau khi được flatten thì được truyền vào làm input của mô hình LSTM, tại mỗi thời điểm, mô hình LSTM cần dự đoán từ tiếp theo trong ảnh là gì.

TransformerOCR

Các bạn có thể tận dụng kiến trúc transformer thay cho mô hình LSTM để dự đoán từ tiếp theo trong ảnh. Chi tiết kiến trúc và cách hoạt động của mô hình transformer mình đã giải thích rất chi tiết tại đây.

Huấn Luyện Mô Hình

Huấn luyện mô hình AttenionOCR hay TransformerOCR hoàn toàn giống với luyện mô hình seq2seq, chúng đều sử dụng cross-entropy loss để tối ưu thay vì sử dụng CTCLoss như mô hình CRNN, tức là tại mỗi thời điểm mô hình dự đoán một từ sau đó so sánh với nhãn để tính loss và cập nhật lại trọng số của mô hình.

Hạn Chế Của Mô Hình Sử dụng CTCLoss

Đối với mô hình CRNN sử dụng CTCloss để làm hàm mục tiêu, số lượng kí tự đối đa có thể dự đoán bằng với widthxheight của feature maps. Do đó, các bạn cần phải cẩn thận điều chỉnh kiến trúc mô hình để có thể dự đoán được số kí tự phù hợp với từng bộ dataset. Đối với mô hình AttentionOCR hoặc TransformerOCR, các bạn không gặp vấn đề này, làm cho các bạn có thể dễ dàng sử dụng lại pretrained model cho các loại dữ liệu khác nhau.
Ngoài ra, AttentionOCR hoặc TransformerOCR đều có kiến trúc của mô hình dịch này seq2seq, do đó các thủ thuật của mô hình này đều có thể ứng dụng cho mô hình của chúng ta.

Thư Viện VietOCR

Thư viện VietOCR được mình xây dựng với mục đích hỗ trợ các bạn có thể sử dụng để giải quyết các bài toán liên quan đến OCR trong công nghiệp. Thư viện cung cấp cả 2 kiến trúc AtentionOCR và TransformerOCR. Tuy kiến trúc TransformerOCR hoạt động khá tốt trong NLP, nhưng theo mình nhận xét thì độ chính không có sự cải thiện đáng kể so với AttentionOCR mà thời gian dự đoán lại chậm hơn khá nhiều.
Mình có cung cấp pretrained model được mình huấn luyện trên tập dữ liệu 10m ảnh để các bạn có thể sử dụng nhanh chóng trong các bài toán mới. Tuy nhiên, mình khuyến khích các bạn huấn luyện mô hình trên tập dữ liệu mới của bản thân nếu muốn sử dụng trong công nghiệp.
Để thử nghiệm nhanh chóng mô hình các bạn có thể tham khảo notebook tại đây
Ở phần tiếp theo, mình sẽ hướng dẫn các bạn cách tạo bộ dataset phù hợp với thư viện, sử dụng thư viện để huấn luyện trên dataset mới, thay đổi cách augmentate, dự đoán cho ảnh mới và một số lưu ý khi sử dụng thư viện.