Automatic Document Classification by Type

Реализация автоматической классификации документов по типу

Классификация документов — первый этап в Document Processing Pipeline: прежде чем извлекать данные, система должна понять, что за документ перед ней. Счёт-фактура, накладная, паспорт, водительские права, акт выполненных работ — для каждого типа свой экстрактор. Точность классификации напрямую влияет на качество всего pipeline.

Мультимодальная классификация

Лучший подход — одновременно использовать визуальные признаки (как выглядит документ) и текстовые (что написано):

from transformers import LayoutLMv3ForSequenceClassification, LayoutLMv3Processor
import torch
import torch.nn as nn

class DocumentClassifier:
    def __init__(self, model_path: str, doc_types: list[str]):
        self.processor = LayoutLMv3Processor.from_pretrained(model_path)
        self.model = LayoutLMv3ForSequenceClassification.from_pretrained(
            model_path,
            num_labels=len(doc_types)
        )
        self.doc_types = doc_types
        self.model.eval()

    @torch.no_grad()
    def classify(self, image_path: str) -> dict:
        from PIL import Image
        image = Image.open(image_path).convert('RGB')

        encoding = self.processor(
            image, return_tensors='pt',
            truncation=True, max_length=512
        )

        outputs = self.model(**encoding)
        probs = torch.softmax(outputs.logits, dim=-1).squeeze()

        top_idx = probs.argmax().item()
        return {
            'document_type': self.doc_types[top_idx],
            'confidence': float(probs[top_idx]),
            'all_scores': {
                self.doc_types[i]: float(probs[i])
                for i in range(len(self.doc_types))
            }
        }

Training без LayoutLM: EfficientNet + BERT

Для быстрого прототипа без доступа к большим моделям:

import timm
from transformers import AutoTokenizer, AutoModel

class LightweightDocClassifier(nn.Module):
    def __init__(self, num_classes: int):
        super().__init__()
        # Visual encoder
        self.visual = timm.create_model('efficientnet_b2',
                                         pretrained=True, num_classes=0)
        # Text encoder
        self.text_encoder = AutoModel.from_pretrained('DeepPavlov/rubert-base-cased')
        self.tokenizer = AutoTokenizer.from_pretrained('DeepPavlov/rubert-base-cased')

        # Fusion
        vis_dim = self.visual.num_features  # 1408
        text_dim = 768
        self.fusion = nn.Sequential(
            nn.Linear(vis_dim + text_dim, 512),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, image_tensor, input_ids, attention_mask):
        vis_features = self.visual(image_tensor)
        text_out = self.text_encoder(input_ids, attention_mask)
        text_features = text_out.pooler_output  # [CLS] token

        combined = torch.cat([vis_features, text_features], dim=-1)
        return self.fusion(combined)

Типовые классы документов

Коллекция признаков для классификации

Для повышения точности — дополнительные признаки помимо визуальных и текстовых:

def extract_document_features(image_path: str, ocr_text: str) -> dict:
    return {
        # Структурные признаки
        'has_table': detect_tables(image_path),
        'has_signature': detect_signature_zone(image_path),
        'has_stamp': detect_stamp(image_path),
        'has_photo': detect_person_photo(image_path),

        # Текстовые паттерны (регулярные выражения)
        'has_inn': bool(re.search(r'\bИНН\b', ocr_text)),
        'has_kpp': bool(re.search(r'\bКПП\b', ocr_text)),
        'has_passport_series': bool(re.search(r'\d{4}\s\d{6}', ocr_text)),
        'has_invoice_number': bool(re.search(r'№\s*\d+', ocr_text)),

        # Метаданные
        'aspect_ratio': get_aspect_ratio(image_path),
        'orientation': detect_orientation(image_path),
    }

Metrics на российских документах

Типичная точность на корпусе российских документов (25 классов):

Сложные случаи: документы одного типа в разных форматах (разные банки, разные поставщики), плохое качество сканов, ламинированные документы.