datasets/ImageNet_LT/select_ood_test_imagenet_1k.py

from genericpath import isdir
import os, argparse

parser = argparse.ArgumentParser(description='PyTorch CIFAR Example')
parser.add_argument('--data_root_path', '--drp', default='/ssd1/haotao/datasets', help='data root path')
args = parser.parse_args()

imagenet_train_dir = os.path.join(args.data_root_path, 'imagenet', 'train')
imagenet1k_wnid_list = os.listdir(imagenet_train_dir)
assert len(imagenet1k_wnid_list)==1000, "%d is not 1000" % len(imagenet1k_wnid_list)

with open("datasets/imagenet_extra_1k_wnid_list.txt", "r") as fp:
	imagenet_extra_1k_wnid_list = fp.read().splitlines()

imagenet10k_dir = os.path.join(args.data_root_path, 'imagenet10k')
imagenet10k_wnid_list = []
for item in os.listdir(imagenet10k_dir):
    if '.tar' in item:
        imagenet10k_wnid_list.append(item.split('.')[0])
print(len(imagenet10k_wnid_list))

imagenet_ood_test_1k_wnid_list = [] 
for wnid in imagenet10k_wnid_list:
    if wnid not in imagenet1k_wnid_list and wnid not in imagenet_extra_1k_wnid_list:
        imagenet_ood_test_1k_wnid_list.append(wnid)
        if len(imagenet_ood_test_1k_wnid_list) == 1000:
            break

fp = open('./datasets/imagenet_ood_test_1k_wnid_list.txt', 'w+')
for wnid in imagenet_ood_test_1k_wnid_list:
    fp.write(wnid)
    fp.write('\n')
    fp.flush()