Improving the orientation classification model

[mislav-zane]

In a project of mine I want to read some randomly rotated numbers from grayscale images.
I am working with db_resnet34 as the detection model, and with crnn_vgg16_bn as the recognition model.

Test Runs

Test Run 1 -- Large Set

I am initializing the ocr_predictor with

predictor = ocr_predictor(
    det_arch="db_resnet34",
    reco_arch="crnn_vgg16_bn",
    pretrained=True,
    assume_straight_pages=False,
    straighten_pages=True,
)

Then I am running the predictor on the preprocessed images. Running this on a dataset of 700 images gives me a success rate of about 52%.

Example from the used data set:

Test Run 2 -- Straightened Set

For comparison I manually rotated a randomly selected subset of 40 images into horizontally straight orientation. Then I ran the predictor initialized with

predictor = ocr_predictor(
    det_arch="db_resnet34",
    reco_arch="crnn_vgg16_bn",
    pretrained=True,
    assume_straight_pages=True,
    straighten_pages=False,
)

on these straightened images. The success rate was at 95%.

Example from the used data set (same image from above, but straightened):

Test Run 3 -- Non-Straightened Set

Running the same subset of 40 images, but non-straightened with the same predictor as Test Run 1 yields a success rate of 55%.

My Question

I have two questions:

1. Is the orientation model designed to work on arbitrary angles, and this kind of images?
2. If yes, do I need to train it on the dataset, or is it supposed to work out of the box?

[felixdittrich92]

Hi @mislav-zane 👋,

1. The page orientation model predicts only the general orientation (-90, 0, 90, 180) afterwards in combination with the predicted segmentation mask from the detection model we calculate the fine grained rotation angle. Mh the kind of image is really special such kind of data was definitely not in the training dataset.
2. Yeah I suggest to fine tune the model on your data in this case see: https://mindee.github.io/doctr/using_doctr/custom_models_training.html#custom-orientation-classification-models
   Only mobilenet_v3_small is actually supported as orientation model additional the dataset should contain only straight images (like the first example you have shared) because we auto-augment the labels by randomly rotating the images and capturing the the angle as label.

Best regards,
Felix

[haimat]

Thanks for your quick response, @felixdittrich92. I have two more questions on the training process:

1. You said our training dataset for the orientation model should only contain straight images, but then you referred to the first example, which is the original, rotated (i.e. non-straight) image. Could you please clarify which images we should use as dataset for the training of the orientation model?
2. You also said the orientation model predicts only general orientation in 90 degrees steps. But as you can see in our examples, we have to deal with other angles as well. When we train the orientation model on our dataset, will this customly trained model then also understand our rotations that are not one of -90, 0, 90, 180?

[felixdittrich92]

Hi @haimat 👋,

1. Yes the original image also if the text is rotated
2. The workflow is as follows:
   1. Predict the general orientation 90 degree steps
   2. detect all text -> segmentation mask
   3. pre-rotate the sample by the given general orientation if min confidence is higher
   4. calculate the fine-grained angle on the segmentation mask
   5. combine both to get the final orientation angle

Hope this helps :)
PS: I think in your case the general page orientation fails you can also try the following before fine-tuning

predictor = ocr_predictor(
    det_arch="db_resnet34",
    reco_arch="crnn_vgg16_bn",
    pretrained=True,
    assume_straight_pages=False,
    straighten_pages=True,
    disable_page_orientation=True,   #  <-- This
    # disable_crop_orientation=True,
)