code: def test_detection_models(
      arch_name, input_shape, output_size, out_prob, train_mode
  ):
      batch_size = 2
      tf.keras.backend.clear_session()
      model = detection.__dict__[arch_name](
          pretrained=True, input_shape=input_shape
      )
      assert isinstance(model, tf.keras.Model)
      input_tensor = tf.random.uniform(
          shape=[batch_size, *input_shape], minval=0, maxval=1
      )
      target = [
          {
              CLASS_NAME: np.array(
                  [[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.8]],
                  dtype=np.float32,
              )
          },
          {
              CLASS_NAME: np.array(
                  [[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.9]],
                  dtype=np.float32,
              )
          },
      ]
      # test training model
      out = model(
          input_tensor,
          target,
          return_model_output=True,
          return_preds=not train_mode,
          training=train_mode,
      )
      assert isinstance(out, dict)
      assert len(out) == 3 if not train_mode else len(out) == 2
      # Check proba map
      assert isinstance(out["out_map"], tf.Tensor)
      assert out["out_map"].dtype == tf.float32
      seg_map = out["out_map"].numpy()
      assert seg_map.shape == (batch_size, *output_size)
      if out_prob:
          assert np.all(np.logical_and(seg_map >= 0, seg_map <= 1))
      # Check boxes
      if not train_mode:
          for boxes_dict in out["preds"]:
              for boxes in boxes_dict.values():
                  assert boxes.shape[1] == 5
                  assert np.all(boxes[:, :2] < boxes[:, 2:4])
                  assert np.all(boxes[:, :4] >= 0) and np.all(
                      boxes[:, :4] <= 1
                  )
      # Check loss
      assert isinstance(out["loss"], tf.Tensor)
      # Target checks
      target = [
          {CLASS_NAME: np.array([[0, 0, 1, 1]], dtype=np.uint8)},
          {CLASS_NAME: np.array([[0, 0, 1, 1]], dtype=np.uint8)},
      ]
      with pytest.raises(AssertionError):
          out = model(input_tensor, target, training=True)
  
      target = [
          {CLASS_NAME: np.array([[0, 0, 1.5, 1.5]], dtype=np.float32)},
          {
              CLASS_NAME: np.array(
                  [[-0.2, -0.3, 1, 1]], dtype=np.float32
              )
          },
      ]
      with pytest.raises(ValueError):
          out = model(input_tensor, target, training=True)
  
      # Check the rotated case
      target = [
          {
              CLASS_NAME: np.array(
                  [
                      [0.75, 0.75, 0.5, 0.5, 0],
                      [0.65, 0.65, 0.3, 0.3, 0],
                  ],
                  dtype=np.float32,
              )
          },
          {
              CLASS_NAME: np.array(
                  [[0.75, 0.75, 0.5, 0.5, 0], [0.65, 0.7, 0.3, 0.4, 0]],
                  dtype=np.float32,
              )
          },
      ]
      loss = model(input_tensor, target, training=True)["loss"]
      assert (
          isinstance(loss, tf.Tensor)
          and ((loss - out["loss"]) / loss).numpy() < 25e-2
      )
  
imports: import os
  import tempfile
  import numpy
  import onnxruntime
  import psutil
  import pytest
  import tensorflow
  from doctr.file_utils import CLASS_NAME
  from doctr.io import DocumentFile
  from doctr.models import detection
  from doctr.models.detection._utils import dilate, erode
  from doctr.models.detection.predictor import DetectionPredictor
  from doctr.models.preprocessor import PreProcessor
  from doctr.models.utils import export_model_to_onnx
label_desc: Use image node_id 1 for calling a global function with example usage: test_detection_models(arch_name, input_shape, output_size, out_prob, train_mode) without return types