code: def test_kiedocumentbuilder():
      num_pages = 2
  
      # Don't resolve lines
      doc_builder = builder.KIEDocumentBuilder(
          resolve_lines=False, resolve_blocks=False
      )
      pages = [np.zeros((100, 200, 3))] * num_pages
      predictions = {
          CLASS_NAME: np.random.rand(words_per_page, 6)
      }  # dict format
      predictions[CLASS_NAME][:2] *= predictions[CLASS_NAME][2:4]
      # Arg consistency check
      with pytest.raises(ValueError):
          doc_builder(
              pages,
              [predictions, predictions],
              [{CLASS_NAME: ("hello", 1.0)}] * 3,
              [(100, 200), (100, 200)],
          )
      out = doc_builder(
          pages,
          [predictions, predictions],
          [{CLASS_NAME: [("hello", 1.0)] * words_per_page}] * num_pages,
          [(100, 200), (100, 200)],
      )
      assert isinstance(out, KIEDocument)
      assert len(out.pages) == num_pages
      assert all(
          isinstance(page.page, np.ndarray) for page in out.pages
      ) and all(page.page.shape == (100, 200, 3) for page in out.pages)
      # 1 Block & 1 line per page
      assert len(out.pages[0].predictions) == 1
      assert len(out.pages[0].predictions[CLASS_NAME]) == words_per_page
  
      # Resolve lines
      doc_builder = builder.KIEDocumentBuilder(
          resolve_lines=True, resolve_blocks=True
      )
      out = doc_builder(
          pages,
          [predictions, predictions],
          [{CLASS_NAME: [("hello", 1.0)] * words_per_page}] * num_pages,
          [(100, 200), (100, 200)],
      )
  
      # No detection
      predictions = {CLASS_NAME: np.zeros((0, 5))}
      out = doc_builder(
          pages,
          [predictions, predictions],
          [{CLASS_NAME: []}, {CLASS_NAME: []}],
          [(100, 200), (100, 200)],
      )
      assert len(out.pages[0].predictions[CLASS_NAME]) == 0
  
      # Rotated boxes to export as straight boxes
      predictions = {
          CLASS_NAME: np.array(
              [
                  [[0.1, 0.1], [0.2, 0.2], [0.15, 0.25], [0.05, 0.15]],
                  [[0.5, 0.5], [0.6, 0.6], [0.55, 0.65], [0.45, 0.55]],
              ]
          )
      }
      doc_builder_2 = builder.KIEDocumentBuilder(
          resolve_blocks=False,
          resolve_lines=False,
          export_as_straight_boxes=True,
      )
      out = doc_builder_2(
          [np.zeros((100, 100, 3))],
          [predictions],
          [{CLASS_NAME: [("hello", 0.99), ("word", 0.99)]}],
          [(100, 100)],
      )
      assert out.pages[0].predictions[CLASS_NAME][0].geometry == (
          (0.05, 0.1),
          (0.2, 0.25),
      )
      assert out.pages[0].predictions[CLASS_NAME][1].geometry == (
          (0.45, 0.5),
          (0.6, 0.65),
      )
  
      # Repr
      assert (
          repr(doc_builder) == "KIEDocumentBuilder(resolve_lines=True, "
          "resolve_blocks=True, paragraph_break=0.035, export_as_straight_boxes=False)"
      )
  
file_path: doctr/tests/common/test_models_builder.py
imports: import numpy
  import pytest
  from doctr.file_utils import CLASS_NAME
  from doctr.io import Document
  from doctr.io.elements import KIEDocument
  from doctr.models import builder
label_desc: Use image node_id 2 for calling a global function with example usage: test_kiedocumentbuilder() without return types