text: #
  # BSD 3-Clause License
  #
  # Copyright (c) 2017 xxxx
  # All rights reserved.
  # Copyright 2021 Huawei Technologies Co., Ltd
  #
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  #
  # * Redistributions of source code must retain the above copyright notice, this
  #   list of conditions and the following disclaimer.
  #
  # * Redistributions in binary form must reproduce the above copyright notice,
  #   this list of conditions and the following disclaimer in the documentation
  #   and/or other materials provided with the distribution.
  #
  # * Neither the name of the copyright holder nor the names of its
  #   contributors may be used to endorse or promote products derived from
  #   this software without specific prior written permission.
  #
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  # ============================================================================
  #from .vision import VisionDataset
  from PIL import Image
  import os
  import os.path
  import six
  import string
  import sys
  from collections import Iterable
  
  if sys.version_info[0] == 2:
      import cPickle as pickle
  else:
      import pickle
  
  from .utils import verify_str_arg, iterable_to_str
  
  
  class LSUNClass(VisionDataset):
      def __init__(self, root, transform=None, target_transform=None):
          import lmdb
          super(LSUNClass, self).__init__(root, transform=transform,
                                          target_transform=target_transform)
  
          self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False,
                               readahead=False, meminit=False)
          with self.env.begin(write=False) as txn:
              self.length = txn.stat()['entries']
          cache_file = '_cache_' + ''.join(c for c in root if c in string.ascii_letters)
          if os.path.isfile(cache_file):
              self.keys = pickle.load(open(cache_file, "rb"))
          else:
              with self.env.begin(write=False) as txn:
                  self.keys = [key for key, _ in txn.cursor()]
              pickle.dump(self.keys, open(cache_file, "wb"))
  
      def __getitem__(self, index):
          img, target = None, None
          env = self.env
          with env.begin(write=False) as txn:
              imgbuf = txn.get(self.keys[index])
  
          buf = six.BytesIO()
          buf.write(imgbuf)
          buf.seek(0)
          img = Image.open(buf).convert('RGB')
  
          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
  
      def __len__(self):
          return self.length
  
  
  class LSUN(VisionDataset):
      """
      `LSUN <http://lsun.cs.princeton.edu>`_ dataset.
  
      Args:
          root (string): Root directory for the database files.
          classes (string or list): One of {'train', 'val', 'test'} or a list of
              categories to load. e,g. ['bedroom_train', 'church_train'].
          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.
      """
  
      def __init__(self, root, classes='train', transform=None, target_transform=None):
          super(LSUN, self).__init__(root, transform=transform,
                                     target_transform=target_transform)
          self.classes = self._verify_classes(classes)
  
          # for each class, create an LSUNClassDataset
          self.dbs = []
          for c in self.classes:
              self.dbs.append(LSUNClass(
                  root=root + '/' + c + '_lmdb',
                  transform=transform))
  
          self.indices = []
          count = 0
          for db in self.dbs:
              count += len(db)
              self.indices.append(count)
  
          self.length = count
  
      def _verify_classes(self, classes):
          categories = ['bedroom', 'bridge', 'church_outdoor', 'classroom',
                        'conference_room', 'dining_room', 'kitchen',
                        'living_room', 'restaurant', 'tower']
          dset_opts = ['train', 'val', 'test']
  
          try:
              verify_str_arg(classes, "classes", dset_opts)
              if classes == 'test':
                  classes = [classes]
              else:
                  classes = [c + '_' + classes for c in categories]
          except ValueError:
              if not isinstance(classes, Iterable):
                  msg = ("Expected type str or Iterable for argument classes, "
                         "but got type {}.")
                  raise ValueError(msg.format(type(classes)))
  
              classes = list(classes)
              msg_fmtstr = ("Expected type str for elements in argument classes, "
                            "but got type {}.")
              for c in classes:
                  verify_str_arg(c, custom_msg=msg_fmtstr.format(type(c)))
                  c_short = c.split('_')
                  category, dset_opt = '_'.join(c_short[:-1]), c_short[-1]
  
                  msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}."
                  msg = msg_fmtstr.format(category, "LSUN class",
                                          iterable_to_str(categories))
                  verify_str_arg(category, valid_values=categories, custom_msg=msg)
  
                  msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts))
                  verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg)
  
          return classes
  
      def __getitem__(self, index):
          """
          Args:
              index (int): Index
  
          Returns:
              tuple: Tuple (image, target) where target is the index of the target category.
          """
          target = 0
          sub = 0
          for ind in self.indices:
              if index < ind:
                  break
              target += 1
              sub = ind
  
          db = self.dbs[target]
          index = index - sub
  
          if self.target_transform is not None:
              target = self.target_transform(target)
  
          img, _ = db[index]
          return img, target
  
      def __len__(self):
          return self.length
  
      def extra_repr(self):
          return "Classes: {classes}".format(**self.__dict__)
  