def check_raw_file(self):
assert self.zipfiles_dir.exists()
if not self.raw_data_folder.exists():
unpack_file(self.zipfiles_dir, self.store_dir, self.logger)
check_path(self.split_config_path, create=True)
if not self.split_new_det_mat_path.exists() or not self.split_new_lab_mat_path.exists():
config_file_dir = self.split_config_path / 'person-re-ranking-master'
if not config_file_dir.exists():
config_file_path = self.store_dir / 'person-re-ranking-master.zip'
if not config_file_path.exists():
urlretrieve(self.split_new_mat_url, config_file_path)
unpack_file(config_file_path, self.split_config_path, self.logger)
if not self.split_new_det_mat_path.exists():
copy_file_to(
self.split_config_path / 'person-re-ranking-master/evaluation/data/CUHK03/cuhk03_new_protocol_config_detected.mat',
self.store_dir)
if not self.split_new_lab_mat_path.exists():
copy_file_to(
self.split_config_path / 'person-re-ranking-master/evaluation/data/CUHK03/cuhk03_new_protocol_config_labeled.mat',
self.store_dir)
if config_file_dir.exists():
remove_folder(config_file_dir)
assert self.raw_mat_path.exists()
def load(self, model_name=None):
if model_name is None:
self.model_parallel.load_state_dict(torch.load(check_path(self.task_dir / 'output/model/model.pkl')))
self.logger.info('Model restored from ' + str(self.task_dir / 'output/model/model.pkl'))
else:
self.model_parallel.load_state_dict(torch.load(check_path(self.task_dir / str('output/model/' + model_name + '.pkl'))))
self.logger.info('Model restored from ' + str(self.task_dir / 'output/model/' + model_name + '.pkl'))
def check_raw_file(self):
if not self.zipfiles_dir.exists():
check_path(self.zipfiles_dir.parent, create=True)
urlretrieve(self.dataset_url, self.zipfiles_dir)
if not self.raw_data_folder.exists():
unpack_file(self.zipfiles_dir, self.store_dir, self.logger)
assert self.split_mat_path.exists()
def check_raw_file(self):
check_path(self.raw_data_folder, create=True)
raw_file_folder = self.zipfiles_dir.parent
raw_data_list = ['bbox_train', 'bbox_test', 'MARS-evaluation-master']
for raw_data in raw_data_list:
if not (self.raw_data_folder / raw_data).exists():
unpack_file(raw_file_folder / str(raw_data + '.zip'), self.raw_data_folder, self.logger)
def __init__(self,
file_folder,
dataset_name,
store_type,
image_dir_list,
track_info,
resize=None,
logger=None):
self.file_folder = Path(file_folder)
self.dataset_name = dataset_name
self.store_type = store_type
self.len = len(image_dir_list)
self.resize_hw = resize
self.logger = logger
if self.resize_hw is not None:
assert isinstance(self.resize_hw, (tuple, list))
assert len(self.resize_hw) == 2
check_path(self.file_folder, create=True)
dataset_dir = list(
self.file_folder.glob(self.dataset_name + '_*.' + self.store_type))
self._store_factory = {
'db': {
'read': self.read_lmdb,
'write': self.write_lmdb
},
'h5': {
'read': self.read_h5,
'write': self.write_h5
}
}
self.store_optical_flow = False # if you want to use optical flow, enable it.
if len(dataset_dir) == 0:
with_optical_flow = False
read_img = self.read_only_img
if track_info is not None and np.max(
track_info[:, -1]) > 1 and self.store_optical_flow:
try:
from .optical_flow_tools import OpticalFlowManager
self.of_generator = OpticalFlowManager(
image_dir_list, track_info, self.resize_hw)
read_img = self.read_img_with_of
with_optical_flow = True
except ImportError as e:
logger.error(e)
self.with_optical_flow = with_optical_flow
self.read_img = read_img
self.dataset_dir = self.write(image_dir_list)
else:
assert len(dataset_dir) == 1
self.dataset_dir = dataset_dir[0]
self.parse()
self.init()
def mode(self, value):
if value == 'Train':
self._mode = value
empty_folder(check_path(self.task_dir / 'output/log', create=True))
empty_folder(
check_path(self.task_dir / 'output/model', create=True))
empty_folder(
check_path(self.task_dir / 'output/result', create=True))
elif value == 'Test':
self._mode = value
else:
raise KeyError
def check_epoch(self, epoch_id):
with open(check_path(self.task_dir / 'test_epoch_id.txt',
create=False)) as f:
info_list_tmp = f.readlines()
info_list = [
int(info_i) for info_i in info_list_tmp if info_i != '\n'
]
assert len(info_list) > 0
idx_mode = info_list[0]
if idx_mode == 0:
assert len(info_list) >= 4
idx_list = list(range(info_list[1], info_list[2],
info_list[3]))
idx_list += info_list[4:]
epoch_size = info_list[2]
idx_list += [epoch_size]
elif idx_mode == 1:
idx_list = info_list[1:]
epoch_size = max(info_list[1:])
else:
raise ValueError
train_flag = True
test_flag = False
if epoch_id >= epoch_size:
train_flag = False
if epoch_id in idx_list:
test_flag = True
return train_flag, test_flag
def __init__(self, name, split_id, store_type, root_dir, logger):
self.name = name
self.split_id = split_id
self.store_type = store_type
self.store_dir = root_dir / str(name)
self.logger = logger
self.dict_dir = Path(check_path(self.store_dir, create=True)) / (str(name) + '_dict.json')
def __init__(self, folder_path, tag, device, logger):
self._folder_path = folder_path
self.logger = logger
self.device = device
self._father_name = self._folder_path.parts[-1]
self._env_name = self.device[
'name'] + '_' + self._father_name + '_' + tag
self._target_dir = check_path(self._folder_path / 'output/log', True)
try:
self._viz = VisionServer(self._env_name, self._target_dir,
self.device)
except ImportError as e:
self.logger.error(e)
self._viz = VisionBuffer(self._env_name, self._target_dir)
self.logger.info(
'Visdom is not installed. The visual data is stored locally.')
except (ConnectionError, ConnectionRefusedError) as e:
self.logger.error(e)
self._viz = VisionBuffer(self._env_name, self._target_dir)
self.logger.info(
'Visdom is not installed, but no server connection. The visual data is stored locally.'
)
else:
self.logger.info(
'Visdom is installed. Using vsidom to display visual data on the web page.'
)
self._mode = None
def __init__(self, task_dir, time_tag, device):
self.father_path = task_dir
self.time_tag = time_tag
self.logger = InfoLogger(
check_path(self.father_path / 'output/log', create=False),
'log.txt', self.time_tag).logger
self.visual = Visual('Visdom', self.father_path, self.time_tag, device,
self.logger)
def __init__(self, env_name, file_path, device):
self._file_path = file_path
self.device = device
self._env_name = env_name
self._target_dir = check_path(self._file_path, True)
self._env_path = check_path(self.device['web_env_dir'], True)
self._env_dir = self._env_path / (self._env_name + '.json')
self._win_box = {}
from visdom import Visdom
self._viz = Visdom(port=self.device['web_port'],
env=self._env_name,
server=self.device['web_host'],
raise_exceptions=True)
self._mode = None
self._flush_epoch = 5
self._flush_id = 0
def __init__(self, task_dir, seed=1, mode='Train'):
self.seed = seed
self.task_dir = task_dir
self.time_tag = time_tag()
self.mode = mode
self._device_dict = {
'pc': {
'name': 'pc',
'root':
Path('/home/username/data'), # Store the extracted files
'rawfiles':
Path('/home/username/rawfiles'
), # The location of the original compressed file
'Model': Path(
'/home/username/model'
), # Store the officially downloaded torch model parameters
'web_env_dir': '/home/username/ignore',
'web_host': "http://localhost",
'web_port': 31094,
'num_workers': 4,
'test_batch_size': 16
},
'server': {
'name': 'server',
'root': Path('/data/data'), # Store the extracted files
'rawfiles':
Path('/data/rawfiles'
), # The location of the original compressed file
'Model': Path(
'/data/model'
), # Store the officially downloaded torch model parameters
'web_env_dir': '/data/ignore',
'web_host': "http://localhost",
'web_port': 31094,
'num_workers': 16,
'test_batch_size': 64
}
}
self._dataset_box = [
'iLIDS-VID', 'PRID-2011', 'LPW', 'MARS', 'VIPeR', 'Market1501',
'CUHK03', 'CUHK01', 'DukeMTMCreID', 'GRID', 'DukeMTMC-VideoReID'
]
self.init_device()
self.recorder = Recorder(check_path(self.task_dir, create=False),
self.time_tag, self.device)
self.logger = self.recorder.logger
self.logger.info('Device: ' + self.device['name'])
self.logger.info('{0:-^60}'.format('Set Seed ' + str(self.seed)))
def store_search_example(father_path, false_example, right_example):
file_path = check_path(father_path / 'output/log', True)
file_dir = file_path / 'search_result.h5'
with h5py.File(file_dir, 'w') as f:
f['num'] = len(false_example)
for i_false, false_i in enumerate(false_example):
f['false_img_' + str(i_false)] = false_i[0].astype(np.string_)
f['false_dst_' + str(i_false)] = false_i[1]
f['false_id_' + str(i_false)] = false_i[2]
f['right_img_' +
str(i_false)] = right_example[i_false][0].astype(np.string_)
f['right_dst_' + str(i_false)] = right_example[i_false][1]
f['right_id_' + str(i_false)] = right_example[i_false][2]
def init_device(self):
if self._device_dict['pc']['rawfiles'].exists():
device = self._device_dict['pc']
else:
device = self._device_dict['server']
check_path(device['root'], create=True)
check_path(device['Model'], create=True)
check_path(device['web_env_dir'], create=True)
self.device = device
def _preprocess(self):
"""
This function is a bit complex and ugly, what it does is
1. Extract data from cuhk-03.mat and save as png images.
2. Create 20 classic splits. (Li et al. CVPR'14)
3. Create new split. (Zhong et al. CVPR'17)
"""
self.logger.info(
"Note: if root path is changed, the previously generated json files need to be re-generated (delete them first)")
if self.imgs_labeled_dir.exists() and \
self.imgs_detected_dir.exists() and \
self.split_classic_det_json_path.exists() and \
self.split_classic_lab_json_path.exists() and \
self.split_new_det_json_path.exists() and \
self.split_new_lab_json_path.exists():
return
check_path(self.imgs_detected_dir, create=True)
check_path(self.imgs_labeled_dir, create=True)
self.logger.info("Extract image data from {} and save as png".format(self.raw_mat_path))
mat = h5py.File(self.raw_mat_path, 'r')
def _deref(ref):
return mat[ref][:].T
def _process_images(img_refs, campid, pid, save_dir):
img_paths = [] # Note: some persons only have images for one view
for imgid, img_ref in enumerate(img_refs):
img = _deref(img_ref)
# skip empty cell
if img.size == 0 or img.ndim < 3: continue
img = Image.fromarray(img, mode='RGB')
# images are saved with the following format, index-1 (ensure uniqueness)
# campid: index of camera pair (1-5)
# pid: index of person in 'campid'-th camera pair
# viewid: index of view, {1, 2}
# imgid: index of image, (1-10)
viewid = 1 if imgid < 5 else 2
img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(campid + 1, pid + 1, viewid, imgid + 1)
img_path = osp.join(save_dir, img_name)
img.save(img_path)
img_paths.append(img_path)
return img_paths
def _extract_img(name):
self.logger.info("Processing {} images (extract and save) ...".format(name))
meta_data = []
imgs_dir = self.imgs_detected_dir if name == 'detected' else self.imgs_labeled_dir
for campid, camp_ref in enumerate(mat[name][0]):
camp = _deref(camp_ref)
num_pids = camp.shape[0]
for pid in range(num_pids):
img_paths = _process_images(camp[pid, :], campid, pid, imgs_dir)
assert len(img_paths) > 0, "campid{}-pid{} has no images".format(campid, pid)
meta_data.append((campid + 1, pid + 1, img_paths))
self.logger.info("done camera pair {} with {} identities".format(campid + 1, num_pids))
return meta_data
meta_detected = _extract_img('detected')
meta_labeled = _extract_img('labeled')
def _extract_classic_split(meta_data, test_split):
train, test = [], []
num_train_pids, num_test_pids = 0, 0
num_train_imgs, num_test_imgs = 0, 0
for i, (campid, pid, img_paths) in enumerate(meta_data):
if [campid, pid] in test_split:
for img_path in img_paths:
camid = int(osp.basename(img_path).split('_')[2])
test.append((img_path, num_test_pids, camid))
num_test_pids += 1
num_test_imgs += len(img_paths)
else:
for img_path in img_paths:
camid = int(osp.basename(img_path).split('_')[2])
train.append((img_path, num_train_pids, camid))
num_train_pids += 1
num_train_imgs += len(img_paths)
return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs
self.logger.info("Creating classic splits (# = 20) ...")
splits_classic_det, splits_classic_lab = [], []
for split_ref in mat['testsets'][0]:
test_split = _deref(split_ref).tolist()
# create split for detected images
train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
_extract_classic_split(meta_detected, test_split)
splits_classic_det.append({
'train': train, 'query': test, 'gallery': test,
'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
})
# create split for labeled images
train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \
_extract_classic_split(meta_labeled, test_split)
splits_classic_lab.append({
'train': train, 'query': test, 'gallery': test,
'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs,
'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs,
'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs,
})
DataPacker.dump(splits_classic_det, self.split_classic_det_json_path, self.logger)
DataPacker.dump(splits_classic_lab, self.split_classic_lab_json_path, self.logger)
mat.close()
def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel):
tmp_set = []
unique_pids = set()
for idx in idxs:
img_name = filelist[idx][0]
camid = int(img_name.split('_')[2])
pid = pids[idx]
if relabel: pid = pid2label[pid]
img_path = osp.join(img_dir, img_name)
tmp_set.append((img_path, int(pid), camid))
unique_pids.add(pid)
return tmp_set, len(unique_pids), len(idxs)
def _extract_new_split(split_dict, img_dir):
train_idxs = split_dict['train_idx'].flatten() - 1 # index-0
pids = split_dict['labels'].flatten()
train_pids = set(pids[train_idxs])
pid2label = {pid: label for label, pid in enumerate(train_pids)}
query_idxs = split_dict['query_idx'].flatten() - 1
gallery_idxs = split_dict['gallery_idx'].flatten() - 1
filelist = split_dict['filelist'].flatten()
train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True)
query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False)
gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False)
return train_info, query_info, gallery_info
self.logger.info("Creating new splits for detected images (767/700) ...")
train_info, query_info, gallery_info = _extract_new_split(
loadmat(self.split_new_det_mat_path),
self.imgs_detected_dir,
)
splits = [{
'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
}]
DataPacker.dump(splits, self.split_new_det_json_path)
self.logger.info("Creating new splits for labeled images (767/700) ...")
train_info, query_info, gallery_info = _extract_new_split(
loadmat(self.split_new_lab_mat_path),
self.imgs_labeled_dir,
)
splits = [{
'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0],
'num_train_pids': train_info[1], 'num_train_imgs': train_info[2],
'num_query_pids': query_info[1], 'num_query_imgs': query_info[2],
'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2],
}]
DataPacker.dump(splits, self.split_new_lab_json_path)
def __init__(self, fpath=None):
self.console = sys.stdout
self.file = None
if fpath is not None:
check_path(osp.dirname(fpath), create=True)
self.file = open(fpath, 'w')
def store_performance(self, cmc_box):
file_path = check_path(
self.task_dir / 'output/result',
create=True) / str('cmc_' + self.time_tag + '.json')
DataPacker.dump(cmc_box, file_path, self.logger)
