def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_ind', type=int, required=True)
    parser.add_argument('--out_root', type=str, default='/scratch/shared/slow/xuji/iid_private')
    parser.add_argument('--IID_model_ind', type=int, required=True)
    parser.add_argument('--max_num_train', type=int, required=True)
    parser.add_argument('--test_code', default=False, action='store_true')
    parser.add_argument('--do_sift', default=False, action='store_true')
    config = parser.parse_args()
    config.out_dir = os.path.join(config.out_root, str(config.model_ind))
    if (not os.path.exists(config.out_dir)):
        os.makedirs(config.out_dir)
    archetype_config_path = os.path.join(config.out_root, str(config.IID_model_ind), 'config.pickle')
    print(('Loading archetype config from: %s' % archetype_config_path))
    with open(archetype_config_path, 'rb') as config_f:
        archetype_config = pickle.load(config_f)
    assert (config.IID_model_ind == archetype_config.model_ind)
    assert (archetype_config.mode == 'IID')
    sample_fn = _get_vectorised_colour_samples
    if config.do_sift:
        sample_fn = _get_vectorised_sift_samples
    archetype_config.include_rgb = True
    archetype_config.no_sobel = True
    if ('Coco' in archetype_config.dataset):
        assert (not archetype_config.using_IR)
        archetype_config.in_channels = 3
    elif (archetype_config.dataset == 'Potsdam'):
        assert archetype_config.using_IR
        archetype_config.in_channels = 4
    if ('Coco' in archetype_config.dataset):
        (dataloaders_head_A, mapping_assignment_dataloader, mapping_test_dataloader) = make_Coco_dataloaders(archetype_config)
    elif (archetype_config.dataset == 'Potsdam'):
        (dataloaders_head_A, mapping_assignment_dataloader, mapping_test_dataloader) = make_Potsdam_dataloaders(archetype_config)
    else:
        raise NotImplementedError
    dataloaders_head_B = dataloaders_head_A
    assert (archetype_config.num_dataloaders == 1)
    dataloader = dataloaders_head_B[0]
    samples = sample_fn(archetype_config, dataloader)
    print('got training samples')
    sys.stdout.flush()
    if config.test_code:
        print('testing code, taking 10000 samples only')
        samples = samples[:10000, :]
    else:
        num_samples_train = min(samples.shape[0], config.max_num_train)
        print(('taking %d samples' % num_samples_train))
        chosen_inds = np.random.choice(samples.shape[0], size=num_samples_train, replace=False)
        samples = samples[chosen_inds, :]
        print(samples.shape)
    sys.stdout.flush()
    kmeans = MiniBatchKMeans(n_clusters=archetype_config.gt_k, verbose=1).fit(samples)
    print('trained kmeans')
    sys.stdout.flush()
    (assign_samples, assign_labels) = sample_fn(archetype_config, mapping_assignment_dataloader)
    num_samples = assign_samples.shape[0]
    assign_preds = kmeans.predict(assign_samples)
    print('finished prediction for mapping assign/test data')
    sys.stdout.flush()
    assign_preds = torch.from_numpy(assign_preds).cuda()
    assign_labels = torch.from_numpy(assign_labels).cuda()
    if (archetype_config.eval_mode == 'hung'):
        match = _hungarian_match(assign_preds, assign_labels, preds_k=archetype_config.gt_k, targets_k=archetype_config.gt_k)
    elif (archetype_config.eval_mode == 'orig'):
        match = _original_match(assign_preds, assign_labels, preds_k=archetype_config.gt_k, targets_k=archetype_config.gt_k)
    elif (archetype_config.eval_mode == 'orig_soft'):
        assert False
    found = torch.zeros(archetype_config.gt_k)
    reordered_preds = torch.zeros(num_samples).to(torch.int32).cuda()
    for (pred_i, target_i) in match:
        reordered_preds[(assign_preds == pred_i)] = target_i
        found[pred_i] = 1
    assert (found.sum() == archetype_config.gt_k)
    acc = _acc(reordered_preds, assign_labels, archetype_config.gt_k)
    print(('got acc %f' % acc))
    config.epoch_acc = [acc]
    config.centroids = kmeans.cluster_centers_
    config.match = match
    with open(os.path.join(config.out_dir, 'config.pickle'), 'w') as outfile:
        pickle.dump(config, outfile)
    with open(os.path.join(config.out_dir, 'config.txt'), 'w') as text_file:
        text_file.write(('%s' % config))
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_ind', type=int, required=True)
    parser.add_argument('--out_root', type=str, default='/scratch/shared/slow/xuji/iid_private')
    parser.add_argument('--IID_model_ind', type=int, required=True)
    parser.add_argument('--max_num_train', type=int, required=True)
    parser.add_argument('--test_code', default=False, action='store_true')
    parser.add_argument('--do_sift', default=False, action='store_true')
    config = parser.parse_args()
    config.out_dir = os.path.join(config.out_root, str(config.model_ind))
    if (not os.path.exists(config.out_dir)):
        os.makedirs(config.out_dir)
    archetype_config_path = os.path.join(config.out_root, str(config.IID_model_ind), 'config.pickle')
    print(('Loading archetype config from: %s' % archetype_config_path))
    with open(archetype_config_path, 'rb') as config_f:
        archetype_config = pickle.load(config_f)
    assert (config.IID_model_ind == archetype_config.model_ind)
    assert (archetype_config.mode == 'IID')
    sample_fn = _get_vectorised_colour_samples
    if config.do_sift:
        sample_fn = _get_vectorised_sift_samples
    archetype_config.include_rgb = True
    archetype_config.no_sobel = True
    if ('Coco' in archetype_config.dataset):
        assert (not archetype_config.using_IR)
        archetype_config.in_channels = 3
    elif (archetype_config.dataset == 'Potsdam'):
        assert archetype_config.using_IR
        archetype_config.in_channels = 4
    if ('Coco' in archetype_config.dataset):
        (dataloaders_head_A, mapping_assignment_dataloader, mapping_test_dataloader) = make_Coco_dataloaders(archetype_config)
    elif (archetype_config.dataset == 'Potsdam'):
        (dataloaders_head_A, mapping_assignment_dataloader, mapping_test_dataloader) = make_Potsdam_dataloaders(archetype_config)
    else:
        raise NotImplementedError
    dataloaders_head_B = dataloaders_head_A
    assert (archetype_config.num_dataloaders == 1)
    dataloader = dataloaders_head_B[0]
    samples = sample_fn(archetype_config, dataloader)
    print('got training samples')
    sys.stdout.flush()
    if config.test_code:
        print('testing code, taking 10000 samples only')
        samples = samples[:10000, :]
    else:
        num_samples_train = min(samples.shape[0], config.max_num_train)
        print(('taking %d samples' % num_samples_train))
        chosen_inds = np.random.choice(samples.shape[0], size=num_samples_train, replace=False)
        samples = samples[chosen_inds, :]
        print(samples.shape)
    sys.stdout.flush()
    kmeans = MiniBatchKMeans(n_clusters=archetype_config.gt_k, verbose=1).fit(samples)
    print('trained kmeans')
    sys.stdout.flush()
    (assign_samples, assign_labels) = sample_fn(archetype_config, mapping_assignment_dataloader)
    num_samples = assign_samples.shape[0]
    assign_preds = kmeans.predict(assign_samples)
    print('finished prediction for mapping assign/test data')
    sys.stdout.flush()
    assign_preds = torch.from_numpy(assign_preds).cuda()
    assign_labels = torch.from_numpy(assign_labels).cuda()
    if (archetype_config.eval_mode == 'hung'):
        match = _hungarian_match(assign_preds, assign_labels, preds_k=archetype_config.gt_k, targets_k=archetype_config.gt_k)
    elif (archetype_config.eval_mode == 'orig'):
        match = _original_match(assign_preds, assign_labels, preds_k=archetype_config.gt_k, targets_k=archetype_config.gt_k)
    elif (archetype_config.eval_mode == 'orig_soft'):
        assert False
    found = torch.zeros(archetype_config.gt_k)
    reordered_preds = torch.zeros(num_samples).to(torch.int32).cuda()
    for (pred_i, target_i) in match:
        reordered_preds[(assign_preds == pred_i)] = target_i
        found[pred_i] = 1
    assert (found.sum() == archetype_config.gt_k)
    acc = _acc(reordered_preds, assign_labels, archetype_config.gt_k)
    print(('got acc %f' % acc))
    config.epoch_acc = [acc]
    config.centroids = kmeans.cluster_centers_
    config.match = match
    with open(os.path.join(config.out_dir, 'config.pickle'), 'w') as outfile:
        pickle.dump(config, outfile)
    with open(os.path.join(config.out_dir, 'config.txt'), 'w') as text_file:
        text_file.write(('%s' % config))