text: # encoding: utf-8
  """
  @author:  Jinkai Zheng
  @contact: 1315673509@qq.com
  """
  
  import os.path as osp
  
  from .bases import ImageDataset
  from ..datasets import DATASET_REGISTRY
  
  
  @DATASET_REGISTRY.register()
  class VeRiWild(ImageDataset):
      """VeRi-Wild.
  
      Reference:
          Lou et al. A Large-Scale Dataset for Vehicle Re-Identification in the Wild. CVPR 2019.
  
      URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
  
      Train dataset statistics:
          - identities: 30671.
          - images: 277797.
      """
      dataset_dir = "VERI-Wild"
      dataset_name = "veriwild"
  
      def __init__(self, root='datasets', query_list='', gallery_list='', **kwargs):
          self.dataset_dir = osp.join(root, self.dataset_dir)
  
          self.image_dir = osp.join(self.dataset_dir, 'images')
          self.train_list = osp.join(
              self.dataset_dir, 'train_test_split/train_list.txt')
          self.vehicle_info = osp.join(
              self.dataset_dir, 'train_test_split/vehicle_info.txt')
          if query_list and gallery_list:
              self.query_list = query_list
              self.gallery_list = gallery_list
          else:
              self.query_list = osp.join(
                  self.dataset_dir, 'train_test_split/test_10000_query.txt')
              self.gallery_list = osp.join(
                  self.dataset_dir, 'train_test_split/test_10000.txt')
  
          required_files = [
              self.image_dir,
              self.train_list,
              self.query_list,
              self.gallery_list,
              self.vehicle_info,
          ]
          self.check_before_run(required_files)
  
          self.imgid2vid, self.imgid2camid, self.imgid2imgpath = self.process_vehicle(
              self.vehicle_info)
  
          train = self.process_dir(self.train_list)
          query = self.process_dir(self.query_list, is_train=False)
          gallery = self.process_dir(self.gallery_list, is_train=False)
  
          super(VeRiWild, self).__init__(train, query, gallery, **kwargs)
  
      def process_dir(self, img_list, is_train=True):
          img_list_lines = open(img_list, 'r').readlines()
  
          dataset = []
          for idx, line in enumerate(img_list_lines):
              line = line.strip()
              vid = int(line.split('/')[0])
              imgid = line.split('/')[1].split('.')[0]
              camid = int(self.imgid2camid[imgid])
              if is_train:
                  vid = f"{self.dataset_name}_{vid}"
                  camid = f"{self.dataset_name}_{camid}"
              dataset.append((self.imgid2imgpath[imgid], vid, camid))
  
          assert len(dataset) == len(img_list_lines)
          return dataset
  
      def process_vehicle(self, vehicle_info):
          imgid2vid = {}
          imgid2camid = {}
          imgid2imgpath = {}
          vehicle_info_lines = open(vehicle_info, 'r').readlines()
  
          for idx, line in enumerate(vehicle_info_lines):
              vid = line.strip().split('/')[0]
              imgid = line.strip().split(';')[0].split('/')[1]
              camid = line.strip().split(';')[1]
              img_path = osp.join(self.image_dir, vid, imgid + '.jpg')
              imgid2vid[imgid] = vid
              imgid2camid[imgid] = camid
              imgid2imgpath[imgid] = img_path
  
          assert len(imgid2vid) == len(vehicle_info_lines)
          return imgid2vid, imgid2camid, imgid2imgpath
  
  
  @DATASET_REGISTRY.register()
  class SmallVeRiWild(VeRiWild):
      """VeRi-Wild.
      Small test dataset statistics:
          - identities: 3000.
          - images: 41861.
      """
  
      def __init__(self, root='datasets', **kwargs):
          dataset_dir = osp.join(root, self.dataset_dir)
          self.query_list = osp.join(
              dataset_dir, 'train_test_split/test_3000_query.txt')
          self.gallery_list = osp.join(
              dataset_dir, 'train_test_split/test_3000.txt')
  
          super(SmallVeRiWild, self).__init__(
              root, self.query_list, self.gallery_list, **kwargs)
  
  
  @DATASET_REGISTRY.register()
  class MediumVeRiWild(VeRiWild):
      """VeRi-Wild.
      Medium test dataset statistics:
          - identities: 5000.
          - images: 69389.
      """
  
      def __init__(self, root='datasets', **kwargs):
          dataset_dir = osp.join(root, self.dataset_dir)
          self.query_list = osp.join(
              dataset_dir, 'train_test_split/test_5000_query.txt')
          self.gallery_list = osp.join(
              dataset_dir, 'train_test_split/test_5000.txt')
  
          super(MediumVeRiWild, self).__init__(
              root, self.query_list, self.gallery_list, **kwargs)
  
  
  @DATASET_REGISTRY.register()
  class LargeVeRiWild(VeRiWild):
      """VeRi-Wild.
      Large test dataset statistics:
          - identities: 10000.
          - images: 138517.
      """
  
      def __init__(self, root='datasets', **kwargs):
          dataset_dir = osp.join(root, self.dataset_dir)
          self.query_list = osp.join(
              dataset_dir, 'train_test_split/test_10000_query.txt')
          self.gallery_list = osp.join(
              dataset_dir, 'train_test_split/test_10000.txt')
  
          super(LargeVeRiWild, self).__init__(
              root, self.query_list, self.gallery_list, **kwargs)
  