mirror of https://github.com/JDAI-CV/fast-reid.git
139 lines
4.7 KiB
Python
139 lines
4.7 KiB
Python
# 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]
|
|
if is_train:
|
|
vid = self.dataset_name + '_' + str(vid)
|
|
dataset.append((self.imgid2imgpath[imgid], vid, 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_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)
|