update prid and grid datasets

fastreid/data/datasets/__init__.py

@@ -21,6 +21,8 @@ from .AirportALERT import AirportALERT
 from .iLIDS import iLIDS
 from .pku import PKU
 from .prai import PRAI
+from .prid import PRID
+from .grid import GRID
 from .saivt import SAIVT
 from .sensereid import SenseReID
 from .sysu_mm import SYSU_mm

fastreid/data/datasets/grid.py

@@ -0,0 +1,44 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+from glob import glob
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['GRID', ]
+
+
+@DATASET_REGISTRY.register()
+class GRID(ImageDataset):
+    """GRID
+    """
+    dataset_dir = "underground_reid"
+    dataset_name = 'grid'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir, 'images')
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        img_paths = glob(os.path.join(train_path, "*.jpeg"))
+
+        for img_path in img_paths:
+            img_name = os.path.basename(img_path)
+            img_info = img_name.split('_')
+            pid = self.dataset_name + "_" + img_info[0]
+            camid = self.dataset_name + "_" + img_info[1]
+            data.append([img_path, pid, camid])
+        return data

fastreid/data/datasets/prid.py

@@ -0,0 +1,41 @@
+# encoding: utf-8
+"""
+@author:  xingyu liao
+@contact: sherlockliao01@gmail.com
+"""
+
+import os
+
+from fastreid.data.datasets import DATASET_REGISTRY
+from fastreid.data.datasets.bases import ImageDataset
+
+__all__ = ['PRID', ]
+
+
+@DATASET_REGISTRY.register()
+class PRID(ImageDataset):
+    """PRID
+    """
+    dataset_dir = "prid_2011"
+    dataset_name = 'prid'
+
+    def __init__(self, root='datasets', **kwargs):
+        self.root = root
+        self.train_path = os.path.join(self.root, self.dataset_dir, 'slim_train')
+
+        required_files = [self.train_path]
+        self.check_before_run(required_files)
+
+        train = self.process_train(self.train_path)
+
+        super().__init__(train, [], [], **kwargs)
+
+    def process_train(self, train_path):
+        data = []
+        for root, dirs, files in os.walk(train_path):
+            for img_name in filter(lambda x: x.endswith('.png'), files):
+                img_path = os.path.join(root, img_name)
+                pid = self.dataset_name + '_' + root.split('/')[-1].split('_')[1]
+                camid = self.dataset_name + '_' + img_name.split('_')[0]
+                data.append([img_path, pid, camid])
+        return data