text: # -*- encoding: utf-8 -*-
  '''
  @File    :   cub2011.py    
  @Contact :   guzhouweihu@163.com
  
  @Modify Time      @Author        @Version    @Desciption
  ------------      -----------    --------    -----------
  2020/10/10 20:33   guzhouweihu      1.0         None
  '''
  
  import os
  
  import pandas as pd
  from torchvision.datasets import VisionDataset
  from torchvision.datasets.folder import default_loader
  from torchvision.datasets.utils import download_file_from_google_drive, download_url
  
  
  class Cub2011(VisionDataset):
      """`CUB-200-2011 <http://www.vision.caltech.edu/visipedia/CUB-200-2011.html>`_ Dataset.
          Args:
              root (string): Root directory of the dataset.
              train (bool, optional): If True, creates dataset from training set, otherwise
                 creates from test set.
              transform (callable, optional): A function/transform that  takes in an PIL image
                 and returns a transformed version. E.g, ``transforms.RandomCrop``
              target_transform (callable, optional): A function/transform that takes in the
                 target and transforms it.
              download (bool, optional): If true, downloads the dataset from the internet and
                 puts it in root directory. If dataset is already downloaded, it is not
                 downloaded again.
      """
      base_folder = 'CUB_200_2011/images'
      # url = 'http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgz'
      url = 'https://drive.google.com/u/0/uc?export=download&confirm=VJfU&id=1hbzc_P1FuxMkcabkgn9ZKinBwW683j45'
      file_id = '1hbzc_P1FuxMkcabkgn9ZKinBwW683j45'
      filename = 'CUB_200_2011.tgz'
      tgz_md5 = '97eceeb196236b17998738112f37df78'
  
      def __init__(self, root, train=True, transform=None, target_transform=None, download=False):
          super(Cub2011, self).__init__(root, transform=transform, target_transform=target_transform)
  
          self.loader = default_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')
  
      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')
          self.data = data.merge(train_test_split, on='img_id')
  
          class_names = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'classes.txt'),
                                    sep=' ', names=['class_name'], usecols=[1])
          self.class_names = class_names['class_name'].to_list()
          if self.train:
              self.data = self.data[self.data.is_training_img == 1]
          else:
              self.data = self.data[self.data.is_training_img == 0]
  
      def _check_integrity(self):
          try:
              self._load_metadata()
          except Exception:
              return False
  
          for index, row in self.data.iterrows():
              filepath = os.path.join(self.root, self.base_folder, row.filepath)
              if not os.path.isfile(filepath):
                  print(filepath)
                  return False
          return True
  
      def _download(self):
          import tarfile
  
          if self._check_integrity():
              print('Files already downloaded and verified')
              return
  
          # download_file_from_google_drive(self.file_id, self.root, self.filename, self.tgz_md5)
          download_url(self.url, self.root, self.filename, md5=self.tgz_md5)
  
          with tarfile.open(os.path.join(self.root, self.filename), "r:gz") as tar:
              tar.extractall(path=self.root)
  
      def __len__(self):
          return len(self.data)
  
      def __getitem__(self, idx):
          sample = self.data.iloc[idx]
          path = os.path.join(self.root, self.base_folder, sample.filepath)
          target = sample.target - 1  # Targets start at 1 by default, so shift to 0
          img = self.loader(path)
  
          if self.transform is not None:
              img = self.transform(img)
          if self.target_transform is not None:
              target = self.target_transform(target)
          return img, target
  
  
  if __name__ == '__main__':
      train_dataset = Cub2011('./cub2011', train=True, download=False)
      test_dataset = Cub2011('./cub2011', train=False, download=False)
  