legacies/doctr

1# Copyright (C) 2021-2024, Mindee.2 3# This program is licensed under the Apache License 2.0.4# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.5 6import numpy as np7import torch8 9from doctr.models import ocr_predictor10from doctr.models.predictor import OCRPredictor11 12DET_ARCHS = [13    "db_resnet50",14    "db_resnet34",15    "db_mobilenet_v3_large",16    "linknet_resnet18",17    "linknet_resnet34",18    "linknet_resnet50",19]20RECO_ARCHS = [21    "crnn_vgg16_bn",22    "crnn_mobilenet_v3_small",23    "crnn_mobilenet_v3_large",24    "master",25    "sar_resnet31",26    "vitstr_small",27    "vitstr_base",28    "parseq",29]30 31 32def load_predictor(33    det_arch: str,34    reco_arch: str,35    assume_straight_pages: bool,36    straighten_pages: bool,37    bin_thresh: float,38    box_thresh: float,39    device: torch.device,40) -> OCRPredictor:41    """Load a predictor from doctr.models42 43    Args:44    ----45        det_arch: detection architecture46        reco_arch: recognition architecture47        assume_straight_pages: whether to assume straight pages or not48        straighten_pages: whether to straighten rotated pages or not49        bin_thresh: binarization threshold for the segmentation map50        device: torch.device, the device to load the predictor on51 52    Returns:53    -------54        instance of OCRPredictor55    """56    predictor = ocr_predictor(57        det_arch,58        reco_arch,59        pretrained=True,60        assume_straight_pages=assume_straight_pages,61        straighten_pages=straighten_pages,62        export_as_straight_boxes=straighten_pages,63        detect_orientation=not assume_straight_pages,64    ).to(device)65    predictor.det_predictor.model.postprocessor.bin_thresh = bin_thresh66    predictor.det_predictor.model.postprocessor.box_thresh = box_thresh67    return predictor68 69 70def forward_image(predictor: OCRPredictor, image: np.ndarray, device: torch.device) -> np.ndarray:71    """Forward an image through the predictor72 73    Args:74    ----75        predictor: instance of OCRPredictor76        image: image to process77        device: torch.device, the device to process the image on78 79    Returns:80    -------81        segmentation map82    """83    with torch.no_grad():84        processed_batches = predictor.det_predictor.pre_processor([image])85        out = predictor.det_predictor.model(processed_batches[0].to(device), return_model_output=True)86        seg_map = out["out_map"].to("cpu").numpy()87 88    return seg_map89