provide predefined element_length_fn for pipeline.bucket_boundaries

calamari_ocr/ocr/dataset/data.py

@@ -1,8 +1,9 @@
 import logging
 import os
-from typing import Type, Optional
+from typing import Callable, Dict, Type, Optional
 
 import tensorflow as tf
+from tfaip.util.tftyping import AnyTensor
 from tfaip.data.data import DataBase
 from tfaip.data.databaseparams import DataPipelineParams
 from tfaip.data.pipeline.datapipeline import DataPipeline
@@ -84,6 +85,12 @@ def _target_layer_specs(self):
             "gt_len": tf.TensorSpec([1], dtype=tf.int32),
         }
 
+    def element_length_fn(self) -> Callable[[Dict[str, AnyTensor]], AnyTensor]:
+        def img_len(x):
+            return x["img_len"]
+
+        return img_len
+
     def create_pipeline(
         self,
         pipeline_params: DataPipelineParams,

calamari_ocr/test/test_prediction.py

@@ -206,6 +206,11 @@ def test_dataset_prediction_voted(self):
 
         predictor.benchmark_results.pretty_print()
 
+    def test_batch_bucketing(self):
+        args = predict_args()
+        args.predictor.pipeline.bucket_boundaries = [20, 50, 100, 200, 400, 800]
+        run(args)
+
 
 if __name__ == "__main__":
     unittest.main()