Детекторы

Модели детекции текста.

class manuscript.detectors.EAST(weights=None, device=None, *, target_size=1280, expand_ratio_w=1.4, expand_ratio_h=1.5, expand_power=0.6, score_thresh=0.6, iou_threshold=0.05, iou_threshold_standard=0.05, score_geo_scale=0.25, quantization=2, axis_aligned_output=True, remove_area_anomalies=False, anomaly_sigma_threshold=5.0, anomaly_min_box_count=30, use_tta=False, tta_iou_thresh=0.1, **kwargs)[исходный код]

Базовые классы: BaseModel

Инициализация текстового детектора EAST с использованием ONNX Runtime.

Параметры:

weights (str or Path, optional) – Путь или идентификатор для весов модели ONNX. Поддерживает: локальный путь, HTTP/HTTPS URL, GitHub release, Google Drive, имя пресета ("east_50_g1"), None — автозагрузка пресета по умолчанию (east_50_g1)
device (str, optional) – Устройство вычислений: "cuda", "coreml" или "cpu". Если None, автоматически выбирается CPU. Для ускорения GPU/CoreML: - CUDA (NVIDIA): pip install onnxruntime-gpu - CoreML (Apple Silicon M1/M2/M3): pip install onnxruntime-silicon По умолчанию None (CPU).
target_size (int, optional) – Размер входного изображения для инференса. Изображения приводятся к (target_size, target_size). Значение по умолчанию — 1280.
expand_ratio_w (float, optional) – Коэффициент горизонтального расширения, применяемый к боксам после NMS. По умолчанию 0.7.
expand_ratio_h (float, optional) – Коэффициент вертикального расширения, применяемый к боксам после NMS. По умолчанию 0.7.
expand_power (float, optional) – Степень нелинейного расширения боксов. Определяет, как масштабируется расширение в зависимости от размера бокса. - 1.0 = линейно (малые и большие боксы расширяются одинаково) - <1.0 = малые боксы расширяются сильнее (например, 0.5, рекомендуется для детекции символов) - >1.0 = большие боксы расширяются сильнее Значение по умолчанию — 0.5.
score_thresh (float, optional) – Порог уверенности для отбора кандидатов перед NMS. Значение по умолчанию — 0.7.
iou_threshold (float, optional) – Порог IoU для фазы объединения locality-aware NMS. По умолчанию 0.2.
iou_threshold_standard (float, optional) – Порог IoU для стандартного NMS после locality-aware объединения. Если None, используется значение iou_threshold. По умолчанию None.
score_geo_scale (float, optional) – Коэффициент масштабирования при декодировании карт геометрии/оценок. По умолчанию 0.25.
quantization (int, optional) – Разрешение квантования координат точек при декодировании. По умолчанию 2.
axis_aligned_output (bool, optional) – Если True, возвращаются осево-ориентированные прямоугольники вместо исходных четырёхугольников. По умолчанию True.
remove_area_anomalies (bool, optional) – Если True, удаляет четырёхугольники с аномально большой площадью относительно распределения. По умолчанию False.
anomaly_sigma_threshold (float, optional) – Порог сигмы для фильтрации аномальных площадей. По умолчанию 5.0.
anomaly_min_box_count (int, optional) – Минимальное количество боксов перед применением фильтрации аномалий. По умолчанию 30.
use_tta (bool, optional) – Enable Test-Time Augmentation (TTA). When enabled, inference is run on both the original and horizontally flipped image, and results are merged. This can improve detection of partially visible or edge text. Default is False.
tta_iou_thresh (float, optional) – IoU threshold for merging boxes from original and flipped images during TTA. Boxes with IoU > threshold are considered matches and merged. Default is 0.1.

Заметки

Класс предоставляет два основных публичных метода:

predict — выполняет инференс на одном изображении и возвращает детекции.
train — высокоуровневая точка входа для обучения модели EAST на пользовательских датасетах.

Детектор использует ONNX Runtime для быстрого инференса на CPU и GPU. Для ускорения на GPU установите: pip install onnxruntime-gpu

Методы

`__call__`(args, *kwargs)	Вызов объекта как функции.
`export`(weights_path, output_path[, ...])	Экспорт модели EAST из PyTorch в формат ONNX.
`predict`(img_or_path[, return_maps, ...])	Запуск инференса EAST и возврат результатов детекции.
`runtime_providers`()	Получение провайдеров выполнения ONNX Runtime в зависимости от устройства.
`train`(train_images, train_anns, val_images, ...)	Обучение модели EAST на пользовательских датасетах.

default_weights_name: Optional[str] = 'east_50_g1'

pretrained_registry: Dict[str, str] = {'east_50_g1': 'github://konstantinkozhin/manuscript-ocr/v0.1.0/east_50_g1.onnx'}

__init__(weights=None, device=None, *, target_size=1280, expand_ratio_w=1.4, expand_ratio_h=1.5, expand_power=0.6, score_thresh=0.6, iou_threshold=0.05, iou_threshold_standard=0.05, score_geo_scale=0.25, quantization=2, axis_aligned_output=True, remove_area_anomalies=False, anomaly_sigma_threshold=5.0, anomaly_min_box_count=30, use_tta=False, tta_iou_thresh=0.1, **kwargs)[исходный код]

Параметры:

weights (str | Path | None)
device (str | None)
target_size (int)
expand_ratio_w (float)
expand_ratio_h (float)
expand_power (float)
score_thresh (float)
iou_threshold (float)
iou_threshold_standard (float | None)
score_geo_scale (float)
quantization (int)
axis_aligned_output (bool)
remove_area_anomalies (bool)
anomaly_sigma_threshold (float)
anomaly_min_box_count (int)
use_tta (bool)
tta_iou_thresh (float)

predict(img_or_path, return_maps=False, sort_reading_order=True, split_into_columns=True, max_columns=10)[исходный код]

Запуск инференса EAST и возврат результатов детекции.

Параметры:

img_or_path (str or pathlib.Path or numpy.ndarray) – Путь к файлу изображения или RGB-изображение в виде массива NumPy формата uint8 с формой (H, W, 3).
return_maps (bool, optional) – Если True, возвращает необработанные карты оценок и геометрии под ключами "score_map" и "geo_map". По умолчанию False.
sort_reading_order (bool, optional) – Если True, сортирует обнаруженные слова в естественном порядке чтения (слева направо, сверху вниз) и группирует их в строки. По умолчанию True.
split_into_columns (bool, optional) – Если True и sort_reading_order=True, сегментирует страницу на колонки (отдельные Blocks). Если False, вся страница рассматривается как одна колонка. Используется только при sort_reading_order=True. По умолчанию True.
max_columns (int, optional) – Максимальное количество колонок для обнаружения при split_into_columns=True. Большие значения позволяют обнаруживать больше колонок. Используется только при sort_reading_order=True и split_into_columns=True. По умолчанию 10.

Результат:

Словарь со следующими ключами: - "page" : Page Разобранный результат детекции в виде объекта Page, содержащего Block(и) со Line(ами) объектов Word. Каждый Word содержит координаты полигона и оценки уверенности. Слова и строки имеют индексы порядка чтения. - "score_map" : numpy.ndarray или None Необработанная карта оценок, если return_maps=True. - "geo_map" : numpy.ndarray или None Необработанная карта геометрии, если return_maps=True.

Тип результата:

dict

Заметки

Метод выполняет: (1) загрузку изображения, (2) изменение размера и нормализацию, (3) инференс модели, (4) декодирование четырёхугольников, (5) NMS, (6) расширение боксов, (7) масштабирование координат обратно к исходному размеру, (8) опциональную сортировку по порядку чтения с группировкой в строки.

Test-Time Augmentation (TTA):

When use_tta=True is set during initialization, the method runs inference on both the original and horizontally flipped image, then merges results. Boxes from both views are matched by IoU and merged by taking the union of coordinates with averaged scores. This can improve detection of text near image edges or partially visible text.

Для визуализации используйте внешнюю утилиту visualize_page:

>>> from manuscript.utils import visualize_page
>>> result = model.predict(img_path)
>>> vis_img = visualize_page(img, result["page"])

Примеры

Выполнить инференс и получить структурированный результат:

>>> from manuscript.detectors import EAST
>>> model = EAST()
>>> img_path = r"example/ocr_example_image.jpg"
>>> result = model.predict(img_path)
>>> page = result["page"]
>>> # Access first line's first word
>>> first_word = page.blocks[0].lines[0].words[0]
>>> print(f"Confidence: {first_word.detection_confidence}")

Визуализировать результаты отдельно:

>>> from manuscript.utils import visualize_page, read_image
>>> result = model.predict(img_path)
>>> img = read_image(img_path)
>>> vis_img = visualize_page(img, result["page"])
>>> vis_img.show()

static train(train_images, train_anns, val_images, val_anns, *, experiment_root='./experiments', model_name='resnet_quad', backbone_name='resnet50', pretrained_backbone=True, freeze_first=True, target_size=1024, score_geo_scale=None, epochs=500, batch_size=3, accumulation_steps=1, lr=0.001, grad_clip=5.0, early_stop=100, use_sam=True, sam_type='asam', use_lookahead=True, use_ema=False, use_multiscale=True, use_ohem=True, ohem_ratio=0.5, use_focal_geo=True, focal_gamma=2.0, resume_from=None, val_interval=1, num_workers=0, device=None)[исходный код]

Обучение модели EAST на пользовательских датасетах.

Параметры:

train_images (str, Path or sequence of paths) – Путь(и) к каталогам с обучающими изображениями.
train_anns (str, Path or sequence of paths) – Путь(и) к JSON-файлам аннотаций в формате COCO, соответствующим train_images.
val_images (str, Path or sequence of paths) – Путь(и) к каталогам с валидационными изображениями.
val_anns (str, Path or sequence of paths) – Путь(и) к JSON-файлам аннотаций в формате COCO, соответствующим val_images.
experiment_root (str, optional) – Базовый каталог, в котором будут создаваться папки экспериментов. По умолчанию "./experiments".
model_name (str, optional) – Имя папки внутри experiment_root для логов и чекпоинтов. По умолчанию "resnet_quad".
backbone_name ({"resnet50", "resnet101"}, optional) – Используемая архитектура backbone. Варианты: - "resnet50" — ResNet-50 (быстрее, меньше параметров) - "resnet101" — ResNet-101 (медленнее, больше ёмкость) По умолчанию "resnet50".
pretrained_backbone (bool, optional) – Использовать веса backbone, предобученные на ImageNet. По умолчанию True.
freeze_first (bool, optional) – Заморозить нижние слои backbone. По умолчанию True.
target_size (int, optional) – Изменять размер короткой стороны изображений до этого значения. По умолчанию 1024.
score_geo_scale (float, optional) – Множитель для восстановления исходных координат из карт оценок/геометрии. Если None, автоматически берётся из модели. По умолчанию None.
epochs (int, optional) – Количество эпох обучения. По умолчанию 500.
batch_size (int, optional) – Размер батча на GPU. По умолчанию 3.
accumulation_steps (int, optional) – Количество шагов накопления градиента. Эффективный размер батча будет batch_size * accumulation_steps. Используйте это для обучения с большими эффективными батчами при ограниченной памяти GPU. Например: - batch_size=2, accumulation_steps=4 → эффективный размер батча = 8 - batch_size=1, accumulation_steps=8 → эффективный размер батча = 8 По умолчанию 1 (без накопления).
lr (float, optional) – Скорость обучения. По умолчанию 1e-3.
grad_clip (float, optional) – Значение отсечения градиента (L2-норма). По умолчанию 5.0.
early_stop (int, optional) – Терпение (число эпох без улучшений) для ранней остановки. По умолчанию 100.
use_sam (bool, optional) – Включить оптимизатор SAM. По умолчанию True.
sam_type ({"sam", "asam"}, optional) – Используемый вариант SAM. По умолчанию "asam".
use_lookahead (bool, optional) – Обернуть оптимизатор в Lookahead. По умолчанию True.
use_ema (bool, optional) – Поддерживать EMA-версию весов модели. По умолчанию False.
use_multiscale (bool, optional) – Случайное мульти-масштабное обучение. По умолчанию True.
use_ohem (bool, optional) – Онлайн-отбор сложных примеров (OHEM). По умолчанию True.
ohem_ratio (float, optional) – Доля сложных негативных примеров для OHEM. По умолчанию 0.5.
use_focal_geo (bool, optional) – Применять focal loss к геометрическим каналам. По умолчанию True.
focal_gamma (float, optional) – Параметр gamma для focal loss геометрии. По умолчанию 2.0.
resume_from (str or Path, optional) – Возобновить обучение из предыдущего эксперимента: a) каталог эксперимента, b) …/checkpoints/, c) прямой путь к last_state.pt. По умолчанию None.
val_interval (int, optional) – Запускать валидацию каждые N эпох. По умолчанию 1.
num_workers (int, optional) – Количество рабочих процессов для загрузки данных. Установите 0 для однопроцессной загрузки (безопаснее на Windows). По умолчанию 0.
device (torch.device, optional) – Устройство CUDA или CPU. Автоматически выбирается, если None.

Результат:

Лучшие веса модели (EMA, если включено, иначе базовая модель).

Тип результата:

torch.nn.Module

Примеры

Обучение на двух датасетах с валидацией:

>>> from manuscript.detectors import EAST
>>>
>>> train_images = [
...     "/data/archive/train_images",
...     "/data/ddi/train_images"
... ]
>>> train_anns = [
...     "/data/archive/train.json",
...     "/data/ddi/train.json"
... ]
>>> val_images = [
...     "/data/archive/test_images",
...     "/data/ddi/test_images"
... ]
>>> val_anns = [
...     "/data/archive/test.json",
...     "/data/ddi/test.json"
... ]
>>>
>>> best_model = EAST.train(
...     train_images=train_images,
...     train_anns=train_anns,
...     val_images=val_images,
...     val_anns=val_anns,
...     backbone_name="resnet50",
...     target_size=256,
...     epochs=20,
...     batch_size=4,
...     use_sam=False,
...     freeze_first=False,
...     val_interval=3,
... )
>>> print("Best checkpoint loaded:", best_model)

static export(weights_path, output_path, backbone_name=None, input_size=1280, opset_version=14, simplify=True)[исходный код]

Экспорт модели EAST из PyTorch в формат ONNX.

Этот метод конвертирует обученную модель EAST из PyTorch в формат ONNX, который может использоваться для более быстрого инференса с ONNX Runtime. Экспортированную модель можно загрузить с помощью EAST(weights_path="model.onnx", use_onnx=True).

Параметры:

weights_path (str or Path) – Путь к файлу весов модели PyTorch (.pth).
output_path (str or Path) – Путь, по которому будет сохранена модель ONNX (.onnx).
backbone_name ({"resnet50", "resnet101"}, optional) – Архитектура backbone модели. Если None, будет автоматически определена из чекпоинта. Должна совпадать с архитектурой, использованной при обучении. По умолчанию None (автоопределение).
input_size (int, optional) – Размер входного изображения (высота и ширина). Модель принимает изображения формы (batch, 3, input_size, input_size). По умолчанию 1280.
opset_version (int, optional) – Версия opset ONNX для экспорта. По умолчанию 14.
simplify (bool, optional) – Если True, применяется упрощение графа ONNX с помощью onnx-simplifier для оптимизации модели. Требуется пакет onnx-simplifier. По умолчанию True.

Результат:

Модель ONNX сохраняется в output_path.

Тип результата:

None

Исключение:

ImportError – Если необходимые пакеты (torch, onnx) не установлены.
FileNotFoundError – Если weights_path не существует.
ValueError – Если backbone_name не соответствует архитектуре чекпоинта.

Заметки

Экспортированная модель ONNX имеет два выхода:

score_map: карта уверенности текста формы (batch, 1, H, W)
geo_map: карта геометрии формы (batch, 8, H, W)

Модель поддерживает динамический размер батча и размеры изображения через настройку динамических осей.

Автоматическое определение backbone:

Метод автоматически определяет архитектуру backbone из чекпоинта, анализируя количество параметров в layer4. Это предотвращает несоответствия между чекпоинтом и архитектурой, которые могли бы привести к некорректному экспорту.

Примеры

Экспорт с автоматическим определением backbone:

>>> from manuscript.detectors import EAST
>>> EAST.export_to_onnx(
...     weights_path="east_resnet50.pth",
...     output_path="east_model.onnx"
... )
Auto-detected backbone: resnet50
Exporting to ONNX (opset 14)...
[OK] ONNX model saved to: east_model.onnx

Экспорт с явным указанием backbone:

>>> EAST.export_to_onnx(
...     weights_path="custom_weights.pth",
...     output_path="custom_model.onnx",
...     backbone_name="resnet101",
...     input_size=1024,
...     simplify=False
... )

Использование экспортированной модели для инференса:

>>> detector = EAST(
...     weights_path="east_model.onnx",
...     use_onnx=True,
...     device="cuda"
... )
>>> result = detector.predict("image.jpg")

См. также

EAST.__init__: Инициализация детектора EAST с поддержкой ONNX с использованием use_onnx=True.