feat: add 2 new rotation flags in the ocr_predictor (#632)

* feat: add 2 new rot flags in ocr_predictor

* fix: style

* feat: add README update

* refcato: requested changes

* fix: kwargs error

README.md

@@ -56,6 +56,20 @@ doc = DocumentFile.from_pdf("path/to/your/doc.pdf").as_images()
 result = model(doc)
 ```
 
+### Dealing with rotated documents
+Should you use docTR on documents that include rotated pages, or pages with multiple box orientations,
+you have multiple options to handle it:
+
+- If you only use straight document pages with straight words (horizontal, same reading direction),
+consider passing `assume_straight_boxes=True` to the ocr_predictor. It will directly fit straight boxes
+on your page and return straight boxes, which makes it the fastest option.
+
+- If you want the predictor to output straight boxes (no matter the orientation of your pages, the final localizations
+will be converted to straight boxes), you need to pass `export_as_straight_boxes=True` in the predictor. Otherwise, if `assume_straight_pages=False`, it will return rotated bounding boxes (potentially with an angle of 0°).
+
+If both options are set to False, the predictor will always fit and return rotated boxes.
+
+
 To interpret your model's predictions, you can visualize them interactively as follows:
 
 ```python

doctr/models/predictor/pytorch.py

@@ -25,6 +25,11 @@ class OCRPredictor(nn.Module, _OCRPredictor):
     Args:
         det_predictor: detection module
         reco_predictor: recognition module
+        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
+            without rotated textual elements.
+        export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions
+            (potentially rotated) as straight bounding boxes.
+
     """
 
     def __init__(

doctr/models/predictor/tensorflow.py

@@ -25,6 +25,10 @@ class OCRPredictor(NestedObject, _OCRPredictor):
     Args:
         det_predictor: detection module
         reco_predictor: recognition module
+        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
+            without rotated textual elements.
+        export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions
+            (potentially rotated) as straight bounding boxes.
     """
     _children_names = ['det_predictor', 'reco_predictor']
 

doctr/models/zoo.py

@@ -12,21 +12,30 @@
 __all__ = ["ocr_predictor"]
 
 
-def _predictor(det_arch: str, reco_arch: str, pretrained: bool, det_bs=2, reco_bs=128) -> OCRPredictor:
+def _predictor(
+    det_arch: str,
+    reco_arch: str,
+    pretrained: bool,
+    det_bs: int = 2,
+    reco_bs: int = 128,
+    **kwargs,
+) -> OCRPredictor:
 
     # Detection
     det_predictor = detection_predictor(det_arch, pretrained=pretrained, batch_size=det_bs)
 
     # Recognition
     reco_predictor = recognition_predictor(reco_arch, pretrained=pretrained, batch_size=reco_bs)
 
-    return OCRPredictor(det_predictor, reco_predictor)
+    return OCRPredictor(det_predictor, reco_predictor, **kwargs)
 
 
 def ocr_predictor(
     det_arch: str = 'db_resnet50',
     reco_arch: str = 'crnn_vgg16_bn',
     pretrained: bool = False,
+    assume_straight_pages: bool = True,
+    export_as_straight_boxes: bool = False,
     **kwargs: Any
 ) -> OCRPredictor:
     """End-to-end OCR architecture using one model for localization, and another for text recognition.
@@ -39,11 +48,23 @@ def ocr_predictor(
         >>> out = model([input_page])
 
     Args:
-        arch: name of the architecture to use ('db_sar_vgg', 'db_sar_resnet', 'db_crnn_vgg', 'db_crnn_resnet')
+        det_arch: name of the detection architecture to use (e.g. 'db_resnet50', 'db_mobilenet_v3_large')
+        reco_arch: name of the recognition architecture to use (e.g. 'crnn_vgg16_bn', 'sar_resnet31')
         pretrained: If True, returns a model pre-trained on our OCR dataset
+        assume_straight_pages: if True, speeds up the inference by assuming you only pass straight pages
+            without rotated textual elements.
+        export_as_straight_boxes: when assume_straight_pages is set to False, export final predictions
+            (potentially rotated) as straight bounding boxes.
 
     Returns:
         OCR predictor
     """
 
-    return _predictor(det_arch, reco_arch, pretrained, **kwargs)
+    return _predictor(
+        det_arch,
+        reco_arch,
+        pretrained,
+        assume_straight_pages=assume_straight_pages,
+        export_as_straight_boxes=export_as_straight_boxes,
+        **kwargs,
+    )