def __init__(self, root='', train=True, meta=True, num_meta=1000, corruption_prob=0, corruption_type='unif', transform=None, target_transform=None, download=False, seed=1):
    self.count = 0
    self.root = root
    self.transform = transform
    self.target_transform = target_transform
    self.train = train
    self.meta = meta
    self.corruption_prob = corruption_prob
    self.num_meta = num_meta
    if download:
        self.download()
    if (not self._check_integrity()):
        raise RuntimeError(('Dataset not found or corrupted.' + ' You can use download=True to download it'))
    if self.train:
        self.train_data = []
        self.train_labels = []
        self.train_coarse_labels = []
        self.train_labels_true = []
        self.soft_labels = []
        for fentry in self.train_list:
            f = fentry[0]
            file = os.path.join(root, self.base_folder, f)
            fo = open(file, 'rb')
            if (sys.version_info[0] == 2):
                entry = pickle.load(fo)
            else:
                entry = pickle.load(fo, encoding='latin1')
            self.train_data.append(entry['data'])
            if ('labels' in entry):
                self.train_labels += entry['labels']
                self.train_labels_true += entry['labels']
                img_num_list = ([int((self.num_meta / 10))] * 10)
                num_classes = 10
            else:
                self.train_labels += entry['fine_labels']
                self.train_labels_true += entry['fine_labels']
                self.train_coarse_labels += entry['coarse_labels']
                img_num_list = ([int((self.num_meta / 100))] * 100)
                num_classes = 100
            fo.close()
        self.train_data = np.concatenate(self.train_data)
        self.train_data = self.train_data.reshape((50000, 3, 32, 32))
        self.train_data = self.train_data.transpose((0, 2, 3, 1))
        data_list_val = {}
        for j in range(num_classes):
            data_list_val[j] = [i for (i, label) in enumerate(self.train_labels) if (label == j)]
        idx_to_meta = []
        idx_to_train = []
        print(img_num_list)
        for (cls_idx, img_id_list) in data_list_val.items():
            np.random.shuffle(img_id_list)
            img_num = img_num_list[int(cls_idx)]
            idx_to_meta.extend(img_id_list[:img_num])
            idx_to_train.extend(img_id_list[img_num:])
        if (meta is True):
            self.train_data = self.train_data[idx_to_meta]
            self.train_labels = list(np.array(self.train_labels)[idx_to_meta])
        else:
            self.train_data = self.train_data[idx_to_train]
            self.train_labels = list(np.array(self.train_labels)[idx_to_train])
            self.train_labels_true = list(np.array(self.train_labels_true)[idx_to_train])
            self.soft_labels = list(np.zeros((len(self.train_data), num_classes), dtype=np.float32))
            self.prediction = np.zeros((len(self.train_data), 10, num_classes), dtype=np.float32)
            clean_labels = self.train_labels
            np.save('clean_labels.npy', clean_labels)
            if (corruption_type == 'unif'):
                C = uniform_mix_C(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'flip'):
                C = flip_labels_C(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'flip2'):
                C = flip_labels_C_two(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'hierarchical'):
                assert (num_classes == 100), 'You must use CIFAR-100 with the hierarchical corruption.'
                coarse_fine = []
                for i in range(20):
                    coarse_fine.append(set())
                for i in range(len(self.train_labels)):
                    coarse_fine[self.train_coarse_labels[i]].add(self.train_labels[i])
                for i in range(20):
                    coarse_fine[i] = list(coarse_fine[i])
                C = (np.eye(num_classes) * (1 - corruption_prob))
                for i in range(20):
                    tmp = np.copy(coarse_fine[i])
                    for j in range(len(tmp)):
                        tmp2 = np.delete(np.copy(tmp), j)
                        C[(tmp[j], tmp2)] += ((corruption_prob * 1) / len(tmp2))
                self.C = C
                print(C)
            elif (corruption_type == 'clabels'):
                net = wrn.WideResNet(40, num_classes, 2, dropRate=0.3).cuda()
                model_name = './cifar{}_labeler'.format(num_classes)
                net.load_state_dict(torch.load(model_name))
                net.eval()
            else:
                assert False, "Invalid corruption type '{}' given. Must be in {'unif', 'flip', 'hierarchical'}".format(corruption_type)
            np.random.seed(seed)
            if (corruption_type == 'clabels'):
                mean = [(x / 255) for x in [125.3, 123.0, 113.9]]
                std = [(x / 255) for x in [63.0, 62.1, 66.7]]
                test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])
                sampling_probs = []
                print('Starting labeling')
                for i in range(((len(self.train_labels) // 64) + 1)):
                    current = self.train_data[(i * 64):((i + 1) * 64)]
                    current = [Image.fromarray(current[i]) for i in range(len(current))]
                    current = torch.cat([test_transform(current[i]).unsqueeze(0) for i in range(len(current))], dim=0)
                    data = V(current).cuda()
                    logits = net(data)
                    smax = F.softmax((logits / 5))
                    sampling_probs.append(smax.data.cpu().numpy())
                sampling_probs = np.concatenate(sampling_probs, 0)
                print('Finished labeling 1')
                new_labeling_correct = 0
                argmax_labeling_correct = 0
                for i in range(len(self.train_labels)):
                    old_label = self.train_labels[i]
                    new_label = np.random.choice(num_classes, p=sampling_probs[i])
                    self.train_labels[i] = new_label
                    if (old_label == new_label):
                        new_labeling_correct += 1
                    if (old_label == np.argmax(sampling_probs[i])):
                        argmax_labeling_correct += 1
                print('Finished labeling 2')
                print('New labeling accuracy:', (new_labeling_correct / len(self.train_labels)))
                print('Argmax labeling accuracy:', (argmax_labeling_correct / len(self.train_labels)))
            else:
                for i in range(len(self.train_labels)):
                    self.train_labels_true[i] = self.train_labels[i]
                for i in range(len(self.train_labels)):
                    self.train_labels[i] = np.random.choice(num_classes, p=C[self.train_labels[i]])
                print('train', len(self.train_labels))
                print('type', type(self.train_labels))
                self.corruption_matrix = C
        noise_labels = self.train_labels
        np.save('noise_labels.npy', noise_labels)
    else:
        f = self.test_list[0][0]
        file = os.path.join(root, self.base_folder, f)
        fo = open(file, 'rb')
        if (sys.version_info[0] == 2):
            entry = pickle.load(fo)
        else:
            entry = pickle.load(fo, encoding='latin1')
        self.test_data = entry['data']
        if ('labels' in entry):
            self.test_labels = entry['labels']
        else:
            self.test_labels = entry['fine_labels']
        fo.close()
        self.test_data = self.test_data.reshape((10000, 3, 32, 32))
        self.test_data = self.test_data.transpose((0, 2, 3, 1))
def __init__(self, root=, train=True, meta=True, num_meta=1000, corruption_prob=0, corruption_type='unif', transform=None, target_transform=None, download=False, seed=1):
    self.count = 0
    self.root = root
    self.transform = transform
    self.target_transform = target_transform
    self.train = train
    self.meta = meta
    self.corruption_prob = corruption_prob
    self.num_meta = num_meta
    if download:
        self.download()
    if (not self._check_integrity()):
        raise RuntimeError(('Dataset not found or corrupted.' + ' You can use download=True to download it'))
    if self.train:
        self.train_data = []
        self.train_labels = []
        self.train_coarse_labels = []
        self.train_labels_true = []
        self.soft_labels = []
        for fentry in self.train_list:
            f = fentry[0]
            file = os.path.join(root, self.base_folder, f)
            fo = open(file, 'rb')
            if (sys.version_info[0] == 2):
                entry = pickle.load(fo)
            else:
                entry = pickle.load(fo, encoding='latin1')
            self.train_data.append(entry['data'])
            if ('labels' in entry):
                self.train_labels += entry['labels']
                self.train_labels_true += entry['labels']
                img_num_list = ([int((self.num_meta / 10))] * 10)
                num_classes = 10
            else:
                self.train_labels += entry['fine_labels']
                self.train_labels_true += entry['fine_labels']
                self.train_coarse_labels += entry['coarse_labels']
                img_num_list = ([int((self.num_meta / 100))] * 100)
                num_classes = 100
            fo.close()
        self.train_data = np.concatenate(self.train_data)
        self.train_data = self.train_data.reshape((50000, 3, 32, 32))
        self.train_data = self.train_data.transpose((0, 2, 3, 1))
        data_list_val = {}
        for j in range(num_classes):
            data_list_val[j] = [i for (i, label) in enumerate(self.train_labels) if (label == j)]
        idx_to_meta = []
        idx_to_train = []
        print(img_num_list)
        for (cls_idx, img_id_list) in data_list_val.items():
            np.random.shuffle(img_id_list)
            img_num = img_num_list[int(cls_idx)]
            idx_to_meta.extend(img_id_list[:img_num])
            idx_to_train.extend(img_id_list[img_num:])
        if (meta is True):
            self.train_data = self.train_data[idx_to_meta]
            self.train_labels = list(np.array(self.train_labels)[idx_to_meta])
        else:
            self.train_data = self.train_data[idx_to_train]
            self.train_labels = list(np.array(self.train_labels)[idx_to_train])
            self.train_labels_true = list(np.array(self.train_labels_true)[idx_to_train])
            self.soft_labels = list(np.zeros((len(self.train_data), num_classes), dtype=np.float32))
            self.prediction = np.zeros((len(self.train_data), 10, num_classes), dtype=np.float32)
            clean_labels = self.train_labels
            np.save('clean_labels.npy', clean_labels)
            if (corruption_type == 'unif'):
                C = uniform_mix_C(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'flip'):
                C = flip_labels_C(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'flip2'):
                C = flip_labels_C_two(self.corruption_prob, num_classes)
                print(C)
                self.C = C
            elif (corruption_type == 'hierarchical'):
                assert (num_classes == 100), 'You must use CIFAR-100 with the hierarchical corruption.'
                coarse_fine = []
                for i in range(20):
                    coarse_fine.append(set())
                for i in range(len(self.train_labels)):
                    coarse_fine[self.train_coarse_labels[i]].add(self.train_labels[i])
                for i in range(20):
                    coarse_fine[i] = list(coarse_fine[i])
                C = (np.eye(num_classes) * (1 - corruption_prob))
                for i in range(20):
                    tmp = np.copy(coarse_fine[i])
                    for j in range(len(tmp)):
                        tmp2 = np.delete(np.copy(tmp), j)
                        C[(tmp[j], tmp2)] += ((corruption_prob * 1) / len(tmp2))
                self.C = C
                print(C)
            elif (corruption_type == 'clabels'):
                net = wrn.WideResNet(40, num_classes, 2, dropRate=0.3).cuda()
                model_name = './cifar{}_labeler'.format(num_classes)
                net.load_state_dict(torch.load(model_name))
                net.eval()
            else:
                assert False, "Invalid corruption type '{}' given. Must be in {'unif', 'flip', 'hierarchical'}".format(corruption_type)
            np.random.seed(seed)
            if (corruption_type == 'clabels'):
                mean = [(x / 255) for x in [125.3, 123.0, 113.9]]
                std = [(x / 255) for x in [63.0, 62.1, 66.7]]
                test_transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])
                sampling_probs = []
                print('Starting labeling')
                for i in range(((len(self.train_labels) // 64) + 1)):
                    current = self.train_data[(i * 64):((i + 1) * 64)]
                    current = [Image.fromarray(current[i]) for i in range(len(current))]
                    current = torch.cat([test_transform(current[i]).unsqueeze(0) for i in range(len(current))], dim=0)
                    data = V(current).cuda()
                    logits = net(data)
                    smax = F.softmax((logits / 5))
                    sampling_probs.append(smax.data.cpu().numpy())
                sampling_probs = np.concatenate(sampling_probs, 0)
                print('Finished labeling 1')
                new_labeling_correct = 0
                argmax_labeling_correct = 0
                for i in range(len(self.train_labels)):
                    old_label = self.train_labels[i]
                    new_label = np.random.choice(num_classes, p=sampling_probs[i])
                    self.train_labels[i] = new_label
                    if (old_label == new_label):
                        new_labeling_correct += 1
                    if (old_label == np.argmax(sampling_probs[i])):
                        argmax_labeling_correct += 1
                print('Finished labeling 2')
                print('New labeling accuracy:', (new_labeling_correct / len(self.train_labels)))
                print('Argmax labeling accuracy:', (argmax_labeling_correct / len(self.train_labels)))
            else:
                for i in range(len(self.train_labels)):
                    self.train_labels_true[i] = self.train_labels[i]
                for i in range(len(self.train_labels)):
                    self.train_labels[i] = np.random.choice(num_classes, p=C[self.train_labels[i]])
                print('train', len(self.train_labels))
                print('type', type(self.train_labels))
                self.corruption_matrix = C
        noise_labels = self.train_labels
        np.save('noise_labels.npy', noise_labels)
    else:
        f = self.test_list[0][0]
        file = os.path.join(root, self.base_folder, f)
        fo = open(file, 'rb')
        if (sys.version_info[0] == 2):
            entry = pickle.load(fo)
        else:
            entry = pickle.load(fo, encoding='latin1')
        self.test_data = entry['data']
        if ('labels' in entry):
            self.test_labels = entry['labels']
        else:
            self.test_labels = entry['fine_labels']
        fo.close()
        self.test_data = self.test_data.reshape((10000, 3, 32, 32))
        self.test_data = self.test_data.transpose((0, 2, 3, 1))