code: def main():
      args = parse_args()
      reset_config(config, args)
      set_cudnn(config)
      seed = config.RANDOM_SEED
      np.random.seed(seed)
      random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
      os.environ["CUDA_VISIBLE_DEVICES"] = config.GPU.strip()
      gpus = list(range(len(config.GPU.strip().split(","))))
  
      logger, final_output_dir, tb_log_dir = create_logger(
          config, args.cfg
      )
      summary_writer = SummaryWriter(log_dir=tb_log_dir)
  
      this_dir = osp.dirname(__file__)
      # backup the source code and the yaml config
      if args.cfg:
          shutil.copy(
              args.cfg,
              osp.join(final_output_dir, osp.basename(args.cfg)),
          )
      if not osp.exists(osp.join(final_output_dir, "lib")):
          shutil.copytree(
              osp.join(this_dir, "../lib/"),
              osp.join(final_output_dir, "lib"),
          )
      for k, v in config.items():
          logger.info(f"{k}: {v}")
  
      # conditional import
      if config.TRAIN.FINETUNE_ROTATER:
          from lib.core.function3 import train, validate, evaluate
      elif config.TRAIN.USE_CYCLE:
          from lib.core.function2 import train, validate, evaluate
      else:
          from lib.core.function1 import train, validate, evaluate
  
      # build model
      logger.info("start building model.")
      if len(gpus) > 1:
          pose_model = torch.nn.DataParallel(
              get_pose_model(config)
          ).cuda(gpus[0])
          discriminator = torch.nn.DataParallel(
              get_discriminator(config)
          ).cuda(gpus[0])
          temp_discriminator = torch.nn.DataParallel(
              get_discriminator(config)
          ).cuda(gpus[0])
      else:
          pose_model = get_pose_model(config).cuda()
          discriminator = get_discriminator(
              config, is_temp=False
          ).cuda()
          temp_discriminator = get_discriminator(
              config, is_temp=True
          ).cuda()
      optimizer_g = get_optimizer(config, pose_model, is_dis=False)
      optimizer_d = get_optimizer(config, discriminator, is_dis=True)
      optimizer_d_temp = get_optimizer(
          config, temp_discriminator, is_dis=True, is_temp=True
      )
      step_size, gamma = (
          config.TRAIN.SCHEDULER_STEP_SIZE,
          config.TRAIN.SCHEDULER_GAMMA,
      )
      scheduler_g = lr_scheduler.StepLR(
          optimizer_g, step_size=step_size, gamma=gamma
      )
      scheduler_d = lr_scheduler.StepLR(
          optimizer_d, step_size=step_size, gamma=gamma
      )
      scheduler_temp = lr_scheduler.StepLR(
          optimizer_d_temp, step_size=step_size, gamma=gamma
      )
      logger.info("finished building model.")
      # print out the model arch
      if config.TRAIN.PRETRAIN_LIFTER:
          print("Load pretrained lifter...")
          state_dict = torch.load(config.TRAIN.LIFTER_PRETRAIN_PATH)[
              "pose_model_state_dict"
          ]
          # state_dict = {k[7:]:v for k, v in state_dict.items()}
          pose_model.load_state_dict(state_dict, strict=False)
  
      if config.DATA.DATASET_NAME == "surreal":
          loader_func = surreal
      else:
          loader_func = (
              h36m if config.DATA.DATASET_NAME == "h36m" else mpiinf
          )
      dataset_train = loader_func(config, is_train=True)
      dataset_test = loader_func(config, is_train=False)
  
      train_loader = DataLoader(
          dataset=dataset_train,
          batch_size=config.BATCH_SIZE,
          shuffle=True,
          drop_last=False,
          pin_memory=True,
          num_workers=config.NUM_WORKERS,
      )
      test_loader = DataLoader(
          dataset=dataset_test,
          batch_size=config.BATCH_SIZE,
          shuffle=False,
          drop_last=False,
          pin_memory=True,
          num_workers=config.NUM_WORKERS,
      )
  
      if args.eval:
          prefix = config.DATA.DATASET_NAME
          # for mode in ['train', 'valid']:
          for mode in ["valid"]:
              is_train = True if mode == "train" else False
              v3d_to_ours = (
                  [
                      3,
                      2,
                      1,
                      4,
                      5,
                      6,
                      16,
                      15,
                      14,
                      11,
                      12,
                      13,
                      8,
                      0,
                      7,
                      9,
                      10,
                  ]
                  if prefix == "h36m"
                  else np.arange(config.DATA.NUM_JOINTS)
              )
              mpi2h36m = [
                  10,
                  9,
                  8,
                  11,
                  12,
                  13,
                  4,
                  3,
                  2,
                  5,
                  6,
                  7,
                  1,
                  14,
                  15,
                  16,
                  0,
              ]
              if prefix == "surreal":
                  indices = np.arange(config.DATA.NUM_JOINTS)
              else:
                  indices = (
                      v3d_to_ours if prefix == "h36m" else mpi2h36m
                  )
              mode = "train" if is_train else "valid"
              read_name = f"../data/{prefix}_{mode}_pred3.h5"
              # read_name = f"../../unsupervised_mesh/data/h36m_{mode}_pred_3d_mesh.h5"
              save_name = f"../data/{prefix}_{mode}_pred_3d.h5"
              if args.eval_suffix is not None:
                  save_name = f"{save_name[:-3]}_{args.eval_suffix}.h5"
  
              # eval mode, load the pretrained model and generate the 3d prediction of all 3ds
              if not config.TRAIN.PRETRAIN_LIFTER:
                  raise Warning(
                      "You are not using a pretrain model... may be you can specify --pretrain flag"
                  )
              dataloader = DataLoader(
                  dataset_train if mode == "train" else dataset_test,
                  batch_size=config.BATCH_SIZE,
                  shuffle=False,
                  drop_last=False,
                  pin_memory=True,
                  num_workers=config.NUM_WORKERS,
              )
              all_out_data = evaluate(
                  dataloader,
                  pose_model,
                  config,
                  is_train=(mode == "train"),
              )
              p1_mpjpe, p2_mpjpe = (
                  all_out_data["p1_mpjpe"],
                  all_out_data["p2_mpjpe"],
              )
              # read out imagenames
              print("Reading imagenames and joints 2d...")
              fin = h5py.File(read_name, "r")
              fout = h5py.File(save_name, "w")
              imagenames = fin["imagename"][:].copy()
              joints_2d_gt = np.array(fin["joint_2d_gt"])
              fout["imagename"] = imagenames
              fout["joint_2d_gt"] = joints_2d_gt[:, indices]
              fout["joint_3d_gt"] = all_out_data["joint_3d_gt"]
              fout["joint_3d_pre"] = all_out_data["joint_3d_pre"]
              possible_same_keys = [
                  "shape",
                  "pose",
                  "original_joint_2d_gt",
                  "joint_2d_pre",
                  "seqlen",
              ]
  
              for key in possible_same_keys:
                  if key in fin.keys():
                      if "joint" in key:
                          fout[key] = np.array(fin[key])[:, indices]
                      else:
                          fout[key] = np.array(fin[key])
              if "seqname" in fin.keys():
                  fout["seqname"] = fin["seqname"][:].copy()
  
              if "auc" in all_out_data.keys():
                  fout["auc"] = all_out_data["auc"]
                  fout["pckh5"] = all_out_data["pckh5"]
                  fout["auc_p2"] = all_out_data["auc_p2"]
                  fout["pckh5_p2"] = all_out_data["pckh5_p2"]
              if "scales" in all_out_data.keys():
                  fout["scale_pre"] = all_out_data["scales"]
              if "scale_mids" in all_out_data.keys():
                  fout["scale_mid_pre"] = all_out_data["scale_mids"]
  
              fin.close()
              fout.close()
              print(
                  "Evaluation on the {} set finished. P1 Mpjpe: {:.3f}, P2 Mpjpe: {:.3f}, saved to {}".format(
                      "training" if is_train else "test",
                      p1_mpjpe,
                      p2_mpjpe,
                      save_name,
                  )
              )
              if prefix == "mpi":
                  print(
                      "PCKh@0.5: {:.3f}, AUC: {:.3f}".format(
                          all_out_data["pckh5"], all_out_data["auc"]
                      )
                  )
                  print(
                      "P2: PCKh@0.5: {:.3f}, AUC: {:.3f}".format(
                          all_out_data["pckh5_p2"],
                          all_out_data["auc_p2"],
                      )
                  )
          # uncomment this if you need to plot images
          # print("Rendering sequences...")
          # subprocess.call(f'python render.py --seq_num 10 --in_filename ../data/{prefix}_valid_pred_3d.h5 --save_dir ../vis', shell=True)
          return
  
      # preparation for visualization & perseq optimization(optional)
      if config.USE_GT:
          # note that the gt here is not the gt above(config.USE_GT)
          train_path = (
              f"../data/{config.DATA.DATASET_NAME}_train_scales.pkl"
          )
          valid_path = (
              f"../data/{config.DATA.DATASET_NAME}_valid_scales.pkl"
          )
      else:
          train_path = (
              f"../data/{config.DATA.DATASET_NAME}_train_scales_pre.pkl"
          )
          valid_path = (
              f"../data/{config.DATA.DATASET_NAME}_valid_scales_pre.pkl"
          )
  
      train_scale_mids_gt = (
          load_pickle(train_path)["scale_mid"]
          if osp.exists(train_path)
          else None
      )
      valid_scale_mids_gt = (
          load_pickle(valid_path)["scale_mid"]
          if osp.exists(valid_path)
          else None
      )
      train_seqnames, valid_seqnames = (
          dataset_train.get_seqnames(),
          dataset_test.get_seqnames(),
      )
      best_p1_mpjpe = best_p2_mpjpe = cur_p1_mpjpe = 10000.0
      best_auc_val = best_pckh5 = 0.0
      best_auc_val_p2 = best_pckh5_p2 = 0.0
  
      for epoch in range(config.TRAIN.NUM_EPOCHS):
          scheduler_d.step()
          scheduler_g.step()
          scheduler_temp.step()
          # scheduler_s.step()
          (
              avg_d_loss,
              avg_g_loss,
              avg_t_loss,
              train_scale_mids_pre,
          ) = train(
              train_loader,
              pose_model,
              discriminator,
              temp_discriminator,
              optimizer_g,
              optimizer_d,
              optimizer_d_temp,
              epoch,
              config,
              summary_writer=summary_writer,
              print_interval=config.PRINT_INTERVAL,
          )
          logger.info(
              "***** Epoch: {}, Avg G Loss: {:.3f}, Avg D Loss: {:.3f} Avg T Loss: {:.3f} *****".format(
                  epoch, avg_g_loss, avg_d_loss, avg_t_loss
              )
          )
          (
              p1_mpjpe,
              p2_mpjpe,
              vis_image,
              valid_scale_mids_pre,
              extra_dict,
          ) = validate(test_loader, pose_model, epoch, config)
          logger.info(
              "Epoch: {}, P1 Mpjpe/Best P1: {:.3f}/{:.3f}, P2 Mpjpe/Best P2/Cur P1: {:.3f}/{:.3f}/{:.3f}".format(
                  epoch,
                  p1_mpjpe,
                  best_p1_mpjpe,
                  p2_mpjpe,
                  best_p2_mpjpe,
                  cur_p1_mpjpe,
              )
          )
          if p2_mpjpe < best_p2_mpjpe:
              best_p2_mpjpe = p2_mpjpe
              cur_p1_mpjpe = p1_mpjpe
              is_best = True
          else:
              is_best = False
  
          if p1_mpjpe < best_p1_mpjpe:
              best_p1_mpjpe = p1_mpjpe
  
          if extra_dict is not None:
              auc_val, pckh5 = extra_dict["auc"], extra_dict["pckh5"]
              auc_val_p2, pckh5_p2 = (
                  extra_dict["auc_p2"],
                  extra_dict["pckh5_p2"],
              )
              if auc_val_p2 > best_auc_val_p2:
                  best_auc_val_p2 = auc_val_p2
                  best_pckh5_p2 = pckh5_p2
                  is_best = True
              else:
                  is_best = False
  
              if auc_val > best_auc_val:
                  best_auc_val = auc_val
                  best_pckh5 = pckh5
              logger.info(
                  "PCKh@0.5(Best): {:.3f}({:.3f}), AUC value(Best): {:.3f}({:.3f})".format(
                      pckh5, best_pckh5, auc_val, best_auc_val
                  )
              )
              logger.info(
                  "P2: PCKh@0.5(Best): {:.3f}({:.3f}), AUC value(Best): {:.3f}({:.3f})".format(
                      pckh5_p2,
                      best_pckh5_p2,
                      auc_val_p2,
                      best_auc_val_p2,
                  )
              )
  
          save_checkpoint(
              {
                  "epoch": epoch,
                  "auc": best_auc_val,
                  "pckh5": best_pckh5,
                  "auc_p2": best_auc_val_p2,
                  "pckh5_p2": best_pckh5_p2,
                  "p1_mpjpe": p1_mpjpe,
                  "p2_mpjpe": p2_mpjpe,
                  "pose_model_state_dict": pose_model.state_dict(),
                  "discriminator_state_dict": discriminator.state_dict(),
                  "temp_discriminator_state_dict": temp_discriminator.state_dict(),
                  "optimizer_d": optimizer_d.state_dict(),
                  "optimizer_g": optimizer_g.state_dict(),
                  "optimizer_d_temp": optimizer_d_temp.state_dict(),
              },
              is_best,
              final_output_dir,
          )
          summary_writer.add_scalar(
              "p1_mpjpe_3d_test/epoch", p1_mpjpe, epoch
          )
          summary_writer.add_scalar(
              "p2_mpjpe_3d_test/epoch", p2_mpjpe, epoch
          )
          summary_writer.add_image(
              "test_joints/epoch", vis_image, epoch
          )
          if extra_dict is not None:
              summary_writer.add_scalar("PCKh0.5/epoch", pckh5, epoch)
              summary_writer.add_scalar("AUC/epoch", auc_val, epoch)
  
          if (
              train_scale_mids_gt is not None
              and train_scale_mids_pre is not None
              and len(train_scale_mids_pre) > 0
          ):
              num_seq = config.VIS.SCALE_MID_NUM_SEQ
              vis_image_scale_mid1 = plot_scalemid_dist(
                  train_scale_mids_pre, train_scale_mids_gt.tolist()
              )
              vis_image_scale_mid1 = (
                  torch.from_numpy(vis_image_scale_mid1)
                  .type(torch.float32)
                  .permute(2, 0, 1)
                  / 255
              )
              vis_image_scale_mid2 = plot_scalemid_seq_dist(
                  train_scale_mids_pre,
                  train_scale_mids_gt.tolist(),
                  train_seqnames,
                  num_seq=num_seq,
              )
              vis_image_scale_mid2 = (
                  torch.from_numpy(vis_image_scale_mid2)
                  .type(torch.float32)
                  .permute(2, 0, 1)
                  / 255
              )
              summary_writer.add_image(
                  "train_scalemid_distribution/epoch",
                  vis_image_scale_mid1,
                  epoch,
              )
              summary_writer.add_image(
                  "train_scalemid_seq_distribution/epoch",
                  vis_image_scale_mid2,
                  epoch,
              )
          if (
              valid_scale_mids_gt is not None
              and valid_scale_mids_pre is not None
              and len(valid_scale_mids_pre) > 0
          ):
              num_seq = config.VIS.SCALE_MID_NUM_SEQ
              vis_image_scale_mid1 = plot_scalemid_dist(
                  valid_scale_mids_pre, valid_scale_mids_gt.tolist()
              )
              vis_image_scale_mid1 = (
                  torch.from_numpy(vis_image_scale_mid1)
                  .type(torch.float32)
                  .permute(2, 0, 1)
                  / 255
              )
              vis_image_scale_mid2 = plot_scalemid_seq_dist(
                  valid_scale_mids_pre,
                  valid_scale_mids_gt.tolist(),
                  valid_seqnames,
                  num_seq=num_seq,
              )
              vis_image_scale_mid2 = (
                  torch.from_numpy(vis_image_scale_mid2)
                  .type(torch.float32)
                  .permute(2, 0, 1)
                  / 255
              )
              summary_writer.add_image(
                  "valid_scalemid_distribution/epoch",
                  vis_image_scale_mid1,
                  epoch,
              )
              summary_writer.add_image(
                  "valid_scalemid_seq_distribution/epoch",
                  vis_image_scale_mid2,
                  epoch,
              )
  
      summary_writer.close()
  
imports: import _init_paths
  import os, sys, datetime, shutil
  import os.path
  import random
  import h5py
  import pickle
  import numpy
  import argparse
  import subprocess
  import warnings
  import torch
  import torch.nn
  import torch.optim
  import torch.backends.cudnn
  from torch.utils.data import TensorDataset, DataLoader
  from torch.optim import lr_scheduler
  from lib.dataloader.mpiinf import MPIINFDataset
  from lib.dataloader.h36m import Human36MDataset
  from lib.dataloader.surreal import SurrealDataset
  from lib.models.model import get_pose_model, get_discriminator
  from lib.core.config import config, update_config, update_dir
  from lib.utils.misc import save_pickle, create_logger, load_pickle
  from lib.utils.utils import load_checkpoint, save_checkpoint, get_optimizer
  from lib.utils.vis import plot_scalemid_dist, plot_scalemid_seq_dist
  from tensorboardX import SummaryWriter
label_desc: Use image node_id 4 for calling a global function with example usage: main() without return types