How can I run the Orientation model solely to determine the page rotation angle?

[sanjay-nit]

Hi Team,

FYI: I'm using docTR (0.10.0) with the following configuration to pass an image and retrieve the page orientation. I'm satisfied with the results from the ocr_predictor. However, to save time, I would like to avoid running the full OCR model. Is there a way to use only the orientation model?
config:

doctr_model = ocr_predictor(
    pretrained=True,
    assume_straight_pages=False,
    straighten_pages=True,
    detect_orientation=True,
    export_as_straight_boxes=True
)

Note: I tried page_orientation_predictor from doctr.models but looks like it only gives one of following angles. 0, 90, 180, 270 (-90) as mentioned in code.

Also, is there any guide available for training a custom orientation model? I checked this page, but it doesn’t clearly explain how to properly prepare the dataset for training, eg: how our images looks like? Thanks!!

[felixdittrich92]

Hi @sanjay-nit 👋,

That's right our orientation prediction depends on the detection model:

1. Estimate general page orientation -> [0, 90, 180, 270 (-90)]
2. Correct the input
3. Detect all text and get the segmentation map
4. Estimate the "fine grained" orientation

Here is a short snippet how you could use it without the recognition part:

import requests
import numpy as np

from doctr.io import DocumentFile
from doctr.models import detection_predictor, page_orientation_predictor
from doctr.models._utils import estimate_orientation


url = "https://www.francetvinfo.fr/pictures/uGwaNE-aJq7zHLhZJdzdCd9nyjE/1200x900/2021/03/16/phpCDwGn0.jpg"

det_predictor = detection_predictor(
    arch="fast_base",
    pretrained=True,
    assume_straight_pages=False,
)  #.cuda().half()  # Uncomment if running on GPU
page_orient_predictor = page_orientation_predictor(pretrained=True)  #.cuda().half()  # Uncomment if running on GPU

det_predictor.model.postprocessor.bin_thresh = 0.3
det_predictor.model.postprocessor.box_thresh = 0.65

docs = DocumentFile.from_images([requests.get(url).content])
loc_preds, out_maps = det_predictor(docs, return_maps=True)
seg_maps = [
    np.where(out_map > getattr(det_predictor.model.postprocessor, "bin_thresh"), 255, 0).astype(np.uint8)
    for out_map in out_maps
]
_, classes, probs = zip(page_orient_predictor(docs))
# Flatten to list of tuples with (value, confidence)
page_orientations = [
    (orientation, prob)
    for page_classes, page_probs in zip(classes, probs)
    for orientation, prob in zip(page_classes, page_probs)
]
origin_pages_orientations = [
    estimate_orientation(seq_map, general_orientation)
    for seq_map, general_orientation in zip(seg_maps, page_orientations)
]
orientations = [
    {"value": orientation, "confidence": prob} for orientation, prob in zip(origin_pages_orientations, probs[0])
]
print(orientations)

[sanjay-nit]

Thanks @felixdittrich92
It works!!!