Add modified vehicle datasets with small, medium and large test datastes

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

fastreid/data/datasets/vehicleid.py

@@ -20,19 +20,21 @@ class VehicleID(ImageDataset):
 
     URL: `<https://pkuml.org/resources/pku-vehicleid.html>`_
 
-    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)

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)

fastreid/data/datasets/veriwild.py

@@ -19,21 +19,24 @@ class VeRiWild(ImageDataset):
 
     URL: `<https://github.com/PKU-IMRE/VERI-Wild>`_
 
-    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)