code: def __init__(
      self,
      cfg,
      is_train=True,
      norm_6dof=True,
      degrees_6dof=False,
      local_rank=0,
  ):
      super(MXFaceDataset, self).__init__()
  
      self.is_train = is_train
      self.data_root = Path(cfg.root_dir)
      self.input_size = cfg.input_size
      self.transform = get_aug_transform(cfg)
      self.local_rank = local_rank
      self.use_trainval = cfg.use_trainval
      self.use_eye = cfg.eyes is not None
      if is_train:
          # self.df = pd.read_csv(osp.join(cfg.cache_dir, 'train_list.csv'), dtype={'subject_id': str, 'facial_action': str, 'img_id': str})
          path_imgrec = os.path.join(cfg.cache_dir, "train.rec")
          path_imgidx = os.path.join(cfg.cache_dir, "train.idx")
          self.imgrec = mx.recordio.MXIndexedRecordIO(
              path_imgidx, path_imgrec, "r"
          )
          self.imgidx = list(self.imgrec.keys)
          self.imggroup = [0] * len(self.imgidx)
          self.size_train = len(self.imgidx)
          if self.use_trainval:
              assert not cfg.sampling_hard
              path_imgrec = os.path.join(cfg.cache_dir, "val.rec")
              path_imgidx = os.path.join(cfg.cache_dir, "val.idx")
              self.imgrec2 = mx.recordio.MXIndexedRecordIO(
                  path_imgidx, path_imgrec, "r"
              )
              imgidx2 = list(self.imgrec2.keys)
              self.imggroup += [1] * len(imgidx2)
              self.imgidx += imgidx2
      else:
          # self.df = pd.read_csv(osp.join(cfg.cache_dir, 'val_list.csv'), dtype={'subject_id': str, 'facial_action': str, 'img_id': str})
          path_imgrec = os.path.join(cfg.cache_dir, "val.rec")
          path_imgidx = os.path.join(cfg.cache_dir, "val.idx")
          self.imgrec = mx.recordio.MXIndexedRecordIO(
              path_imgidx, path_imgrec, "r"
          )
          self.imgidx = list(self.imgrec.keys)
          self.imggroup = [0] * len(self.imgidx)
      self.imgidx = np.array(self.imgidx)
      self.imggroup = np.array(self.imggroup)
      if cfg.sampling_hard and is_train:
          meta = np.load(os.path.join(cfg.cache_dir, "train.meta.npy"))
          assert meta.shape[0] == len(self.imgidx)
          new_imgidx = []
          for i in range(len(self.imgidx)):
              idx = self.imgidx[i]
              assert i == idx
              pose = np.abs(meta[i, :2])
              # repeat = np.sum(pose>=35)*3+1
              if np.max(pose) < 15:
                  repeat = 2
              else:
                  repeat = 1
              new_imgidx += [idx] * repeat
          if local_rank == 0:
              print("new-imgidx:", len(self.imgidx), len(new_imgidx))
          self.imgidx = np.array(new_imgidx)
      self.label_6dof_mean = [
          -0.018197,
          -0.017891,
          0.025348,
          -0.005368,
          0.001176,
          -0.532206,
      ]  # mean of pitch, yaw, roll, tx, ty, tz
      self.label_6dof_std = [
          0.314015,
          0.271809,
          0.081881,
          0.022173,
          0.048839,
          0.065444,
      ]  # std of pitch, yaw, roll, tx, ty, tz
      txt_path = self.data_root / "resources/projection_matrix.txt"
      self.M_proj = np.loadtxt(txt_path, dtype=np.float32)
      self.M1 = np.array(
          [
              [400.0, 0, 0, 0],
              [0, 400.0, 0, 0],
              [0, 0, 1, 0],
              [400.0, 400.0, 0, 1],
          ]
      )
      self.dst_pts = np.float32(
          [[0, 0], [0, cfg.input_size - 1], [cfg.input_size - 1, 0]]
      )
      self.norm_6dof = norm_6dof
      self.degrees_6dof = degrees_6dof
      self.task = cfg.task
      self.num_verts = cfg.num_verts
      self.loss_pip = cfg.loss_pip
      self.net_stride = 32
      if local_rank == 0:
          logging.info("data_transform_list:%s" % self.transform)
          logging.info("len:%d" % len(self.imgidx))
          logging.info("glen:%d" % len(self.imggroup))
      self.is_test_aug = False
  
      self.enable_flip = cfg.enable_flip
      self.flipindex = cfg.flipindex.copy()
      self.verts3d_central_index = cfg.verts3d_central_index
      self.eye_dataset = None
      self.use_eye = False
      if cfg.eyes is not None:
          # from eye_dataset import EyeDataset
          # self.eye_dataset = EyeDataset(cfg.eyes['root'], load_data=False)
          self.use_eye = True
  
imports: import numbers
  import os
  import os.path
  import pickle
  import queue
  import threading
  import logging
  import numbers
  import math
  import pandas
  from scipy.spatial.transform import Rotation
  import mxnet
  from pathlib import Path
  import numpy
  import torch
  from torch.utils.data import DataLoader, Dataset
  from skimage import transform
  import cv2
  import albumentations
  from albumentations.pytorch import ToTensorV2
  from augs import *
label_desc: Use image node_id 1 to create a new MXFaceDataset object from inherited base classes: Dataset with example: obj = MXFaceDataset(cfg, is_train, norm_6dof, degrees_6dof, local_rank)