text: from __future__ import print_function
  import torch
  import torch.utils.data as data
  import torchvision
  from torchvision import transforms
  import random
  import os
  import numpy as np
  from PIL import Image
  from torch.nn.functional import softmax
  class Base_Dataset(data.Dataset):
      def __init__(self, root, partition, target_ratio=0.0):
          super(Base_Dataset, self).__init__()
          # set dataset info
          self.root = root
          self.partition = partition
          self.target_ratio = target_ratio
          # self.target_ratio=0 no mixup
          mean_pix = [0.485, 0.456, 0.406]
          std_pix = [0.229, 0.224, 0.225]
          normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
  
          if self.partition == 'train':
              self.transformer = transforms.Compose([transforms.Resize(256),
                                                     transforms.RandomHorizontalFlip(),
                                                     transforms.RandomCrop(224),
                                                     transforms.ToTensor(),
                                                     normalize])
          else:
              self.transformer = transforms.Compose([transforms.Resize(256),
                                                     transforms.CenterCrop(224),
                                                     transforms.ToTensor(),
                                                     normalize])
  
      def __len__(self):
  
          if self.partition == 'train':
              if self.label_flag is None:
                  return int(min(self.alpha))
              else:
                  return int(self.num_labeled_target/10)
                  # return int((self.num_labeled_target) / (self.num_class - 1))
          elif self.partition == 'test':
              return int(len(self.target_image) / (self.num_class - 1))
          elif self.partition == 'tune':
              return int(len(self.target_image))
          elif self.partition == 'vis':
              return int(len(self.target_image))
  
  
      def __getitem__(self, item):
  
          image_data = []
          label_data = []
  
          # have a labeled part and a random sampled part
          random_real_label = []
          class_index_target = []
  
          # domain_label = []
          ST_split = [] # Mask of targets to be evaluated
          # select index for support class
  
          # Phase 1
          if self.label_flag is None:
              class_index_source = list(range(self.num_class - 1))
              random.shuffle(class_index_source)
              if self.partition == 'train':
                  # load labeled source samples by class
                  for classes in class_index_source:
                      image = Image.open(random.choice(self.source_image[classes])).convert('RGB')
                      if self.transformer is not None:
                          image = self.transformer(image)
                      image_data.append(image)
                      label_data.append(classes)
                      ST_split.append(0)
                  # randomly load source
  
                  for i in range(self.num_class - 1):
                      index = random.choice(list(range(len(self.source_image))))
                      # index = random.choice(list(range(len(self.label_flag))))
                      source_image = Image.open(self.all_source_images[index]).convert('RGB')
                      if self.transformer is not None:
                          source_image = self.transformer(source_image)
                      image_data.append(source_image)
                      label_data.append(self.all_source_labels[index])
                      random_real_label.append(self.all_source_labels[index])
                      # domain_label.append(0)
                      ST_split.append(0)
  
              elif self.partition == 'test':
                  # randomly load target samples 
                  for i in range(self.num_class - 1):
                      index = random.choice(list(range(len(self.target_image))))
                      # index = random.choice(list(range(len(self.label_flag))))
                      target_image = Image.open(self.target_image[index]).convert('RGB')
                      if self.transformer is not None:
                          target_image = self.transformer(target_image)
                      image_data.append(target_image)
                      label_data.append(0) # just to ignore
                      ST_split.append(1)
                  # load target by index
                  for i in range(self.num_class - 1):
                      target_image = Image.open(self.target_image[item * (self.num_class - 1) + i]).convert('RGB')
                      if self.transformer is not None:
                          target_image = self.transformer(target_image)
                      image_data.append(target_image)
                      label_data.append(self.num_class) # mark as target samples to evaluate
                      random_real_label.append(self.target_label[item * (self.num_class - 1) + i])
                      ST_split.append(1)
              
              elif self.partition == 'tune':
                  # randomly load target samples 
                  index = random.choice(list(range(len(self.target_image))))
                  # index = random.choice(list(range(len(self.label_flag))))
                  target_image = Image.open(self.target_image[index]).convert('RGB')
                  if self.transformer is not None:
                      target_image = self.transformer(target_image)
                  image_data.append(target_image)
                  label_data.append(0) # just to ignore
                  ST_split.append(1) # depre
  
                  # load target by index
                  target_image = Image.open(self.target_image[item]).convert('RGB')
                  if self.transformer is not None:
                      target_image = self.transformer(target_image)
                  image_data.append(target_image)
                  label_data.append(self.num_class) # mark as target samples to evaluate
                  random_real_label.append(self.target_label[item])
                  ST_split.append(1)
  
  
  
  
          # Phase 2
          else:
              # num_class_index_target = int(self.target_ratio * (self.num_class - 1))
              if self.partition == 'vis':
                  for classes in range(self.num_class - 1):
                      image = Image.open(random.choice(self.source_image[classes])).convert('RGB')
                      if self.transformer is not None:
                          image = self.transformer(image)
                      image_data.append(image)
                      label_data.append(classes)
                      ST_split.append(0)
                      
                  for classes in range(self.num_class - 1):    
                      index = random.choice(list(range(len(self.target_image))))
                      target_image = Image.open(self.target_image[index]).convert('RGB')
                      if self.transformer is not None:
                          target_image = self.transformer(target_image)
                      image_data.append(target_image)
                      label_data.append(self.target_label[index])
                      random_real_label.append(self.target_label[index])
                      ST_split.append(1)
          
                  
                  
              else:
  
                  if self.target_ratio > 0:
  
                      class_index_target = self.available_index
  
                      # class_index_target = random.sample(available_index, min(num_class_index_target, len(available_index)))
  
                  for classes in class_index_target:
                      # select support samples from source domain or target domain
                      image = Image.open(random.choice(self.target_image_list[classes])).convert('RGB')
  
                      if self.transformer is not None:
                          image = self.transformer(image)
                      image_data.append(image)
                      label_data.append(classes)
                      # domain_label.append(0)
                      ST_split.append(0)
                      # target_real_label.append(classes)
  
                  # add target samples
                  for i in range(self.num_class - 1):
  
                      if self.partition == 'train':
                          index = random.choice(list(range(len(self.label_flag))))
                          # index = random.choice(list(range(len(self.target_image))))
                          target_image = Image.open(self.target_image[index]).convert('RGB')
                          if self.transformer is not None:
                              target_image = self.transformer(target_image)
                          image_data.append(target_image)
                          label_data.append(self.label_flag[index])
                          random_real_label.append(self.target_label[index])
                          # domain_label.append(0)
                          ST_split.append(0)
  
                      elif self.partition == 'test':
                          target_image = Image.open(self.target_image[item * (self.num_class - 1) + i]).convert('RGB')
                          if self.transformer is not None:
                              target_image = self.transformer(target_image)
                          image_data.append(target_image)
                          label_data.append(self.num_class)
                          random_real_label.append(self.target_label[item * (self.num_class - 1) + i])
                          # domain_label.append(0)
                          ST_split.append(1)
  
          image_data = torch.stack(image_data)
          label_data = torch.tensor(label_data).type(torch.LongTensor)
          real_label_data = torch.tensor(random_real_label)
          # domain_label = torch.tensor(domain_label)
          ST_split = torch.tensor(ST_split)
          return image_data, label_data, real_label_data, ST_split
  
      def load_dataset(self):
          source_image_list = {key: [] for key in range(self.num_class - 1)}
          target_image_list = []
          target_label_list = []
          all_source_images = []
          all_source_labels = []
  
          with open(self.source_path) as f:
              for ind, line in enumerate(f.readlines()):
                  image_dir, label = line.split(' ')
                  label = label.strip()
                  if label == str(self.num_class-1):
                      continue
                  source_image_list[int(label)].append(image_dir)
                  # source_image_list.append(image_dir)
                  all_source_images.append(image_dir)
                  all_source_labels.append(int(label))
  
  
          with open(self.target_path) as f:
              for ind, line in enumerate(f.readlines()):
                  image_dir, label = line.split(' ')
                  label = label.strip()
                  # target_image_list[int(label)].append(image_dir)
                  target_image_list.append(image_dir)
                  target_label_list.append(int(label))
  
          return source_image_list, target_image_list, target_label_list, all_source_images, all_source_labels
  
  
  class Office_Dataset(Base_Dataset):
  
      def __init__(self, root, partition, label_flag=None, source='A', target='W', target_ratio=0.0, confidence_ratio=None):
          super(Office_Dataset, self).__init__(root, partition, target_ratio)
          # set dataset info
          src_name, tar_name = self.getFilePath(source, target)
          self.source_path = os.path.join(root, src_name)
          self.target_path = os.path.join(root, tar_name)
          self.class_name = ["back_pack", "bike", "bike_helmet", "bookcase", "bottle",
                             "calculator", "desk_chair", "desk_lamp", "desktop_computer", "file_cabinet", "unk"]
          self.num_class = len(self.class_name)
          self.source_image, self.target_image, self.target_label, self.all_source_images, self.all_source_labels = self.load_dataset()
          self.alpha = [len(self.source_image[key]) for key in self.source_image.keys()]
          self.label_flag = label_flag
  
          # source only
          # if self.label_flag is None:
              # self.label_flag = torch.ones(len(self.target_image)) * self.num_class
  
          if self.label_flag is not None:
              # if pseudo label comes
              self.target_image_list = {key: [] for key in range(self.num_class + 1)}
              for i in range(len(self.label_flag)):
                  self.target_image_list[self.label_flag[i].item()].append(self.target_image[i])
              self.available_index = [key for key in self.target_image_list.keys() if len(self.target_image_list[key]) > 0
                                     and key < self.num_class - 1]
              self.num_labeled_target = sum([len(self.target_image_list[key]) for key in self.target_image_list.keys()[:-2]])
              if confidence_ratio is not None:
                  self.class_ratio = 1 - torch.tensor(confidence_ratio[:-1])
              else:
                  self.class_ratio = torch.tensor([self.num_labeled_target - len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]]).float()
              self.class_ratio = softmax(self.class_ratio)
              print("class weighted reconsidered")
              print(self.class_ratio)
  
          if self.target_ratio > 0:
              self.alpha_value = [len(self.source_image[key]) + len(self.target_image_list[key]) for key in self.source_image.keys()]
          else:
              self.alpha_value = self.alpha
  
          self.alpha_value = np.array(self.alpha_value)
          self.alpha_value = (self.alpha_value.max() + 1 - self.alpha_value) / self.alpha_value.mean()
          self.alpha_value = torch.tensor(self.alpha_value).float().cuda()
  
      def getFilePath(self, source, target):
  
          if source == 'A':
              src_name = 'amazon_src_list.txt'
          elif source == 'W':
              src_name = 'webcam_src_list.txt'
          elif source == 'D':
              src_name = 'dslr_src_list.txt'
          else:
              print("Unknown Source Type, only supports A W D.")
  
          if target == 'A':
              tar_name = 'amazon_tar_list.txt'
          elif target == 'W':
              tar_name = 'webcam_tar_list.txt'
          elif target == 'D':
              tar_name = 'dslr_tar_list.txt'
          else:
              print("Unknown Target Type, only supports A W D.")
  
          return src_name, tar_name
  
  
  
  
  class Home_Dataset(Base_Dataset):
      def __init__(self, root, partition, label_flag=None, source='A', target='R', target_ratio=0.0, confidence_ratio=None):
          super(Home_Dataset, self).__init__(root, partition, target_ratio)
          src_name, tar_name = self.getFilePath(source, target)
          self.source_path = os.path.join(root, src_name)
          self.target_path = os.path.join(root, tar_name)
          self.class_name = ['Alarm_Clock', 'Backpack', 'Batteries', 'Bed', 'Bike', 'Bottle', 'Bucket', 'Calculator',
                             'Calendar', 'Candles', 'Chair', 'Clipboards', 'Computer', 'Couch', 'Curtains', 'Desk_Lamp',
                             'Drill', 'Eraser', 'Exit_Sign', 'Fan', 'File_Cabinet', 'Flipflops', 'Flowers', 'Folder',
                             'Fork', 'unk']
          self.num_class = len(self.class_name)
  
          self.source_image, self.target_image, self.target_label, self.all_source_images, self.all_source_labels = self.load_dataset()
          self.alpha = [len(self.source_image[key]) for key in self.source_image.keys()]
          self.label_flag = label_flag
  
          # create the unlabeled tag
  
              # self.label_flag = torch.ones(len(self.target_image)) * self.num_class
  
          if self.label_flag is not None:
              # if pseudo label comes
              self.target_image_list = {key: [] for key in range(self.num_class + 1)}
              for i in range(len(self.label_flag)):
                  self.target_image_list[self.label_flag[i].item()].append(self.target_image[i])
              self.available_index = [key for key in self.target_image_list.keys() if len(self.target_image_list[key]) > 0
                                 and key < self.num_class - 1]
              self.num_labeled_target = sum(
                  [len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]])
              if confidence_ratio is not None:
                  self.class_ratio = 1 - torch.tensor(confidence_ratio[:-1])
              else:
                  self.class_ratio = torch.tensor([self.num_labeled_target - len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]]).float()
              self.class_ratio = softmax(self.class_ratio)
              print("class weighted reconsidered")
              print(self.class_ratio)
  
  
          # if self.target_ratio > 0:
          #     self.alpha_value = [len(self.source_image[key]) + len(self.target_image_list[key]) for key in
          #                         self.source_image.keys()]
          # else:
          #     self.alpha_value = self.alpha
          #
          # self.alpha_value = np.array(self.alpha_value)
          # self.alpha_value = (self.alpha_value.max() + 1 - self.alpha_value) / self.alpha_value.mean()
          # self.alpha_value = torch.tensor(self.alpha_value).float().cuda()
  
      def getFilePath(self, source, target):
  
          if source == 'A':
              src_name = 'art_source.txt'
          elif source == 'C':
              src_name = 'clip_source.txt'
          elif source == 'P':
              src_name = 'product_source.txt'
          elif source == 'R':
              src_name = 'real_source.txt'
          else:
              print("Unknown Source Type, only supports A C P R.")
  
          if target == 'A':
              tar_name = 'art_tar.txt'
          elif target == 'C':
              tar_name = 'clip_tar.txt'
          elif target == 'P':
              tar_name = 'product_tar.txt'
          elif target == 'R':
              tar_name = 'real_tar.txt'
          else:
              print("Unknown Target Type, only supports A C P R.")
  
          return src_name, tar_name
  
  
  class Visda_Dataset(Base_Dataset):
      def __init__(self, root, partition, label_flag=None, target_ratio=0.0, confidence_ratio=None):
          super(Visda_Dataset, self).__init__(root, partition, target_ratio)
          # set dataset info
          self.source_path = os.path.join(root, 'source_list.txt')
          self.target_path = os.path.join(root, 'target_list.txt')
          self.class_name = ["bicycle", "bus", "car", "motorcycle", "train", "truck", 'unk']
          self.num_class = len(self.class_name)
          self.source_image, self.target_image, self.target_label, self.all_source_images, self.all_source_labels = self.load_dataset()
          self.alpha = [len(self.source_image[key]) for key in self.source_image.keys()]
          self.label_flag = label_flag
  
          # source-only stage
          # if self.label_flag is None:
          #     self.label_flag = torch.ones(len(self.target_image)) * self.num_class
  
          if self.label_flag is not None:
              # if pseudo label comes
              self.target_image_list = {key: [] for key in range(self.num_class + 1)}
              for i in range(len(self.label_flag)):
                  self.target_image_list[self.label_flag[i].item()].append(self.target_image[i])
  
              self.available_index = [key for key in self.target_image_list.keys() if len(self.target_image_list[key]) > 0
                                     and key < self.num_class - 1]
              self.num_labeled_target = sum(
                  [len(self.target_image_list[key]) for key in self.target_image_list.keys()][:-2])
              if confidence_ratio is not None:
                  self.class_ratio = 1 - torch.tensor(confidence_ratio[:-1])
              else:
                  self.class_ratio = torch.tensor([self.num_labeled_target - len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]]).float()
              self.class_ratio = softmax(self.class_ratio)
              print("class weighted reconsidered")
              print(self.class_ratio)
  
  
  class Visda18_Dataset(Base_Dataset):
      def __init__(self, root, partition, label_flag=None, target_ratio=0.0, confidence_ratio=None):
          super(Visda18_Dataset, self).__init__(root, partition, target_ratio)
          # set dataset info
          self.source_path = os.path.join(root, 'source_list_k.txt')
          self.target_path = os.path.join(root, 'target_list.txt')
          self.class_name = ["areoplane","bicycle", "bus", "car", "horse", "knife", "motorcycle", "person", "plant",
                             "skateboard", "train", "truck", 'unk']
          self.num_class = len(self.class_name)
          self.source_image, self.target_image, self.target_label, self.all_source_images, self.all_source_labels = self.load_dataset()
          self.alpha = [len(self.source_image[key]) for key in self.source_image.keys()]
          self.label_flag = label_flag
  
          # create the unlabeled tag
  
              # self.label_flag = torch.ones(len(self.target_image)) * self.num_class
  
  
          if self.label_flag is not None and self.label_flag is not 'vis':
              # if pseudo label comes
              self.target_image_list = {key: [] for key in range(self.num_class + 1)}
              for i in range(len(self.label_flag)):
                  self.target_image_list[self.label_flag[i].item()].append(self.target_image[i])
              self.available_index = [key for key in self.target_image_list.keys() if len(self.target_image_list[key]) > 0
                                 and key < self.num_class - 1]
              self.num_labeled_target = sum(
                  [len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]])
              if confidence_ratio is not None:
                  self.class_ratio = 1 - torch.tensor(confidence_ratio[:-1])
              else:
                  self.class_ratio = torch.tensor([self.num_labeled_target - len(self.target_image_list[key]) for key in list(self.target_image_list.keys())[:-2]]).float()
              self.class_ratio = softmax(self.class_ratio)
              print("class weighted reconsidered")
              print(self.class_ratio)
  
  