code: def main(args):
      print(args)
  
      if args.push_to_hub:
          login_to_hub()
  
      if not isinstance(args.workers, int):
          args.workers = min(16, mp.cpu_count())
  
      system_available_memory = int(
          psutil.virtual_memory().available / 1024**3
      )
  
      # AMP
      if args.amp:
          mixed_precision.set_global_policy("mixed_float16")
  
      st = time.time()
      val_set = DetectionDataset(
          img_folder=os.path.join(args.val_path, "images"),
          label_path=os.path.join(args.val_path, "labels.json"),
          sample_transforms=T.SampleCompose(
              (
                  [
                      T.Resize(
                          (args.input_size, args.input_size),
                          preserve_aspect_ratio=True,
                          symmetric_pad=True,
                      )
                  ]
                  if not args.rotation or args.eval_straight
                  else []
              )
              + (
                  [
                      T.Resize(
                          args.input_size, preserve_aspect_ratio=True
                      ),  # This does not pad
                      T.RandomApply(
                          T.RandomRotate(90, expand=True), 0.5
                      ),
                      T.Resize(
                          (args.input_size, args.input_size),
                          preserve_aspect_ratio=True,
                          symmetric_pad=True,
                      ),
                  ]
                  if args.rotation and not args.eval_straight
                  else []
              )
          ),
          use_polygons=args.rotation and not args.eval_straight,
      )
      val_loader = DataLoader(
          val_set,
          batch_size=args.batch_size,
          shuffle=False,
          drop_last=False,
          num_workers=args.workers,
      )
      print(
          f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
          f"{val_loader.num_batches} batches)"
      )
      with open(os.path.join(args.val_path, "labels.json"), "rb") as f:
          val_hash = hashlib.sha256(f.read()).hexdigest()
  
      batch_transforms = T.Compose(
          [
              T.Normalize(
                  mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287)
              ),
          ]
      )
  
      # Load doctr model
      model = detection.__dict__[args.arch](
          pretrained=args.pretrained,
          input_shape=(args.input_size, args.input_size, 3),
          assume_straight_pages=not args.rotation,
          class_names=val_set.class_names,
      )
  
      # Resume weights
      if isinstance(args.resume, str):
          model.load_weights(args.resume)
  
      if isinstance(args.pretrained_backbone, str):
          print("Loading backbone weights.")
          model = load_backbone(model, args.pretrained_backbone)
          print("Done.")
  
      # Metrics
      val_metric = LocalizationConfusion(
          use_polygons=args.rotation and not args.eval_straight,
          mask_shape=(args.input_size, args.input_size),
          use_broadcasting=True
          if system_available_memory > 62
          else False,
      )
      if args.test_only:
          print("Running evaluation")
          val_loss, recall, precision, mean_iou = evaluate(
              model, val_loader, batch_transforms, val_metric
          )
          print(
              f"Validation loss: {val_loss:.6} (Recall: {recall:.2%} | Precision: {precision:.2%} | "
              f"Mean IoU: {mean_iou:.2%})"
          )
          return
  
      st = time.time()
      # Load both train and val data generators
      train_set = DetectionDataset(
          img_folder=os.path.join(args.train_path, "images"),
          label_path=os.path.join(args.train_path, "labels.json"),
          img_transforms=T.Compose(
              [
                  # Augmentations
                  T.RandomApply(T.ColorInversion(), 0.1),
                  T.RandomJpegQuality(60),
                  T.RandomApply(
                      T.GaussianNoise(mean=0.1, std=0.1), 0.1
                  ),
                  T.RandomApply(T.RandomShadow(), 0.4),
                  T.RandomApply(
                      T.GaussianBlur(kernel_shape=3, std=(0.1, 0.1)),
                      0.3,
                  ),
                  T.RandomSaturation(0.3),
                  T.RandomContrast(0.3),
                  T.RandomBrightness(0.3),
                  T.RandomApply(T.ToGray(num_output_channels=3), 0.1),
              ]
          ),
          sample_transforms=T.SampleCompose(
              (
                  [
                      T.Resize(
                          (args.input_size, args.input_size),
                          preserve_aspect_ratio=True,
                          symmetric_pad=True,
                      )
                  ]
                  if not args.rotation
                  else []
              )
              + (
                  [
                      T.Resize(
                          args.input_size, preserve_aspect_ratio=True
                      ),  # This does not pad
                      T.RandomApply(
                          T.RandomRotate(90, expand=True), 0.5
                      ),
                      T.Resize(
                          (args.input_size, args.input_size),
                          preserve_aspect_ratio=True,
                          symmetric_pad=True,
                      ),
                  ]
                  if args.rotation
                  else []
              )
          ),
          use_polygons=args.rotation,
      )
      train_loader = DataLoader(
          train_set,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=True,
          num_workers=args.workers,
      )
      print(
          f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
          f"{train_loader.num_batches} batches)"
      )
      with open(
          os.path.join(args.train_path, "labels.json"), "rb"
      ) as f:
          train_hash = hashlib.sha256(f.read()).hexdigest()
  
      if args.show_samples:
          x, target = next(iter(train_loader))
          plot_samples(x, target)
          return
  
      # Optimizer
      scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
          args.lr,
          decay_steps=args.epochs * len(train_loader),
          decay_rate=1 / (25e4),  # final lr as a fraction of initial lr
          staircase=False,
          name="ExponentialDecay",
      )
      optimizer = tf.keras.optimizers.Adam(
          learning_rate=scheduler,
          beta_1=0.95,
          beta_2=0.99,
          epsilon=1e-6,
          clipnorm=5,
      )
      if args.amp:
          optimizer = mixed_precision.LossScaleOptimizer(optimizer)
      # LR Finder
      if args.find_lr:
          lrs, losses = record_lr(
              model,
              train_loader,
              batch_transforms,
              optimizer,
              amp=args.amp,
          )
          plot_recorder(lrs, losses)
          return
  
      # Tensorboard to monitor training
      current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      exp_name = (
          f"{args.arch}_{current_time}"
          if args.name is None
          else args.name
      )
  
      config = {
          "learning_rate": args.lr,
          "epochs": args.epochs,
          "batch_size": args.batch_size,
          "architecture": args.arch,
          "input_size": args.input_size,
          "optimizer": optimizer.name,
          "framework": "tensorflow",
          "scheduler": scheduler.name,
          "train_hash": train_hash,
          "val_hash": val_hash,
          "pretrained": args.pretrained,
          "rotation": args.rotation,
      }
  
      # W&B
      if args.wb:
          import wandb
  
          run = wandb.init(
              name=exp_name, project="text-detection", config=config
          )
  
      # ClearML
      if args.clearml:
          from clearml import Task
  
          task = Task.init(
              project_name="docTR/text-detection",
              task_name=exp_name,
              reuse_last_task_id=False,
          )
          task.upload_artifact("config", config)
  
      if args.freeze_backbone:
          for layer in model.feat_extractor.layers:
              layer.trainable = False
  
      min_loss = np.inf
      if args.early_stop:
          early_stopper = EarlyStopper(
              patience=args.early_stop_epochs,
              min_delta=args.early_stop_delta,
          )
  
      # Training loop
      for epoch in range(args.epochs):
          fit_one_epoch(
              model, train_loader, batch_transforms, optimizer, args.amp
          )
          # Validation loop at the end of each epoch
          val_loss, recall, precision, mean_iou = evaluate(
              model, val_loader, batch_transforms, val_metric
          )
          if val_loss < min_loss:
              print(
                  f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
              )
              model.save_weights(f"./{exp_name}/weights")
              min_loss = val_loss
          log_msg = f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} "
          if any(val is None for val in (recall, precision, mean_iou)):
              log_msg += "(Undefined metric value, caused by empty GTs or predictions)"
          else:
              log_msg += f"(Recall: {recall:.2%} | Precision: {precision:.2%} | Mean IoU: {mean_iou:.2%})"
          print(log_msg)
          # W&B
          if args.wb:
              wandb.log(
                  {
                      "val_loss": val_loss,
                      "recall": recall,
                      "precision": precision,
                      "mean_iou": mean_iou,
                  }
              )
  
          # ClearML
          if args.clearml:
              from clearml import Logger
  
              logger = Logger.current_logger()
              logger.report_scalar(
                  title="Validation Loss",
                  series="val_loss",
                  value=val_loss,
                  iteration=epoch,
              )
              logger.report_scalar(
                  title="Precision Recall",
                  series="recall",
                  value=recall,
                  iteration=epoch,
              )
              logger.report_scalar(
                  title="Precision Recall",
                  series="precision",
                  value=precision,
                  iteration=epoch,
              )
              logger.report_scalar(
                  title="Mean IoU",
                  series="mean_iou",
                  value=mean_iou,
                  iteration=epoch,
              )
          if args.early_stop and early_stopper.early_stop(val_loss):
              print(
                  "Training halted early due to reaching patience limit."
              )
              break
      if args.wb:
          run.finish()
  
      if args.push_to_hub:
          push_to_hf_hub(
              model, exp_name, task="detection", run_config=args
          )
  
imports: import os
  import datetime
  import hashlib
  import multiprocessing
  import time
  import numpy
  import psutil
  import tensorflow
  from tensorflow.keras import mixed_precision
  from tqdm.auto import tqdm
  from doctr.models import login_to_hub, push_to_hf_hub
  from doctr import transforms
  from doctr.datasets import DataLoader, DetectionDataset
  from doctr.models import detection
  from doctr.utils.metrics import LocalizationConfusion
  from utils import EarlyStopper, load_backbone, plot_recorder, plot_samples
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types