diff --git a/fastreid/data/build.py b/fastreid/data/build.py index 3f9d8d2..2e05c1d 100644 --- a/fastreid/data/build.py +++ b/fastreid/data/build.py @@ -76,7 +76,7 @@ def build_reid_test_loader(cfg, dataset_name): test_loader = DataLoader( test_set, batch_sampler=batch_sampler, - num_workers=2, # save some memory + num_workers=4, # save some memory collate_fn=fast_batch_collator, pin_memory=True, ) diff --git a/fastreid/data/common.py b/fastreid/data/common.py index 82feb3b..e1d349e 100644 --- a/fastreid/data/common.py +++ b/fastreid/data/common.py @@ -33,7 +33,10 @@ class CommDataset(Dataset): return len(self.img_items) def __getitem__(self, index): - img_path, pid, camid = self.img_items[index] + img_item = self.img_items[index] + img_path = img_item[0] + pid = img_item[1] + camid = img_item[2] img = read_image(img_path) if self.transform is not None: img = self.transform(img) if self.relabel: diff --git a/fastreid/data/datasets/saivt.py b/fastreid/data/datasets/saivt.py new file mode 100644 index 0000000..c1f2d07 --- /dev/null +++ b/fastreid/data/datasets/saivt.py @@ -0,0 +1,45 @@ +# 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__ = ['SAIVT', ] + + +@DATASET_REGISTRY.register() +class SAIVT(ImageDataset): + dataset_dir = "SAIVT-SoftBio" + dataset_name = "saivt" + + 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) + + super().__init__(train, [], [], **kwargs) + + def process_train(self, train_path): + data = [] + + pid_path = os.path.join(train_path, "cropped_images") + pid_list = os.listdir(pid_path) + + for pid_name in pid_list: + pid = self.dataset_name + '_' + pid_name + img_list = glob(os.path.join(pid_path, pid_name, "*.jpeg")) + for img_path in img_list: + img_name = os.path.basename(img_path) + camid = self.dataset_name + '_' + img_name.split('-')[2] + data.append([img_path, pid, camid]) + return data diff --git a/fastreid/data/transforms/build.py b/fastreid/data/transforms/build.py index 244d62a..cc5acf5 100644 --- a/fastreid/data/transforms/build.py +++ b/fastreid/data/transforms/build.py @@ -52,6 +52,7 @@ def build_transforms(cfg, is_train=True): if do_autoaug: res.append(T.RandomApply([AutoAugment()], p=autoaug_prob)) + res.append(T.Resize(size_train, interpolation=3)) if do_flip: res.append(T.RandomHorizontalFlip(p=flip_prob))