-
Notifications
You must be signed in to change notification settings - Fork 0
/
sensereid.py
52 lines (40 loc) · 1.66 KB
/
sensereid.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
# encoding: utf-8
"""
@author: xingyu liao
@contact: [email protected]
"""
import os
from glob import glob
from fastreid.data.datasets import DATASET_REGISTRY
from fastreid.data.datasets.bases import ImageDataset
__all__ = ['SenseReID', ]
@DATASET_REGISTRY.register()
class SenseReID(ImageDataset):
"""Sense reid
"""
dataset_dir = "SenseReID"
dataset_name = "SenseReID"
def __init__(self, root='datasets', **kwargs):
self.root = root
self.train_path = os.path.join(self.root, self.dataset_dir)
required_files = [self.train_path]
self.check_before_run(required_files)
#train = self.process_train(self.train_path)
#change by hby, add query and gallery
train = self.process_train(self.train_path, file_path_list=['test_gallery', 'test_probe'])
query = self.process_train(self.train_path, file_path_list=['test_probe'])
gallery = self.process_train(self.train_path, file_path_list=['test_gallery'])
super().__init__(train, query, gallery, **kwargs)
def process_train(self, train_path, file_path_list):
data = []
#file_path_list = ['test_gallery', 'test_probe']
for file_path in file_path_list:
sub_file = os.path.join(train_path, file_path)
img_name = glob(os.path.join(sub_file, "*.jpg"))
for img_path in img_name:
img_name = img_path.split('/')[-1]
img_info = img_name.split('_')
pid = self.dataset_name + "_" + img_info[0]
camid = self.dataset_name + "_" + img_info[1].split('.')[0]
data.append([img_path, pid, camid])
return data