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())
  
      torch.backends.cudnn.benchmark = True
      system_available_memory = int(
          psutil.virtual_memory().available / 1024**3
      )
  
      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,
          drop_last=False,
          num_workers=args.workers,
          sampler=SequentialSampler(val_set),
          pin_memory=torch.cuda.is_available(),
          collate_fn=val_set.collate_fn,
      )
      print(
          f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
          f"{len(val_loader)} batches)"
      )
      with open(os.path.join(args.val_path, "labels.json"), "rb") as f:
          val_hash = hashlib.sha256(f.read()).hexdigest()
  
      batch_transforms = 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,
          assume_straight_pages=not args.rotation,
          class_names=val_set.class_names,
      )
  
      # Resume weights
      if isinstance(args.resume, str):
          print(f"Resuming {args.resume}")
          checkpoint = torch.load(args.resume, map_location="cpu")
          model.load_state_dict(checkpoint)
  
      # GPU
      if isinstance(args.device, int):
          if not torch.cuda.is_available():
              raise AssertionError(
                  "PyTorch cannot access your GPU. Please investigate!"
              )
          if args.device >= torch.cuda.device_count():
              raise ValueError("Invalid device index")
      # Silent default switch to GPU if available
      elif torch.cuda.is_available():
          args.device = 0
      else:
          logging.warning("No accessible GPU, targe device set to CPU.")
      if torch.cuda.is_available():
          torch.cuda.set_device(args.device)
          model = model.cuda()
  
      # 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,
              amp=args.amp,
          )
          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=Compose(
              [
                  # Augmentations
                  T.RandomApply(T.ColorInversion(), 0.1),
                  T.RandomApply(
                      T.GaussianNoise(mean=0.1, std=0.1), 0.1
                  ),
                  T.RandomApply(T.RandomShadow(), 0.4),
                  T.RandomApply(GaussianBlur(kernel_size=3), 0.3),
                  RandomPhotometricDistort(p=0.1),
                  RandomGrayscale(p=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
                      ),
                      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,
          drop_last=True,
          num_workers=args.workers,
          sampler=RandomSampler(train_set),
          pin_memory=torch.cuda.is_available(),
          collate_fn=train_set.collate_fn,
      )
      print(
          f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
          f"{len(train_loader)} 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
  
      # Backbone freezing
      if args.freeze_backbone:
          for p in model.feat_extractor.parameters():
              p.requires_grad = False
  
      # Optimizer
      optimizer = torch.optim.Adam(
          [p for p in model.parameters() if p.requires_grad],
          args.lr,
          betas=(0.95, 0.99),
          eps=1e-6,
          weight_decay=args.weight_decay,
      )
      # LR Finder
      if args.find_lr:
          lrs, losses = record_lr(
              model,
              train_loader,
              batch_transforms,
              optimizer,
              amp=args.amp,
          )
          plot_recorder(lrs, losses)
          return
      # Scheduler
      if args.sched == "cosine":
          scheduler = CosineAnnealingLR(
              optimizer,
              args.epochs * len(train_loader),
              eta_min=args.lr / 25e4,
          )
      elif args.sched == "onecycle":
          scheduler = OneCycleLR(
              optimizer, args.lr, args.epochs * len(train_loader)
          )
  
      # Training monitoring
      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
      )
  
      # W&B
      if args.wb:
          run = wandb.init(
              name=exp_name,
              project="text-detection",
              config={
                  "learning_rate": args.lr,
                  "epochs": args.epochs,
                  "weight_decay": args.weight_decay,
                  "batch_size": args.batch_size,
                  "architecture": args.arch,
                  "input_size": args.input_size,
                  "optimizer": "adam",
                  "framework": "pytorch",
                  "scheduler": args.sched,
                  "train_hash": train_hash,
                  "val_hash": val_hash,
                  "pretrained": args.pretrained,
                  "rotation": args.rotation,
                  "amp": args.amp,
              },
          )
  
      # Create loss queue
      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,
              scheduler,
              amp=args.amp,
          )
          # Validation loop at the end of each epoch
          val_loss, recall, precision, mean_iou = evaluate(
              model,
              val_loader,
              batch_transforms,
              val_metric,
              amp=args.amp,
          )
          if val_loss < min_loss:
              print(
                  f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
              )
              torch.save(model.state_dict(), f"./{exp_name}.pt")
              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,
                  }
              )
          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 logging
  import multiprocessing
  import time
  import numpy
  import psutil
  import torch
  import wandb
  from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  from torchvision.transforms.v2 import Compose, GaussianBlur, Normalize, RandomGrayscale, RandomPhotometricDistort
  from tqdm.auto import tqdm
  from doctr import transforms
  from doctr.datasets import DetectionDataset
  from doctr.models import detection, login_to_hub, push_to_hf_hub
  from doctr.utils.metrics import LocalizationConfusion
  from utils import EarlyStopper, plot_recorder, plot_samples
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types