text: import logging
  
  import numpy as np
  from PIL import Image
  from torchvision import datasets
  from torchvision import transforms
  
  from .randaugment import RandAugmentMC
  
  logger = logging.getLogger(__name__)
  
  cifar10_mean = (0.4914, 0.4822, 0.4465)
  cifar10_std = (0.2471, 0.2435, 0.2616)
  cifar100_mean = (0.5071, 0.4867, 0.4408)
  cifar100_std = (0.2675, 0.2565, 0.2761)
  normal_mean = (0.5, 0.5, 0.5)
  normal_std = (0.5, 0.5, 0.5)
  
  
  def get_cifar10(root, num_labeled, num_expand_x, num_expand_u):
      transform_labeled = transforms.Compose([
          transforms.RandomHorizontalFlip(),
          transforms.RandomCrop(size=32,
                                padding=int(32*0.125),
                                padding_mode='reflect'),
          transforms.ToTensor(),
          transforms.Normalize(mean=cifar10_mean, std=cifar10_std)
      ])
      transform_val = transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize(mean=cifar10_mean, std=cifar10_std)
      ])
      base_dataset = datasets.CIFAR10(root, train=True, download=True)
  
      train_labeled_idxs, train_unlabeled_idxs = x_u_split(
          base_dataset.targets, num_labeled, num_expand_x, num_expand_u, num_classes=10)
  
      train_labeled_dataset = CIFAR10SSL(
          root, train_labeled_idxs, train=True,
          transform=transform_labeled)
  
      train_unlabeled_dataset = CIFAR10SSL(
          root, train_unlabeled_idxs, train=True,
          transform=TransformFix(mean=cifar10_mean, std=cifar10_std))
  
      test_dataset = datasets.CIFAR10(
          root, train=False, transform=transform_val, download=False)
      logger.info("Dataset: CIFAR10")
      logger.info(f"Labeled examples: {len(train_labeled_idxs)}"
                  f" Unlabeled examples: {len(train_unlabeled_idxs)}")
  
      return train_labeled_dataset, train_unlabeled_dataset, test_dataset
  
  
  def get_cifar100(root, num_labeled, num_expand_x, num_expand_u):
  
      transform_labeled = transforms.Compose([
          transforms.RandomHorizontalFlip(),
          transforms.RandomCrop(size=32,
                                padding=int(32*0.125),
                                padding_mode='reflect'),
          transforms.ToTensor(),
          transforms.Normalize(mean=cifar100_mean, std=cifar100_std)])
  
      transform_val = transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize(mean=cifar100_mean, std=cifar100_std)])
  
      base_dataset = datasets.CIFAR100(
          root, train=True, download=True)
  
      train_labeled_idxs, train_unlabeled_idxs = x_u_split(
          base_dataset.targets, num_labeled, num_classes=100)
  
      train_labeled_dataset = CIFAR100SSL(
          root, train_labeled_idxs, train=True,
          transform=transform_labeled)
  
      train_unlabeled_dataset = CIFAR100SSL(
          root, train_unlabeled_idxs, train=True,
          transform=TransformFix(mean=cifar100_mean, std=cifar100_std))
  
      test_dataset = datasets.CIFAR100(
          root, train=False, transform=transform_val, download=False)
  
      logger.info("Dataset: CIFAR100")
      logger.info(f"Labeled examples: {len(train_labeled_idxs)}"
                  f" Unlabeled examples: {len(train_unlabeled_idxs)}")
  
      return train_labeled_dataset, train_unlabeled_dataset, test_dataset
  
  
  def x_u_split(labels,
                num_labeled,
                num_expand_x,
                num_expand_u,
                num_classes):
      label_per_class = num_labeled // num_classes
      labels = np.array(labels)
      labeled_idx = []
      unlabeled_idx = []
      for i in range(num_classes):
          idx = np.where(labels == i)[0]
          np.random.shuffle(idx)
          labeled_idx.extend(idx[:label_per_class])
          unlabeled_idx.extend(idx[label_per_class:])
  
      exapand_labeled = num_expand_x // len(labeled_idx)
      exapand_unlabeled = num_expand_u // len(unlabeled_idx)
      labeled_idx = np.hstack(
          [labeled_idx for _ in range(exapand_labeled)])
      unlabeled_idx = np.hstack(
          [unlabeled_idx for _ in range(exapand_unlabeled)])
  
      if len(labeled_idx) < num_expand_x:
          diff = num_expand_x - len(labeled_idx)
          labeled_idx = np.hstack(
              (labeled_idx, np.random.choice(labeled_idx, diff)))
      else:
          assert len(labeled_idx) == num_expand_x
  
      if len(unlabeled_idx) < num_expand_u:
          diff = num_expand_u - len(unlabeled_idx)
          unlabeled_idx = np.hstack(
              (unlabeled_idx, np.random.choice(unlabeled_idx, diff)))
      else:
          assert len(unlabeled_idx) == num_expand_u
  
      return labeled_idx, unlabeled_idx
  
  
  class TransformFix(object):
      def __init__(self, mean, std):
          self.weak = transforms.Compose([
              transforms.RandomHorizontalFlip(),
              transforms.RandomCrop(size=32,
                                    padding=int(32*0.125),
                                    padding_mode='reflect')])
          self.strong = transforms.Compose([
              transforms.RandomHorizontalFlip(),
              transforms.RandomCrop(size=32,
                                    padding=int(32*0.125),
                                    padding_mode='reflect'),
              RandAugmentMC(n=2, m=10)])
          self.normalize = transforms.Compose([
              transforms.ToTensor(),
              transforms.Normalize(mean=mean, std=std)])
  
      def __call__(self, x):
          weak = self.weak(x)
          strong = self.strong(x)
          return self.normalize(weak), self.normalize(strong)
  
  
  class CIFAR10SSL(datasets.CIFAR10):
      def __init__(self, root, indexs, train=True,
                   transform=None, target_transform=None,
                   download=False):
          super().__init__(root, train=train,
                           transform=transform,
                           target_transform=target_transform,
                           download=download)
          if indexs is not None:
              self.data = self.data[indexs]
              self.targets = np.array(self.targets)[indexs]
  
      def __getitem__(self, index):
          img, target = self.data[index], self.targets[index]
          img = Image.fromarray(img)
  
          if self.transform is not None:
              img = self.transform(img)
  
          if self.target_transform is not None:
              target = self.target_transform(target)
  
          return img, target
  
  
  class CIFAR100SSL(datasets.CIFAR100):
      def __init__(self, root, indexs, train=True,
                   transform=None, target_transform=None,
                   download=False):
          super().__init__(root, train=train,
                           transform=transform,
                           target_transform=target_transform,
                           download=download)
          if indexs is not None:
              self.data = self.data[indexs]
              self.targets = np.array(self.targets)[indexs]
  
      def __getitem__(self, index):
          img, target = self.data[index], self.targets[index]
          img = Image.fromarray(img)
  
          if self.transform is not None:
              img = self.transform(img)
  
          if self.target_transform is not None:
              target = self.target_transform(target)
  
          return img, target
  