text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  
  """
  DensePose Training Script.
  
  This script is similar to the training script in detectron2/tools.
  
  It is an example of how a user might use detectron2 for a new project.
  """
  
  import logging
  import os
  from collections import OrderedDict
  
  import detectron2.utils.comm as comm
  from detectron2.checkpoint import DetectionCheckpointer
  from detectron2.config import CfgNode, get_cfg
  from detectron2.engine import DefaultTrainer, default_argument_parser, default_setup, hooks, launch
  from detectron2.evaluation import COCOEvaluator, DatasetEvaluators, verify_results
  from detectron2.modeling import DatasetMapperTTA
  from detectron2.utils.logger import setup_logger
  
  from densepose import (
      DensePoseCOCOEvaluator,
      DensePoseGeneralizedRCNNWithTTA,
      add_dataset_category_config,
      add_densepose_config,
      load_from_cfg,
  )
  from densepose.data import DatasetMapper, build_detection_test_loader, build_detection_train_loader
  
  
  class Trainer(DefaultTrainer):
      @classmethod
      def build_evaluator(cls, cfg: CfgNode, dataset_name, output_folder=None):
          if output_folder is None:
              output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
          evaluators = [COCOEvaluator(dataset_name, cfg, True, output_folder)]
          if cfg.MODEL.DENSEPOSE_ON:
              evaluators.append(DensePoseCOCOEvaluator(dataset_name, True, output_folder))
          return DatasetEvaluators(evaluators)
  
      @classmethod
      def build_test_loader(cls, cfg: CfgNode, dataset_name):
          return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
  
      @classmethod
      def build_train_loader(cls, cfg: CfgNode):
          return build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
  
      @classmethod
      def test_with_TTA(cls, cfg: CfgNode, model):
          logger = logging.getLogger("detectron2.trainer")
          # In the end of training, run an evaluation with TTA
          # Only support some R-CNN models.
          logger.info("Running inference with test-time augmentation ...")
          transform_data = load_from_cfg(cfg)
          model = DensePoseGeneralizedRCNNWithTTA(cfg, model, transform_data, DatasetMapperTTA(cfg))
          evaluators = [
              cls.build_evaluator(
                  cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
              )
              for name in cfg.DATASETS.TEST
          ]
          res = cls.test(cfg, model, evaluators)
          res = OrderedDict({k + "_TTA": v for k, v in res.items()})
          return res
  
  
  def setup(args):
      cfg = get_cfg()
      add_dataset_category_config(cfg)
      add_densepose_config(cfg)
      cfg.merge_from_file(args.config_file)
      cfg.merge_from_list(args.opts)
      cfg.freeze()
      default_setup(cfg, args)
      # Setup logger for "densepose" module
      setup_logger(output=cfg.OUTPUT_DIR, distributed_rank=comm.get_rank(), name="densepose")
      return cfg
  
  
  def main(args):
      cfg = setup(args)
  
      if args.eval_only:
          model = Trainer.build_model(cfg)
          DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
              cfg.MODEL.WEIGHTS, resume=args.resume
          )
          res = Trainer.test(cfg, model)
          if cfg.TEST.AUG.ENABLED:
              res.update(Trainer.test_with_TTA(cfg, model))
          if comm.is_main_process():
              verify_results(cfg, res)
          return res
  
      trainer = Trainer(cfg)
      trainer.resume_or_load(resume=args.resume)
      if cfg.TEST.AUG.ENABLED:
          trainer.register_hooks(
              [hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
          )
      return trainer.train()
  
  
  if __name__ == "__main__":
      args = default_argument_parser().parse_args()
      print("Command Line Args:", args)
      launch(
          main,
          args.num_gpus,
          num_machines=args.num_machines,
          machine_rank=args.machine_rank,
          dist_url=args.dist_url,
          args=(args,),
      )
  