code: def __init__(self, root_dir, is_train):
      super(GazeDataset, self).__init__()
  
      # self.local_rank = local_rank
      self.is_train = is_train
      self.input_size = 160
      # self.num_kps = 68
      transform_list = []
      if is_train:
          transform_list += [
              A.ColorJitter(brightness=0.3, contrast=0.3, p=0.5),
              A.ToGray(p=0.1),
              A.ISONoise(p=0.1),
              A.MedianBlur(blur_limit=(1, 7), p=0.1),
              A.GaussianBlur(blur_limit=(1, 7), p=0.1),
              A.MotionBlur(blur_limit=(5, 13), p=0.1),
              A.ImageCompression(
                  quality_lower=10, quality_upper=90, p=0.05
              ),
              A.ShiftScaleRotate(
                  shift_limit=0.1,
                  scale_limit=0.2,
                  rotate_limit=30,
                  interpolation=cv2.INTER_LINEAR,
                  border_mode=cv2.BORDER_CONSTANT,
                  value=0,
                  mask_value=0,
                  p=0.6,
              ),
              # A.HorizontalFlip(p=0.5),
              RectangleBorderAugmentation(
                  limit=0.2, fill_value=0, p=0.1
              ),
          ]
      transform_list += [
          A.geometric.resize.Resize(
              self.input_size,
              self.input_size,
              interpolation=cv2.INTER_LINEAR,
              always_apply=True,
          ),
          A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
          ToTensorV2(),
      ]
      self.transform = A.ReplayCompose(
          transform_list,
          keypoint_params=A.KeypointParams(
              format="xy", remove_invisible=False
          ),
      )
      self.root_dir = root_dir
      if is_train:
          path_imgrec = os.path.join(root_dir, "train.rec")
          path_imgidx = os.path.join(root_dir, "train.idx")
      else:
          path_imgrec = os.path.join(root_dir, "val.rec")
          path_imgidx = os.path.join(root_dir, "val.idx")
      self.imgrec = mx.recordio.MXIndexedRecordIO(
          path_imgidx, path_imgrec, "r"
      )
      self.imgidx = np.array(list(self.imgrec.keys))
      logging.info("len:%d" % len(self.imgidx))
      print("!!!len:%d" % len(self.imgidx))
      # self.num_face = 1103
      self.num_eye = 481
  
imports: import os
  import os.path
  import queue
  import mxnet
  import pickle
  import threading
  import logging
  import numpy
  import insightface
  from insightface.utils import face_align
  import torch
  from torch.utils.data import DataLoader, Dataset
  from torchvision import transforms
  import cv2
  import albumentations
  from albumentations.pytorch import ToTensorV2
  from .augs import RectangleBorderAugmentation
label_desc: Use image node_id 1 to create a new GazeDataset object from inherited base classes: Dataset with example: obj = GazeDataset(root_dir, is_train)