import os

# set the root directory
root_dir = "/Checkpoint/user/data/tree_data/imagenet_images_256/train/entertainment center/"

name = "entertainment center"
attributes = []
for subdir, dirs, files in os.walk(root_dir):
    for file in files:
        # check if file is an image file
        if file.endswith('.jpg') or file.endswith('.png'):
            # get the full path to the image file
            image_path = os.path.join(subdir, file)
            # extract the attribute from the path
            attribute = subdir[subdir.find(name)+len(name)+1:]
            print(attribute)
            # map the category to its index
            # category_index = category_to_index[category]
            # append to the respective lists
            # self.images.append(image_path)
            # self.y.append(category_index)
            attributes.append(attribute)