code: def main():
      args = parse_args()
      batch_size = args.batch_size
      use_cuda = torch.cuda.is_available()
      hyperparams = vars(args)
      pprint(hyperparams)
  
      active_set, test_set = get_datasets(
          hyperparams["initial_pool"], hyperparams["data_path"]
      )
  
      # We will use the FocalLoss
      criterion = FocalLoss(gamma=2, alpha=0.25)
  
      # Our model is a simple Unet
      model = smp.Unet(
          encoder_name="resnext50_32x4d",
          encoder_depth=5,
          encoder_weights="imagenet",
          decoder_use_batchnorm=False,
          classes=len(pascal_voc_ids),
      )
      # Add a Dropout layerto use MC-Dropout
      add_dropout(model, classes=len(pascal_voc_ids), activation=None)
  
      # This will enable Dropout at test time.
      model = MCDropoutModule(model)
  
      # Put everything on GPU.
      if use_cuda:
          model.cuda()
  
      # Make an optimizer
      optimizer = optim.SGD(
          model.parameters(),
          lr=hyperparams["lr"],
          momentum=0.9,
          weight_decay=5e-4,
      )
      # Keep a copy of the original weights
      initial_weights = deepcopy(model.state_dict())
  
      # Add metrics
      model = ModelWrapper(model, criterion)
      model.add_metric(
          "cls_report",
          lambda: ClassificationReport(len(pascal_voc_ids)),
      )
  
      # Which heuristic you want to use?
      # We will use our custom reduction function.
      heuristic = get_heuristic(
          hyperparams["heuristic"], reduction=mean_regions
      )
  
      # The ALLoop is in charge of predicting the uncertainty and
      loop = ActiveLearningLoop(
          active_set,
          model.predict_on_dataset_generator,
          heuristic=heuristic,
          query_size=hyperparams["query_size"],
          # Instead of predicting on the entire pool, only a subset is used
          max_sample=1000,
          batch_size=batch_size,
          iterations=hyperparams["iterations"],
          use_cuda=use_cuda,
      )
      acc = []
      for epoch in tqdm(range(args.al_step)):
          # Following Gal et al. 2016, we reset the weights.
          model.load_state_dict(initial_weights)
          # Train 50 epochs before sampling.
          model.train_on_dataset(
              active_set,
              optimizer,
              batch_size,
              hyperparams["learning_epoch"],
              use_cuda,
          )
  
          # Validation!
          model.test_on_dataset(test_set, batch_size, use_cuda)
          should_continue = loop.step()
  
          logs = model.get_metrics()
          pprint(logs)
          acc.append(logs)
          if not should_continue:
              break
  
imports: import argparse
  from copy import deepcopy
  from pprint import pprint
  import torch.backends
  from PIL import Image
  from torch import optim
  from torchvision.transforms import transforms
  from tqdm import tqdm
  from baal import get_heuristic, ActiveLearningLoop
  from baal.bayesian.dropout import MCDropoutModule
  from baal import ModelWrapper
  from baal import ClassificationReport
  from baal import PILToLongTensor
  from utils import pascal_voc_ids, active_pascal, add_dropout, FocalLoss
  import torch
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: main() without return types