From 25aecab2daed8089fd07e698150b45e82d37f768 Mon Sep 17 00:00:00 2001 From: KaiyangZhou Date: Thu, 24 Jan 2019 22:12:02 +0000 Subject: [PATCH] bug fix --- torchreid/datasets/msmt17.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torchreid/datasets/msmt17.py b/torchreid/datasets/msmt17.py index 2c20a0c..5b8cf43 100644 --- a/torchreid/datasets/msmt17.py +++ b/torchreid/datasets/msmt17.py @@ -22,7 +22,7 @@ from .bases import BaseImageDataset # Log: # 22.01.2019: v1 and v2 only differ in dir names _TRAIN_DIR_KEY = 'train_dir' -_TEST_DIR_KEY = 'train_dir' +_TEST_DIR_KEY = 'test_dir' _VERSION = { 'MSMT17_V1': { _TRAIN_DIR_KEY: 'train', @@ -57,8 +57,8 @@ class MSMT17(BaseImageDataset): has_main_dir = False for main_dir in _VERSION: if osp.exists(osp.join(self.dataset_dir, main_dir)): - train_dir = _VERSION[main_dir][TRAIN_DIR_KEY] - test_dir = _VERSION[main_dir][TEST_DIR_KEY] + train_dir = _VERSION[main_dir][_TRAIN_DIR_KEY] + test_dir = _VERSION[main_dir][_TEST_DIR_KEY] has_main_dir = True break assert has_main_dir, "Dataset folder not found"