veriwild.py

Go to the documentation of this file.

# 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_dirdataset_dir = osp.join(root, self.dataset_dirdataset_dir)
 
        self.image_dir = osp.join(self.dataset_dirdataset_dir, 'images')
        self.train_list = osp.join(self.dataset_dirdataset_dir, 'train_test_split/train_list.txt')
        self.vehicle_info = osp.join(self.dataset_dirdataset_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_dirdataset_dir, 'train_test_split/test_10000_query.txt')
            self.gallery_list = osp.join(self.dataset_dirdataset_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]
            if is_train:
                vid = self.dataset_name + "_" + str(vid)
            dataset.append((self.imgid2imgpath[imgid], vid, int(self.imgid2camid[imgid])))
 
        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[1:]):
            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) - 1
        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_dirdataset_dir)
        self.query_listquery_list = osp.join(dataset_dir, 'train_test_split/test_3000_query.txt')
        self.gallery_listgallery_list = osp.join(dataset_dir, 'train_test_split/test_3000.txt')
 
        super(SmallVeRiWild, self).__init__(root, self.query_listquery_list, self.gallery_listgallery_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_dirdataset_dir)
        self.query_listquery_list = osp.join(dataset_dir, 'train_test_split/test_5000_query.txt')
        self.gallery_listgallery_list = osp.join(dataset_dir, 'train_test_split/test_5000.txt')
 
        super(MediumVeRiWild, self).__init__(root, self.query_listquery_list, self.gallery_listgallery_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_dirdataset_dir)
        self.query_listquery_list = osp.join(dataset_dir, 'train_test_split/test_10000_query.txt')
        self.gallery_listgallery_list = osp.join(dataset_dir, 'train_test_split/test_10000.txt')
 
        super(LargeVeRiWild, self).__init__(root, self.query_listquery_list, self.gallery_listgallery_list, **kwargs)