Class CustomCub2011

data/cub.py:14–133

import torch
import torch.nn.functional as F
class CustomCub2011(Dataset):
    base_folder = 'CUB_200_2011/images'
    url = 'http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgz'
    filename = 'CUB_200_2011.tgz'
    tgz_md5 = '97eceeb196236b17998738112f37df78'

    def __init__(self, root, train=True, transform=None, target_transform=None, loader=default_loader, download=True):

        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform

        self.loader = loader
        self.train = train


        if download:
            self._download()

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')

        self.uq_idxs = np.array(range(len(self)))

    def _load_metadata(self):
        images = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'images.txt'), sep=' ',
                             names=['img_id', 'filepath'])
        image_class_labels = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'image_class_labels.txt'),
                                         sep=' ', names=['img_id', 'target'])
        train_test_split = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'train_test_split.txt'),
                                       sep=' ', names=['img_id', 'is_training_img'])

        data = images.merge(image_class_labels, on='img_id')
        data = data.merge(train_test_split, on='img_id')

        name_attribute = pd.read_csv(os.path.join(self.root, 'attributes.txt'), sep=' ',
                                     names=['attribute_id', 'attribute_name'])
        from collections import defaultdict
        dict_attribute = defaultdict(list)
        for _i, _name in zip(name_attribute['attribute_id'], name_attribute['attribute_name']):
            dict_attribute[_name.split('::')[0]].append(_name.split('::')[1])

        names_ = list(dict_attribute.keys())
        processed_attribute_file = os.path.join(self.root,'CUB_200_2011','processed_attributes.txt')
        A_all = pd.read_csv(processed_attribute_file, sep=' ', names=names_)
        A_all.insert(0, 'img_id', list(range(1, len(A_all)+1)))
        self.data = data.merge(A_all, on='img_id')
        self.dict_attribute = dict_attribute
        class_attributes_file = os.path.join(self.root,'CUB_200_2011','attributes',
                                             'class_attribute_labels_continuous.txt')
        C_A = np.zeros((200, 312))
        class_attr_rf = open(class_attributes_file, 'r')
        i = 0
        for line in class_attr_rf.readlines():
            strs = line.strip().split(' ')
            for j in range(len(strs)):
                C_A[i][j] = 0.0 if strs[j] == '0.0' else float(strs[j]) * 0.01