diff --git a/fastreid/data/datasets/__init__.py b/fastreid/data/datasets/__init__.py index eb47a5e..2f10264 100644 --- a/fastreid/data/datasets/__init__.py +++ b/fastreid/data/datasets/__init__.py @@ -17,5 +17,5 @@ from .dukemtmcreid import DukeMTMC from .market1501 import Market1501 from .msmt17 import MSMT17 from .veri import VeRi -from .vehicleid import VehicleID -from .veriwild import VeRiWild +from .vehicleid import VehicleID, SmallVehicleID, MediumVehicleID, LargeVehicleID +from .veriwild import VeRiWild, SmallVeRiWild, MediumVeRiWild, LargeVeRiWild diff --git a/fastreid/data/datasets/vehicleid.py b/fastreid/data/datasets/vehicleid.py index 01d0a44..a7158c7 100644 --- a/fastreid/data/datasets/vehicleid.py +++ b/fastreid/data/datasets/vehicleid.py @@ -20,19 +20,21 @@ class VehicleID(ImageDataset): URL: ``_ - Dataset statistics: - - identities: 26267. - - images: 221763. + Train dataset statistics: + - identities: 13164. + - images: 113346. """ dataset_dir = 'vehicleid' - dataset_url = None - def __init__(self, root='/home/liuxinchen3/notespace/data', **kwargs): + def __init__(self, root='datasets', test_list='', **kwargs): self.dataset_dir = osp.join(root, self.dataset_dir) self.image_dir = osp.join(self.dataset_dir, 'image') self.train_list = osp.join(self.dataset_dir, 'train_test_split/train_list.txt') - self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_2400.txt') + if test_list: + self.test_list = test_list + else: + self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_13164.txt') required_files = [ self.dataset_dir, @@ -64,9 +66,57 @@ class VehicleID(ImageDataset): return dataset else: query = [] + gallery = [] for sample in dataset: if sample[1] not in vid_container: vid_container.add(sample[1]) + gallery.append(sample) + else: query.append(sample) - return query, dataset + return query, gallery + + +@DATASET_REGISTRY.register() +class SmallVehicleID(VehicleID): + """VehicleID. + Small test dataset statistics: + - identities: 800. + - images: 6493. + """ + + def __init__(self, root='datasets', **kwargs): + self.dataset_dir = osp.join(root, self.dataset_dir) + self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_800.txt') + + super(SmallVehicleID, self).__init__(root, self.test_list, **kwargs) + + +@DATASET_REGISTRY.register() +class MediumVehicleID(VehicleID): + """VehicleID. + Medium test dataset statistics: + - identities: 1600. + - images: 13377. + """ + + def __init__(self, root='datasets', **kwargs): + self.dataset_dir = osp.join(root, self.dataset_dir) + self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_1600.txt') + + super(MediumVehicleID, self).__init__(root, self.test_list, **kwargs) + + +@DATASET_REGISTRY.register() +class LargeVehicleID(VehicleID): + """VehicleID. + Large test dataset statistics: + - identities: 2400. + - images: 19777. + """ + + def __init__(self, root='datasets', **kwargs): + self.dataset_dir = osp.join(root, self.dataset_dir) + self.test_list = osp.join(self.dataset_dir, 'train_test_split/test_list_2400.txt') + + super(LargeVehicleID, self).__init__(root, self.test_list, **kwargs) diff --git a/fastreid/data/datasets/veri.py b/fastreid/data/datasets/veri.py index d17826f..5a0f0c0 100644 --- a/fastreid/data/datasets/veri.py +++ b/fastreid/data/datasets/veri.py @@ -26,7 +26,6 @@ class VeRi(ImageDataset): - images: 37778 (train) + 1678 (query) + 11579 (gallery). """ dataset_dir = 'veri' - dataset_url = None def __init__(self, root='/home/liuxinchen3/notespace/data', **kwargs): self.dataset_dir = osp.join(root, self.dataset_dir) diff --git a/fastreid/data/datasets/veriwild.py b/fastreid/data/datasets/veriwild.py index 1dc0879..b7e7af7 100644 --- a/fastreid/data/datasets/veriwild.py +++ b/fastreid/data/datasets/veriwild.py @@ -19,21 +19,24 @@ class VeRiWild(ImageDataset): URL: ``_ - Dataset statistics: - - identities: 40671. - - images: 416314. + Train dataset statistics: + - identities: 30671. + - images: 277797. """ dataset_dir = 'VERI-Wild' - dataset_url = None - def __init__(self, root='/home/liuxinchen3/notespace/data', **kwargs): + 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.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') 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, @@ -82,3 +85,51 @@ class VeRiWild(ImageDataset): 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): + self.dataset_dir = osp.join(root, self.dataset_dir) + self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_3000_query.txt') + self.gallery_list = osp.join(self.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): + self.dataset_dir = osp.join(root, self.dataset_dir) + self.query_list = osp.join(self.dataset_dir, 'train_test_split/test_5000_query.txt') + self.gallery_list = osp.join(self.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): + self.dataset_dir = osp.join(root, self.dataset_dir) + 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') + + super(LargeVeRiWild, self).__init__(root, self.query_list, self.gallery_list, **kwargs)