code: def __init__(
      self, root_dir, local_rank, aug_modes="brightness=0.1+mask=0.1"
  ):
      super(MXFaceDataset, self).__init__()
      default_aug_probs = {
          "brightness": 0.2,
          "blur": 0.1,
          "mask": 0.1,
      }
  
      aug_mode_list = aug_modes.lower().split("+")
      aug_mode_map = {}
      for aug_mode_str in aug_mode_list:
          _aug = aug_mode_str.split("=")
          aug_key = _aug[0]
          if len(_aug) > 1:
              aug_prob = float(_aug[1])
          else:
              aug_prob = default_aug_probs[aug_key]
          aug_mode_map[aug_key] = aug_prob
  
      transform_list = []
      self.mask_aug = False
      self.mask_prob = 0.0
      key = "mask"
      if key in aug_mode_map:
          self.mask_aug = True
          self.mask_prob = aug_mode_map[key]
          transform_list.append(
              MaskAugmentation(
                  mask_names=[
                      "mask_white",
                      "mask_blue",
                      "mask_black",
                      "mask_green",
                  ],
                  mask_probs=[0.4, 0.4, 0.1, 0.1],
                  h_low=0.33,
                  h_high=0.4,
                  p=self.mask_prob,
              )
          )
      if local_rank == 0:
          print("data_transform_list:", transform_list)
          print("mask:", self.mask_aug, self.mask_prob)
      key = "brightness"
      if key in aug_mode_map:
          prob = aug_mode_map[key]
          transform_list.append(
              A.RandomBrightnessContrast(
                  brightness_limit=0.125, contrast_limit=0.05, p=prob
              )
          )
      key = "blur"
      if key in aug_mode_map:
          prob = aug_mode_map[key]
          transform_list.append(
              A.ImageCompression(
                  quality_lower=30, quality_upper=80, p=prob
              )
          )
          transform_list.append(A.MedianBlur(blur_limit=(1, 7), p=prob))
          transform_list.append(
              A.MotionBlur(blur_limit=(5, 12), p=prob)
          )
      transform_list += [
          A.HorizontalFlip(p=0.5),
          A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
          ToTensorV2(),
      ]
      # here, the input for A transform is rgb cv2 img
      self.transform = A.Compose(transform_list)
      self.root_dir = root_dir
      self.local_rank = local_rank
      path_imgrec = os.path.join(root_dir, "train.rec")
      path_imgidx = os.path.join(root_dir, "train.idx")
      self.imgrec = mx.recordio.MXIndexedRecordIO(
          path_imgidx, path_imgrec, "r"
      )
      s = self.imgrec.read_idx(0)
      header, _ = mx.recordio.unpack(s)
      # print(header)
      # print(len(self.imgrec.keys))
      if header.flag > 0:
          if len(header.label) == 2:
              self.imgidx = np.array(range(1, int(header.label[0])))
          else:
              self.imgidx = np.array(list(self.imgrec.keys))
      else:
          self.imgidx = np.array(list(self.imgrec.keys))
  
imports: import numbers
  import os
  import queue
  import threading
  import mxnet
  import numpy
  import torch
  from torch.utils.data import DataLoader, Dataset
  from torchvision import transforms
  import cv2
  import albumentations
  from albumentations.pytorch import ToTensorV2
  from insightface.app import MaskAugmentation
label_desc: Use image node_id 1 to create a new MXFaceDataset object from inherited base classes: Dataset with example: obj = MXFaceDataset(root_dir, local_rank, aug_modes)