1844 lines
78 KiB
Python
1844 lines
78 KiB
Python
from __future__ import print_function, absolute_import
|
|
import os
|
|
import glob
|
|
import re
|
|
import sys
|
|
import urllib
|
|
import tarfile
|
|
import zipfile
|
|
import os.path as osp
|
|
from scipy.io import loadmat
|
|
import numpy as np
|
|
import h5py
|
|
from scipy.misc import imsave
|
|
|
|
from utils import mkdir_if_missing, write_json, read_json
|
|
|
|
"""Image ReID"""
|
|
|
|
class Market1501(object):
|
|
"""
|
|
Market1501
|
|
|
|
Reference:
|
|
Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
|
|
|
|
URL: http://www.liangzheng.org/Project/project_reid.html
|
|
|
|
Dataset statistics:
|
|
# identities: 1501 (+1 for background)
|
|
# images: 12936 (train) + 3368 (query) + 15913 (gallery)
|
|
"""
|
|
dataset_dir = 'market1501'
|
|
|
|
def __init__(self, root='data', **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
|
|
self.query_dir = osp.join(self.dataset_dir, 'query')
|
|
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
|
|
|
|
self._check_before_run()
|
|
|
|
train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
|
|
query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
|
|
gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
|
|
|
|
print("=> Market1501 loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.train_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
|
if not osp.exists(self.query_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.query_dir))
|
|
if not osp.exists(self.gallery_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
|
|
|
|
def _process_dir(self, dir_path, relabel=False):
|
|
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
|
pattern = re.compile(r'([-\d]+)_c(\d)')
|
|
|
|
pid_container = set()
|
|
for img_path in img_paths:
|
|
pid, _ = map(int, pattern.search(img_path).groups())
|
|
if pid == -1: continue # junk images are just ignored
|
|
pid_container.add(pid)
|
|
pid2label = {pid:label for label, pid in enumerate(pid_container)}
|
|
|
|
dataset = []
|
|
for img_path in img_paths:
|
|
pid, camid = map(int, pattern.search(img_path).groups())
|
|
if pid == -1: continue # junk images are just ignored
|
|
assert 0 <= pid <= 1501 # pid == 0 means background
|
|
assert 1 <= camid <= 6
|
|
camid -= 1 # index starts from 0
|
|
if relabel: pid = pid2label[pid]
|
|
dataset.append((img_path, pid, camid))
|
|
|
|
num_pids = len(pid_container)
|
|
num_imgs = len(dataset)
|
|
return dataset, num_pids, num_imgs
|
|
|
|
class CUHK03(object):
|
|
"""
|
|
CUHK03
|
|
|
|
Reference:
|
|
Li et al. DeepReID: Deep Filter Pairing Neural Network for Person Re-identification. CVPR 2014.
|
|
|
|
URL: http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html#!
|
|
|
|
Dataset statistics:
|
|
# identities: 1360
|
|
# images: 13164
|
|
# cameras: 6
|
|
# splits: 20 (classic)
|
|
|
|
Args:
|
|
split_id (int): split index (default: 0)
|
|
cuhk03_labeled (bool): whether to load labeled images; if false, detected images are loaded (default: False)
|
|
"""
|
|
dataset_dir = 'cuhk03'
|
|
|
|
def __init__(self, root='data', split_id=0, cuhk03_labeled=False, cuhk03_classic_split=False, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.data_dir = osp.join(self.dataset_dir, 'cuhk03_release')
|
|
self.raw_mat_path = osp.join(self.data_dir, 'cuhk-03.mat')
|
|
|
|
self.imgs_detected_dir = osp.join(self.dataset_dir, 'images_detected')
|
|
self.imgs_labeled_dir = osp.join(self.dataset_dir, 'images_labeled')
|
|
|
|
self.split_classic_det_json_path = osp.join(self.dataset_dir, 'splits_classic_detected.json')
|
|
self.split_classic_lab_json_path = osp.join(self.dataset_dir, 'splits_classic_labeled.json')
|
|
|
|
self.split_new_det_json_path = osp.join(self.dataset_dir, 'splits_new_detected.json')
|
|
self.split_new_lab_json_path = osp.join(self.dataset_dir, 'splits_new_labeled.json')
|
|
|
|
self.split_new_det_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_detected.mat')
|
|
self.split_new_lab_mat_path = osp.join(self.dataset_dir, 'cuhk03_new_protocol_config_labeled.mat')
|
|
|
|
self._check_before_run()
|
|
self._preprocess()
|
|
|
|
if cuhk03_labeled:
|
|
image_type = 'labeled'
|
|
split_path = self.split_classic_lab_json_path if cuhk03_classic_split else self.split_new_lab_json_path
|
|
else:
|
|
image_type = 'detected'
|
|
split_path = self.split_classic_det_json_path if cuhk03_classic_split else self.split_new_det_json_path
|
|
|
|
splits = read_json(split_path)
|
|
assert split_id < len(splits), "Condition split_id ({}) < len(splits) ({}) is false".format(split_id, len(splits))
|
|
split = splits[split_id]
|
|
print("Split index = {}".format(split_id))
|
|
|
|
train = split['train']
|
|
query = split['query']
|
|
gallery = split['gallery']
|
|
|
|
num_train_pids = split['num_train_pids']
|
|
num_query_pids = split['num_query_pids']
|
|
num_gallery_pids = split['num_gallery_pids']
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
|
|
num_train_imgs = split['num_train_imgs']
|
|
num_query_imgs = split['num_query_imgs']
|
|
num_gallery_imgs = split['num_gallery_imgs']
|
|
num_total_imgs = num_train_imgs + num_query_imgs
|
|
|
|
print("=> CUHK03 ({}) loaded".format(image_type))
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.data_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.data_dir))
|
|
if not osp.exists(self.raw_mat_path):
|
|
raise RuntimeError("'{}' is not available".format(self.raw_mat_path))
|
|
if not osp.exists(self.split_new_det_mat_path):
|
|
raise RuntimeError("'{}' is not available".format(self.split_new_det_mat_path))
|
|
if not osp.exists(self.split_new_lab_mat_path):
|
|
raise RuntimeError("'{}' is not available".format(self.split_new_lab_mat_path))
|
|
|
|
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)
|
|
"""
|
|
print("Note: if root path is changed, the previously generated json files need to be re-generated (delete them first)")
|
|
if osp.exists(self.imgs_labeled_dir) and \
|
|
osp.exists(self.imgs_detected_dir) and \
|
|
osp.exists(self.split_classic_det_json_path) and \
|
|
osp.exists(self.split_classic_lab_json_path) and \
|
|
osp.exists(self.split_new_det_json_path) and \
|
|
osp.exists(self.split_new_lab_json_path):
|
|
return
|
|
|
|
mkdir_if_missing(self.imgs_detected_dir)
|
|
mkdir_if_missing(self.imgs_labeled_dir)
|
|
|
|
print("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
|
|
# 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)
|
|
imsave(img_path, img)
|
|
img_paths.append(img_path)
|
|
return img_paths
|
|
|
|
def _extract_img(name):
|
|
print("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))
|
|
print("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
|
|
|
|
print("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,
|
|
})
|
|
|
|
write_json(splits_classic_det, self.split_classic_det_json_path)
|
|
write_json(splits_classic_lab, self.split_classic_lab_json_path)
|
|
|
|
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
|
|
|
|
print("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],
|
|
}]
|
|
write_json(splits, self.split_new_det_json_path)
|
|
|
|
print("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],
|
|
}]
|
|
write_json(splits, self.split_new_lab_json_path)
|
|
|
|
class DukeMTMCreID(object):
|
|
"""
|
|
DukeMTMC-reID
|
|
|
|
Reference:
|
|
1. Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.
|
|
2. Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.
|
|
|
|
URL: https://github.com/layumi/DukeMTMC-reID_evaluation
|
|
|
|
Dataset statistics:
|
|
# identities: 1404 (train + query)
|
|
# images:16522 (train) + 2228 (query) + 17661 (gallery)
|
|
# cameras: 8
|
|
"""
|
|
dataset_dir = 'dukemtmc-reid'
|
|
|
|
def __init__(self, root='data', **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.train_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_train')
|
|
self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/query')
|
|
self.gallery_dir = osp.join(self.dataset_dir, 'DukeMTMC-reID/bounding_box_test')
|
|
|
|
self._check_before_run()
|
|
|
|
train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
|
|
query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
|
|
gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
|
|
|
|
print("=> DukeMTMC-reID loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.train_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
|
if not osp.exists(self.query_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.query_dir))
|
|
if not osp.exists(self.gallery_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
|
|
|
|
def _process_dir(self, dir_path, relabel=False):
|
|
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
|
pattern = re.compile(r'([-\d]+)_c(\d)')
|
|
|
|
pid_container = set()
|
|
for img_path in img_paths:
|
|
pid, _ = map(int, pattern.search(img_path).groups())
|
|
pid_container.add(pid)
|
|
pid2label = {pid:label for label, pid in enumerate(pid_container)}
|
|
|
|
dataset = []
|
|
for img_path in img_paths:
|
|
pid, camid = map(int, pattern.search(img_path).groups())
|
|
assert 1 <= camid <= 8
|
|
camid -= 1 # index starts from 0
|
|
if relabel: pid = pid2label[pid]
|
|
dataset.append((img_path, pid, camid))
|
|
|
|
num_pids = len(pid_container)
|
|
num_imgs = len(dataset)
|
|
return dataset, num_pids, num_imgs
|
|
|
|
class MSMT17(object):
|
|
"""
|
|
MSMT17
|
|
|
|
Reference:
|
|
Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018.
|
|
|
|
URL: http://www.pkuvmc.com/publications/msmt17.html
|
|
|
|
Dataset statistics:
|
|
# identities: 4101
|
|
# images: 32621 (train) + 11659 (query) + 82161 (gallery)
|
|
# cameras: 15
|
|
"""
|
|
dataset_dir = 'msmt17'
|
|
|
|
def __init__(self, root='data', **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.train_dir = osp.join(self.dataset_dir, 'MSMT17_V1/train')
|
|
self.test_dir = osp.join(self.dataset_dir, 'MSMT17_V1/test')
|
|
self.list_train_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_train.txt')
|
|
self.list_val_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_val.txt')
|
|
self.list_query_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_query.txt')
|
|
self.list_gallery_path = osp.join(self.dataset_dir, 'MSMT17_V1/list_gallery.txt')
|
|
|
|
self._check_before_run()
|
|
train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, self.list_train_path)
|
|
#val, num_val_pids, num_val_imgs = self._process_dir(self.train_dir, self.list_val_path)
|
|
query, num_query_pids, num_query_imgs = self._process_dir(self.test_dir, self.list_query_path)
|
|
gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.test_dir, self.list_gallery_path)
|
|
|
|
#train += val
|
|
#num_train_imgs += num_val_imgs
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
|
|
|
|
print("=> MSMT17 loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.train_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
|
if not osp.exists(self.test_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.test_dir))
|
|
|
|
def _process_dir(self, dir_path, list_path):
|
|
with open(list_path, 'r') as txt:
|
|
lines = txt.readlines()
|
|
dataset = []
|
|
pid_container = set()
|
|
for img_idx, img_info in enumerate(lines):
|
|
img_path, pid = img_info.split(' ')
|
|
pid = int(pid) # no need to relabel
|
|
camid = int(img_path.split('_')[2])
|
|
img_path = osp.join(dir_path, img_path)
|
|
dataset.append((img_path, pid, camid))
|
|
pid_container.add(pid)
|
|
num_imgs = len(dataset)
|
|
num_pids = len(pid_container)
|
|
# check if pid starts from 0 and increments with 1
|
|
for idx, pid in enumerate(pid_container):
|
|
assert idx == pid, "See code comment for explanation"
|
|
return dataset, num_pids, num_imgs
|
|
|
|
class VIPeR(object):
|
|
"""
|
|
VIPeR
|
|
|
|
Reference:
|
|
Gray et al. Evaluating appearance models for recognition, reacquisition, and tracking. PETS 2007.
|
|
|
|
URL: https://vision.soe.ucsc.edu/node/178
|
|
|
|
Dataset statistics:
|
|
# identities: 632
|
|
# images: 632 x 2 = 1264
|
|
# cameras: 2
|
|
"""
|
|
dataset_dir = 'viper'
|
|
|
|
def __init__(self, root='data', split_id=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.dataset_url = 'http://users.soe.ucsc.edu/~manduchi/VIPeR.v1.0.zip'
|
|
self.cam_a_path = osp.join(self.dataset_dir, 'VIPeR', 'cam_a')
|
|
self.cam_b_path = osp.join(self.dataset_dir, 'VIPeR', 'cam_b')
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
|
|
self._download_data()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
|
|
train = split['train']
|
|
query = split['query'] # query and gallery share the same images
|
|
gallery = split['gallery']
|
|
|
|
train = [tuple(item) for item in train]
|
|
query = [tuple(item) for item in query]
|
|
gallery = [tuple(item) for item in gallery]
|
|
|
|
num_train_pids = split['num_train_pids']
|
|
num_query_pids = split['num_query_pids']
|
|
num_gallery_pids = split['num_gallery_pids']
|
|
|
|
num_train_imgs = len(train)
|
|
num_query_imgs = len(query)
|
|
num_gallery_imgs = len(gallery)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs
|
|
|
|
print("=> VIPeR loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _download_data(self):
|
|
if osp.exists(self.dataset_dir):
|
|
print("This dataset has been downloaded.")
|
|
return
|
|
|
|
print("Creating directory {}".format(self.dataset_dir))
|
|
mkdir_if_missing(self.dataset_dir)
|
|
fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
|
|
|
|
print("Downloading VIPeR dataset")
|
|
urllib.urlretrieve(self.dataset_url, fpath)
|
|
|
|
print("Extracting files")
|
|
zip_ref = zipfile.ZipFile(fpath, 'r')
|
|
zip_ref.extractall(self.dataset_dir)
|
|
zip_ref.close()
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.cam_a_path):
|
|
raise RuntimeError("'{}' is not available".format(self.cam_a_path))
|
|
if not osp.exists(self.cam_b_path):
|
|
raise RuntimeError("'{}' is not available".format(self.cam_b_path))
|
|
|
|
def _prepare_split(self):
|
|
if not osp.exists(self.split_path):
|
|
print("Creating 10 random splits")
|
|
|
|
cam_a_imgs = sorted(glob.glob(osp.join(self.cam_a_path, '*.bmp')))
|
|
cam_b_imgs = sorted(glob.glob(osp.join(self.cam_b_path, '*.bmp')))
|
|
assert len(cam_a_imgs) == len(cam_b_imgs)
|
|
num_pids = len(cam_a_imgs)
|
|
print("Number of identities: {}".format(num_pids))
|
|
num_train_pids = num_pids // 2
|
|
|
|
splits = []
|
|
for _ in range(10):
|
|
order = np.arange(num_pids)
|
|
np.random.shuffle(order)
|
|
train_idxs = order[:num_train_pids]
|
|
test_idxs = order[num_train_pids:]
|
|
assert not bool(set(train_idxs) & set(test_idxs)), "Error: train and test overlap"
|
|
|
|
train = []
|
|
for pid, idx in enumerate(train_idxs):
|
|
cam_a_img = cam_a_imgs[idx]
|
|
cam_b_img = cam_b_imgs[idx]
|
|
train.append((cam_a_img, pid, 0))
|
|
train.append((cam_b_img, pid, 1))
|
|
|
|
test = []
|
|
for pid, idx in enumerate(test_idxs):
|
|
cam_a_img = cam_a_imgs[idx]
|
|
cam_b_img = cam_b_imgs[idx]
|
|
test.append((cam_a_img, pid, 0))
|
|
test.append((cam_b_img, pid, 1))
|
|
|
|
split = {'train': train, 'query': test, 'gallery': test,
|
|
'num_train_pids': num_train_pids,
|
|
'num_query_pids': num_pids - num_train_pids,
|
|
'num_gallery_pids': num_pids - num_train_pids
|
|
}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created".format(len(splits)))
|
|
write_json(splits, self.split_path)
|
|
print("Split file saved to {}".format(self.split_path))
|
|
|
|
print("Splits created")
|
|
|
|
class GRID(object):
|
|
"""
|
|
GRID
|
|
|
|
Reference:
|
|
Loy et al. Multi-camera activity correlation analysis. CVPR 2009.
|
|
|
|
URL: http://personal.ie.cuhk.edu.hk/~ccloy/downloads_qmul_underground_reid.html
|
|
|
|
Dataset statistics:
|
|
# identities: 250
|
|
# images: 1275
|
|
# cameras: 8
|
|
"""
|
|
dataset_dir = 'grid'
|
|
|
|
def __init__(self, root='data', split_id=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.dataset_url = 'http://personal.ie.cuhk.edu.hk/~ccloy/files/datasets/underground_reid.zip'
|
|
self.probe_path = osp.join(self.dataset_dir, 'underground_reid', 'probe')
|
|
self.gallery_path = osp.join(self.dataset_dir, 'underground_reid', 'gallery')
|
|
self.split_mat_path = osp.join(self.dataset_dir, 'underground_reid', 'features_and_partitions.mat')
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
|
|
self._download_data()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
|
|
train = split['train']
|
|
query = split['query']
|
|
gallery = split['gallery']
|
|
|
|
train = [tuple(item) for item in train]
|
|
query = [tuple(item) for item in query]
|
|
gallery = [tuple(item) for item in gallery]
|
|
|
|
num_train_pids = split['num_train_pids']
|
|
num_query_pids = split['num_query_pids']
|
|
num_gallery_pids = split['num_gallery_pids']
|
|
|
|
num_train_imgs = len(train)
|
|
num_query_imgs = len(query)
|
|
num_gallery_imgs = len(gallery)
|
|
|
|
num_total_pids = num_train_pids + num_gallery_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
|
|
|
|
print("=> GRID loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.probe_path):
|
|
raise RuntimeError("'{}' is not available".format(self.probe_path))
|
|
if not osp.exists(self.gallery_path):
|
|
raise RuntimeError("'{}' is not available".format(self.gallery_path))
|
|
if not osp.exists(self.split_mat_path):
|
|
raise RuntimeError("'{}' is not available".format(self.split_mat_path))
|
|
|
|
def _download_data(self):
|
|
if osp.exists(self.dataset_dir):
|
|
print("This dataset has been downloaded.")
|
|
return
|
|
|
|
print("Creating directory {}".format(self.dataset_dir))
|
|
mkdir_if_missing(self.dataset_dir)
|
|
fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
|
|
|
|
print("Downloading GRID dataset")
|
|
urllib.urlretrieve(self.dataset_url, fpath)
|
|
|
|
print("Extracting files")
|
|
zip_ref = zipfile.ZipFile(fpath, 'r')
|
|
zip_ref.extractall(self.dataset_dir)
|
|
zip_ref.close()
|
|
|
|
def _prepare_split(self):
|
|
if not osp.exists(self.split_path):
|
|
print("Creating 10 random splits")
|
|
split_mat = loadmat(self.split_mat_path)
|
|
trainIdxAll = split_mat['trainIdxAll'][0] # length = 10
|
|
probe_img_paths = sorted(glob.glob(osp.join(self.probe_path, '*.jpeg')))
|
|
gallery_img_paths = sorted(glob.glob(osp.join(self.gallery_path, '*.jpeg')))
|
|
|
|
splits = []
|
|
for split_idx in range(10):
|
|
train_idxs = trainIdxAll[split_idx][0][0][2][0].tolist()
|
|
assert len(train_idxs) == 125
|
|
idx2label = {idx: label for label, idx in enumerate(train_idxs)}
|
|
|
|
train, query, gallery = [], [], []
|
|
|
|
# processing probe folder
|
|
for img_path in probe_img_paths:
|
|
img_name = osp.basename(img_path)
|
|
img_idx = int(img_name.split('_')[0])
|
|
camid = int(img_name.split('_')[1])
|
|
if img_idx in train_idxs:
|
|
# add to train data
|
|
train.append((img_path, idx2label[img_idx], camid))
|
|
else:
|
|
# add to query data
|
|
query.append((img_path, img_idx, camid))
|
|
|
|
# process gallery folder
|
|
for img_path in gallery_img_paths:
|
|
img_name = osp.basename(img_path)
|
|
img_idx = int(img_name.split('_')[0])
|
|
camid = int(img_name.split('_')[1])
|
|
if img_idx in train_idxs:
|
|
# add to train data
|
|
train.append((img_path, idx2label[img_idx], camid))
|
|
else:
|
|
# add to gallery data
|
|
gallery.append((img_path, img_idx, camid))
|
|
|
|
split = {'train': train, 'query': query, 'gallery': gallery,
|
|
'num_train_pids': 125,
|
|
'num_query_pids': 125,
|
|
'num_gallery_pids': 900,
|
|
}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created".format(len(splits)))
|
|
write_json(splits, self.split_path)
|
|
print("Split file saved to {}".format(self.split_path))
|
|
|
|
print("Splits created")
|
|
|
|
class CUHK01(object):
|
|
"""
|
|
CUHK01
|
|
|
|
Reference:
|
|
Li et al. Human Reidentification with Transferred Metric Learning. ACCV 2012.
|
|
|
|
URL: http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html
|
|
|
|
Dataset statistics:
|
|
# identities: 971
|
|
# images: 3884
|
|
# cameras: 4
|
|
"""
|
|
dataset_dir = 'cuhk01'
|
|
|
|
def __init__(self, root='data', split_id=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.zip_path = osp.join(self.dataset_dir, 'CUHK01.zip')
|
|
self.campus_dir = osp.join(self.dataset_dir, 'campus')
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
|
|
self._extract_file()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
|
|
train = split['train']
|
|
query = split['query']
|
|
gallery = split['gallery']
|
|
|
|
train = [tuple(item) for item in train]
|
|
query = [tuple(item) for item in query]
|
|
gallery = [tuple(item) for item in gallery]
|
|
|
|
num_train_pids = split['num_train_pids']
|
|
num_query_pids = split['num_query_pids']
|
|
num_gallery_pids = split['num_gallery_pids']
|
|
|
|
num_train_imgs = len(train)
|
|
num_query_imgs = len(query)
|
|
num_gallery_imgs = len(gallery)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs
|
|
|
|
print("=> CUHK01 loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _extract_file(self):
|
|
if not osp.exists(self.campus_dir):
|
|
print("Extracting files")
|
|
zip_ref = zipfile.ZipFile(self.zip_path, 'r')
|
|
zip_ref.extractall(self.dataset_dir)
|
|
zip_ref.close()
|
|
print("Files extracted")
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.campus_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.campus_dir))
|
|
|
|
def _prepare_split(self):
|
|
"""
|
|
Image name format: 0001001.png, where first four digits represent identity
|
|
and last four digits represent cameras. Camera 1&2 are considered the same
|
|
view and camera 3&4 are considered the same view.
|
|
"""
|
|
if not osp.exists(self.split_path):
|
|
print("Creating 10 random splits")
|
|
img_paths = sorted(glob.glob(osp.join(self.campus_dir, '*.png')))
|
|
img_list = []
|
|
pid_container = set()
|
|
for img_path in img_paths:
|
|
img_name = osp.basename(img_path)
|
|
pid = int(img_name[:4]) - 1
|
|
camid = (int(img_name[4:7]) - 1) // 2
|
|
img_list.append((img_path, pid, camid))
|
|
pid_container.add(pid)
|
|
|
|
num_pids = len(pid_container)
|
|
num_train_pids = num_pids // 2
|
|
|
|
splits = []
|
|
for _ in range(10):
|
|
order = np.arange(num_pids)
|
|
np.random.shuffle(order)
|
|
train_idxs = order[:num_train_pids]
|
|
train_idxs = np.sort(train_idxs)
|
|
idx2label = {idx: label for label, idx in enumerate(train_idxs)}
|
|
|
|
train, test = [], []
|
|
|
|
for img_path, pid, camid in img_list:
|
|
if pid in train_idxs:
|
|
train.append((img_path, idx2label[pid], camid))
|
|
else:
|
|
test.append((img_path, pid, camid))
|
|
|
|
split = {'train': train, 'query': test, 'gallery': test,
|
|
'num_train_pids': num_train_pids,
|
|
'num_query_pids': num_pids - num_train_pids,
|
|
'num_gallery_pids': num_pids - num_train_pids,
|
|
}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created".format(len(splits)))
|
|
write_json(splits, self.split_path)
|
|
print("Split file saved to {}".format(self.split_path))
|
|
|
|
print("Splits created")
|
|
|
|
class PRID450S(object):
|
|
"""
|
|
PRID450S
|
|
|
|
Reference:
|
|
Roth et al. Mahalanobis Distance Learning for Person Re-Identification. PR 2014.
|
|
|
|
URL: https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/prid450s/
|
|
|
|
Dataset statistics:
|
|
# identities: 450
|
|
# images: 900
|
|
# cameras: 2
|
|
"""
|
|
dataset_dir = 'prid450s'
|
|
|
|
def __init__(self, root='data', split_id=0, min_seq_len=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.dataset_url = 'https://files.icg.tugraz.at/f/8c709245bb/?raw=1'
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
self.cam_a_path = osp.join(self.dataset_dir, 'cam_a')
|
|
self.cam_b_path = osp.join(self.dataset_dir, 'cam_b')
|
|
|
|
self._download_data()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
|
|
train = split['train']
|
|
query = split['query']
|
|
gallery = split['gallery']
|
|
|
|
train = [tuple(item) for item in train]
|
|
query = [tuple(item) for item in query]
|
|
gallery = [tuple(item) for item in gallery]
|
|
|
|
num_train_pids = split['num_train_pids']
|
|
num_query_pids = split['num_query_pids']
|
|
num_gallery_pids = split['num_gallery_pids']
|
|
|
|
num_train_imgs = len(train)
|
|
num_query_imgs = len(query)
|
|
num_gallery_imgs = len(gallery)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs
|
|
|
|
print("=> PRID450S loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.cam_a_path):
|
|
raise RuntimeError("'{}' is not available".format(self.cam_a_path))
|
|
if not osp.exists(self.cam_b_path):
|
|
raise RuntimeError("'{}' is not available".format(self.cam_b_path))
|
|
|
|
def _download_data(self):
|
|
if osp.exists(self.dataset_dir):
|
|
print("This dataset has been downloaded.")
|
|
return
|
|
|
|
print("Creating directory {}".format(self.dataset_dir))
|
|
mkdir_if_missing(self.dataset_dir)
|
|
fpath = osp.join(self.dataset_dir, 'prid_450s.zip')
|
|
|
|
print("Downloading PRID450S dataset")
|
|
urllib.urlretrieve(self.dataset_url, fpath)
|
|
|
|
print("Extracting files")
|
|
zip_ref = zipfile.ZipFile(fpath, 'r')
|
|
zip_ref.extractall(self.dataset_dir)
|
|
zip_ref.close()
|
|
|
|
def _prepare_split(self):
|
|
if not osp.exists(self.split_path):
|
|
cam_a_imgs = sorted(glob.glob(osp.join(self.cam_a_path, 'img_*.png')))
|
|
cam_b_imgs = sorted(glob.glob(osp.join(self.cam_b_path, 'img_*.png')))
|
|
assert len(cam_a_imgs) == len(cam_b_imgs)
|
|
|
|
num_pids = len(cam_a_imgs)
|
|
num_train_pids = num_pids // 2
|
|
|
|
splits = []
|
|
for _ in range(10):
|
|
order = np.arange(num_pids)
|
|
np.random.shuffle(order)
|
|
train_idxs = np.sort(order[:num_train_pids])
|
|
idx2label = {idx: label for label, idx in enumerate(train_idxs)}
|
|
|
|
train, test = [], []
|
|
|
|
# processing camera a
|
|
for img_path in cam_a_imgs:
|
|
img_name = osp.basename(img_path)
|
|
img_idx = int(img_name.split('_')[1].split('.')[0])
|
|
if img_idx in train_idxs:
|
|
train.append((img_path, idx2label[img_idx], 0))
|
|
else:
|
|
test.append((img_path, img_idx, 0))
|
|
|
|
# processing camera b
|
|
for img_path in cam_b_imgs:
|
|
img_name = osp.basename(img_path)
|
|
img_idx = int(img_name.split('_')[1].split('.')[0])
|
|
if img_idx in train_idxs:
|
|
train.append((img_path, idx2label[img_idx], 1))
|
|
else:
|
|
test.append((img_path, img_idx, 1))
|
|
|
|
split = {'train': train, 'query': test, 'gallery': test,
|
|
'num_train_pids': num_train_pids,
|
|
'num_query_pids': num_pids - num_train_pids,
|
|
'num_gallery_pids': num_pids - num_train_pids,
|
|
}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created".format(len(splits)))
|
|
write_json(splits, self.split_path)
|
|
print("Split file saved to {}".format(self.split_path))
|
|
|
|
print("Splits created")
|
|
|
|
class iLIDS(object):
|
|
"""
|
|
iLIDS (for single shot setting)
|
|
|
|
Reference:
|
|
Wang et al. Person Re-Identification by Video Ranking. ECCV 2014.
|
|
|
|
URL: http://www.eecs.qmul.ac.uk/~xiatian/downloads_qmul_iLIDS-VID_ReID_dataset.html
|
|
|
|
Dataset statistics:
|
|
# identities: 300
|
|
# images: 600
|
|
# cameras: 2
|
|
"""
|
|
dataset_dir = 'ilids-vid'
|
|
|
|
def __init__(self, root='data', split_id=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.dataset_url = 'http://www.eecs.qmul.ac.uk/~xiatian/iLIDS-VID/iLIDS-VID.tar'
|
|
self.data_dir = osp.join(self.dataset_dir, 'i-LIDS-VID')
|
|
self.split_dir = osp.join(self.dataset_dir, 'train-test people splits')
|
|
self.split_mat_path = osp.join(self.split_dir, 'train_test_splits_ilidsvid.mat')
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
self.cam_1_path = osp.join(self.dataset_dir, 'i-LIDS-VID/images/cam1') # differ from video
|
|
self.cam_2_path = osp.join(self.dataset_dir, 'i-LIDS-VID/images/cam2')
|
|
|
|
self._download_data()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
train_dirs, test_dirs = split['train'], split['test']
|
|
print("# train identites: {}, # test identites {}".format(len(train_dirs), len(test_dirs)))
|
|
|
|
train, num_train_imgs, num_train_pids = self._process_data(train_dirs, cam1=True, cam2=True)
|
|
query, num_query_imgs, num_query_pids = self._process_data(test_dirs, cam1=True, cam2=False)
|
|
gallery, num_gallery_imgs, num_gallery_pids = self._process_data(test_dirs, cam1=False, cam2=True)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_imgs = num_train_imgs + num_query_imgs
|
|
|
|
print("=> PRID450S loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # images")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _download_data(self):
|
|
if osp.exists(self.dataset_dir):
|
|
print("This dataset has been downloaded.")
|
|
return
|
|
|
|
mkdir_if_missing(self.dataset_dir)
|
|
fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
|
|
|
|
print("Downloading iLIDS-VID dataset")
|
|
urllib.urlretrieve(self.dataset_url, fpath)
|
|
|
|
print("Extracting files")
|
|
tar = tarfile.open(fpath)
|
|
tar.extractall(path=self.dataset_dir)
|
|
tar.close()
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.data_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.data_dir))
|
|
if not osp.exists(self.split_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.split_dir))
|
|
|
|
def _prepare_split(self):
|
|
if not osp.exists(self.split_path):
|
|
print("Creating splits")
|
|
mat_split_data = loadmat(self.split_mat_path)['ls_set']
|
|
|
|
num_splits = mat_split_data.shape[0]
|
|
num_total_ids = mat_split_data.shape[1]
|
|
assert num_splits == 10
|
|
assert num_total_ids == 300
|
|
num_ids_each = num_total_ids/2
|
|
|
|
# pids in mat_split_data are indices, so we need to transform them
|
|
# to real pids
|
|
person_cam1_dirs = sorted(glob.glob(osp.join(self.cam_1_path, '*')))
|
|
person_cam2_dirs = sorted(glob.glob(osp.join(self.cam_2_path, '*')))
|
|
|
|
person_cam1_dirs = [osp.basename(item) for item in person_cam1_dirs]
|
|
person_cam2_dirs = [osp.basename(item) for item in person_cam2_dirs]
|
|
|
|
# make sure persons in one camera view can be found in the other camera view
|
|
assert set(person_cam1_dirs) == set(person_cam2_dirs)
|
|
|
|
splits = []
|
|
for i_split in range(num_splits):
|
|
# first 50% for testing and the remaining for training, following Wang et al. ECCV'14.
|
|
train_idxs = sorted(list(mat_split_data[i_split,num_ids_each:]))
|
|
test_idxs = sorted(list(mat_split_data[i_split,:num_ids_each]))
|
|
|
|
train_idxs = [int(i)-1 for i in train_idxs]
|
|
test_idxs = [int(i)-1 for i in test_idxs]
|
|
|
|
# transform pids to person dir names
|
|
train_dirs = [person_cam1_dirs[i] for i in train_idxs]
|
|
test_dirs = [person_cam1_dirs[i] for i in test_idxs]
|
|
|
|
split = {'train': train_dirs, 'test': test_dirs}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created, following Wang et al. ECCV'14".format(len(splits)))
|
|
print("Split file is saved to {}".format(self.split_path))
|
|
write_json(splits, self.split_path)
|
|
|
|
def _process_data(self, dirnames, cam1=True, cam2=True):
|
|
dirname2pid = {dirname:i for i, dirname in enumerate(dirnames)}
|
|
dataset = []
|
|
|
|
for i, dirname in enumerate(dirnames):
|
|
if cam1:
|
|
pdir = osp.join(self.cam_1_path, dirname)
|
|
img_path = glob.glob(osp.join(pdir, '*.png'))
|
|
# only one image is available in one folder
|
|
assert len(img_path) == 1
|
|
img_path = img_path[0]
|
|
pid = dirname2pid[dirname]
|
|
dataset.append((img_path, pid, 0))
|
|
|
|
if cam2:
|
|
pdir = osp.join(self.cam_2_path, dirname)
|
|
img_path = glob.glob(osp.join(pdir, '*.png'))
|
|
# only one image is available in one folder
|
|
assert len(img_path) == 1
|
|
img_path = img_path[0]
|
|
pid = dirname2pid[dirname]
|
|
dataset.append((img_path, pid, 1))
|
|
|
|
num_imgs = len(dataset)
|
|
num_pids = len(dirnames)
|
|
|
|
return dataset, num_imgs, num_pids
|
|
|
|
"""Video ReID"""
|
|
|
|
class Mars(object):
|
|
"""
|
|
MARS
|
|
|
|
Reference:
|
|
Zheng et al. MARS: A Video Benchmark for Large-Scale Person Re-identification. ECCV 2016.
|
|
|
|
URL: http://www.liangzheng.com.cn/Project/project_mars.html
|
|
|
|
Dataset statistics:
|
|
# identities: 1261
|
|
# tracklets: 8298 (train) + 1980 (query) + 9330 (gallery)
|
|
# cameras: 6
|
|
"""
|
|
dataset_dir = 'mars'
|
|
|
|
def __init__(self, root='data', min_seq_len=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.train_name_path = osp.join(self.dataset_dir, 'info/train_name.txt')
|
|
self.test_name_path = osp.join(self.dataset_dir, 'info/test_name.txt')
|
|
self.track_train_info_path = osp.join(self.dataset_dir, 'info/tracks_train_info.mat')
|
|
self.track_test_info_path = osp.join(self.dataset_dir, 'info/tracks_test_info.mat')
|
|
self.query_IDX_path = osp.join(self.dataset_dir, 'info/query_IDX.mat')
|
|
|
|
self._check_before_run()
|
|
|
|
# prepare meta data
|
|
train_names = self._get_names(self.train_name_path)
|
|
test_names = self._get_names(self.test_name_path)
|
|
track_train = loadmat(self.track_train_info_path)['track_train_info'] # numpy.ndarray (8298, 4)
|
|
track_test = loadmat(self.track_test_info_path)['track_test_info'] # numpy.ndarray (12180, 4)
|
|
query_IDX = loadmat(self.query_IDX_path)['query_IDX'].squeeze() # numpy.ndarray (1980,)
|
|
query_IDX -= 1 # index from 0
|
|
track_query = track_test[query_IDX,:]
|
|
gallery_IDX = [i for i in range(track_test.shape[0]) if i not in query_IDX]
|
|
track_gallery = track_test[gallery_IDX,:]
|
|
|
|
train, num_train_tracklets, num_train_pids, num_train_imgs = \
|
|
self._process_data(train_names, track_train, home_dir='bbox_train', relabel=True, min_seq_len=min_seq_len)
|
|
|
|
query, num_query_tracklets, num_query_pids, num_query_imgs = \
|
|
self._process_data(test_names, track_query, home_dir='bbox_test', relabel=False, min_seq_len=min_seq_len)
|
|
|
|
gallery, num_gallery_tracklets, num_gallery_pids, num_gallery_imgs = \
|
|
self._process_data(test_names, track_gallery, home_dir='bbox_test', relabel=False, min_seq_len=min_seq_len)
|
|
|
|
num_imgs_per_tracklet = num_train_imgs + num_query_imgs + num_gallery_imgs
|
|
min_num = np.min(num_imgs_per_tracklet)
|
|
max_num = np.max(num_imgs_per_tracklet)
|
|
avg_num = np.mean(num_imgs_per_tracklet)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
|
|
|
|
print("=> MARS loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # tracklets")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
|
|
print(" number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.train_name_path):
|
|
raise RuntimeError("'{}' is not available".format(self.train_name_path))
|
|
if not osp.exists(self.test_name_path):
|
|
raise RuntimeError("'{}' is not available".format(self.test_name_path))
|
|
if not osp.exists(self.track_train_info_path):
|
|
raise RuntimeError("'{}' is not available".format(self.track_train_info_path))
|
|
if not osp.exists(self.track_test_info_path):
|
|
raise RuntimeError("'{}' is not available".format(self.track_test_info_path))
|
|
if not osp.exists(self.query_IDX_path):
|
|
raise RuntimeError("'{}' is not available".format(self.query_IDX_path))
|
|
|
|
def _get_names(self, fpath):
|
|
names = []
|
|
with open(fpath, 'r') as f:
|
|
for line in f:
|
|
new_line = line.rstrip()
|
|
names.append(new_line)
|
|
return names
|
|
|
|
def _process_data(self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0):
|
|
assert home_dir in ['bbox_train', 'bbox_test']
|
|
num_tracklets = meta_data.shape[0]
|
|
pid_list = list(set(meta_data[:,2].tolist()))
|
|
num_pids = len(pid_list)
|
|
|
|
if relabel: pid2label = {pid:label for label, pid in enumerate(pid_list)}
|
|
tracklets = []
|
|
num_imgs_per_tracklet = []
|
|
|
|
for tracklet_idx in range(num_tracklets):
|
|
data = meta_data[tracklet_idx,...]
|
|
start_index, end_index, pid, camid = data
|
|
if pid == -1: continue # junk images are just ignored
|
|
assert 1 <= camid <= 6
|
|
if relabel: pid = pid2label[pid]
|
|
camid -= 1 # index starts from 0
|
|
img_names = names[start_index-1:end_index]
|
|
|
|
# make sure image names correspond to the same person
|
|
pnames = [img_name[:4] for img_name in img_names]
|
|
assert len(set(pnames)) == 1, "Error: a single tracklet contains different person images"
|
|
|
|
# make sure all images are captured under the same camera
|
|
camnames = [img_name[5] for img_name in img_names]
|
|
assert len(set(camnames)) == 1, "Error: images are captured under different cameras!"
|
|
|
|
# append image names with directory information
|
|
img_paths = [osp.join(self.dataset_dir, home_dir, img_name[:4], img_name) for img_name in img_names]
|
|
if len(img_paths) >= min_seq_len:
|
|
img_paths = tuple(img_paths)
|
|
tracklets.append((img_paths, pid, camid))
|
|
num_imgs_per_tracklet.append(len(img_paths))
|
|
|
|
num_tracklets = len(tracklets)
|
|
|
|
return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
|
|
|
|
class iLIDSVID(object):
|
|
"""
|
|
iLIDS-VID
|
|
|
|
Reference:
|
|
Wang et al. Person Re-Identification by Video Ranking. ECCV 2014.
|
|
|
|
URL: http://www.eecs.qmul.ac.uk/~xiatian/downloads_qmul_iLIDS-VID_ReID_dataset.html
|
|
|
|
Dataset statistics:
|
|
# identities: 300
|
|
# tracklets: 600
|
|
# cameras: 2
|
|
"""
|
|
dataset_dir = 'ilids-vid'
|
|
|
|
def __init__(self, root='data', split_id=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.dataset_url = 'http://www.eecs.qmul.ac.uk/~xiatian/iLIDS-VID/iLIDS-VID.tar'
|
|
self.data_dir = osp.join(self.dataset_dir, 'i-LIDS-VID')
|
|
self.split_dir = osp.join(self.dataset_dir, 'train-test people splits')
|
|
self.split_mat_path = osp.join(self.split_dir, 'train_test_splits_ilidsvid.mat')
|
|
self.split_path = osp.join(self.dataset_dir, 'splits.json')
|
|
self.cam_1_path = osp.join(self.dataset_dir, 'i-LIDS-VID/sequences/cam1')
|
|
self.cam_2_path = osp.join(self.dataset_dir, 'i-LIDS-VID/sequences/cam2')
|
|
|
|
self._download_data()
|
|
self._check_before_run()
|
|
|
|
self._prepare_split()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
train_dirs, test_dirs = split['train'], split['test']
|
|
print("# train identites: {}, # test identites {}".format(len(train_dirs), len(test_dirs)))
|
|
|
|
train, num_train_tracklets, num_train_pids, num_imgs_train = \
|
|
self._process_data(train_dirs, cam1=True, cam2=True)
|
|
query, num_query_tracklets, num_query_pids, num_imgs_query = \
|
|
self._process_data(test_dirs, cam1=True, cam2=False)
|
|
gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
|
|
self._process_data(test_dirs, cam1=False, cam2=True)
|
|
|
|
num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
|
|
min_num = np.min(num_imgs_per_tracklet)
|
|
max_num = np.max(num_imgs_per_tracklet)
|
|
avg_num = np.mean(num_imgs_per_tracklet)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
|
|
|
|
print("=> iLIDS-VID loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # tracklets")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
|
|
print(" number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _download_data(self):
|
|
if osp.exists(self.dataset_dir):
|
|
print("This dataset has been downloaded.")
|
|
return
|
|
|
|
mkdir_if_missing(self.dataset_dir)
|
|
fpath = osp.join(self.dataset_dir, osp.basename(self.dataset_url))
|
|
|
|
print("Downloading iLIDS-VID dataset")
|
|
urllib.urlretrieve(self.dataset_url, fpath)
|
|
|
|
print("Extracting files")
|
|
tar = tarfile.open(fpath)
|
|
tar.extractall(path=self.dataset_dir)
|
|
tar.close()
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.data_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.data_dir))
|
|
if not osp.exists(self.split_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.split_dir))
|
|
|
|
def _prepare_split(self):
|
|
if not osp.exists(self.split_path):
|
|
print("Creating splits")
|
|
mat_split_data = loadmat(self.split_mat_path)['ls_set']
|
|
|
|
num_splits = mat_split_data.shape[0]
|
|
num_total_ids = mat_split_data.shape[1]
|
|
assert num_splits == 10
|
|
assert num_total_ids == 300
|
|
num_ids_each = num_total_ids/2
|
|
|
|
# pids in mat_split_data are indices, so we need to transform them
|
|
# to real pids
|
|
person_cam1_dirs = sorted(glob.glob(osp.join(self.cam_1_path, '*')))
|
|
person_cam2_dirs = sorted(glob.glob(osp.join(self.cam_2_path, '*')))
|
|
|
|
person_cam1_dirs = [osp.basename(item) for item in person_cam1_dirs]
|
|
person_cam2_dirs = [osp.basename(item) for item in person_cam2_dirs]
|
|
|
|
# make sure persons in one camera view can be found in the other camera view
|
|
assert set(person_cam1_dirs) == set(person_cam2_dirs)
|
|
|
|
splits = []
|
|
for i_split in range(num_splits):
|
|
# first 50% for testing and the remaining for training, following Wang et al. ECCV'14.
|
|
train_idxs = sorted(list(mat_split_data[i_split,num_ids_each:]))
|
|
test_idxs = sorted(list(mat_split_data[i_split,:num_ids_each]))
|
|
|
|
train_idxs = [int(i)-1 for i in train_idxs]
|
|
test_idxs = [int(i)-1 for i in test_idxs]
|
|
|
|
# transform pids to person dir names
|
|
train_dirs = [person_cam1_dirs[i] for i in train_idxs]
|
|
test_dirs = [person_cam1_dirs[i] for i in test_idxs]
|
|
|
|
split = {'train': train_dirs, 'test': test_dirs}
|
|
splits.append(split)
|
|
|
|
print("Totally {} splits are created, following Wang et al. ECCV'14".format(len(splits)))
|
|
print("Split file is saved to {}".format(self.split_path))
|
|
write_json(splits, self.split_path)
|
|
|
|
print("Splits created")
|
|
|
|
def _process_data(self, dirnames, cam1=True, cam2=True):
|
|
tracklets = []
|
|
num_imgs_per_tracklet = []
|
|
dirname2pid = {dirname:i for i, dirname in enumerate(dirnames)}
|
|
|
|
for dirname in dirnames:
|
|
if cam1:
|
|
person_dir = osp.join(self.cam_1_path, dirname)
|
|
img_names = glob.glob(osp.join(person_dir, '*.png'))
|
|
assert len(img_names) > 0
|
|
img_names = tuple(img_names)
|
|
pid = dirname2pid[dirname]
|
|
tracklets.append((img_names, pid, 0))
|
|
num_imgs_per_tracklet.append(len(img_names))
|
|
|
|
if cam2:
|
|
person_dir = osp.join(self.cam_2_path, dirname)
|
|
img_names = glob.glob(osp.join(person_dir, '*.png'))
|
|
assert len(img_names) > 0
|
|
img_names = tuple(img_names)
|
|
pid = dirname2pid[dirname]
|
|
tracklets.append((img_names, pid, 1))
|
|
num_imgs_per_tracklet.append(len(img_names))
|
|
|
|
num_tracklets = len(tracklets)
|
|
num_pids = len(dirnames)
|
|
|
|
return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
|
|
|
|
class PRID2011(object):
|
|
"""
|
|
PRID2011
|
|
|
|
Reference:
|
|
Hirzer et al. Person Re-Identification by Descriptive and Discriminative Classification. SCIA 2011.
|
|
|
|
URL: https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/
|
|
|
|
Dataset statistics:
|
|
# identities: 200
|
|
# tracklets: 400
|
|
# cameras: 2
|
|
"""
|
|
dataset_dir = 'prid2011'
|
|
|
|
def __init__(self, root='data', split_id=0, min_seq_len=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.split_path = osp.join(self.dataset_dir, 'splits_prid2011.json')
|
|
self.cam_a_path = osp.join(self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_a')
|
|
self.cam_b_path = osp.join(self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_b')
|
|
|
|
self._check_before_run()
|
|
splits = read_json(self.split_path)
|
|
if split_id >= len(splits):
|
|
raise ValueError("split_id exceeds range, received {}, but expected between 0 and {}".format(split_id, len(splits)-1))
|
|
split = splits[split_id]
|
|
train_dirs, test_dirs = split['train'], split['test']
|
|
print("# train identites: {}, # test identites {}".format(len(train_dirs), len(test_dirs)))
|
|
|
|
train, num_train_tracklets, num_train_pids, num_imgs_train = \
|
|
self._process_data(train_dirs, cam1=True, cam2=True)
|
|
query, num_query_tracklets, num_query_pids, num_imgs_query = \
|
|
self._process_data(test_dirs, cam1=True, cam2=False)
|
|
gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
|
|
self._process_data(test_dirs, cam1=False, cam2=True)
|
|
|
|
num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
|
|
min_num = np.min(num_imgs_per_tracklet)
|
|
max_num = np.max(num_imgs_per_tracklet)
|
|
avg_num = np.mean(num_imgs_per_tracklet)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
|
|
|
|
print("=> PRID2011 loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # tracklets")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
|
|
print(" number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
|
|
def _process_data(self, dirnames, cam1=True, cam2=True):
|
|
tracklets = []
|
|
num_imgs_per_tracklet = []
|
|
dirname2pid = {dirname:i for i, dirname in enumerate(dirnames)}
|
|
|
|
for dirname in dirnames:
|
|
if cam1:
|
|
person_dir = osp.join(self.cam_a_path, dirname)
|
|
img_names = glob.glob(osp.join(person_dir, '*.png'))
|
|
assert len(img_names) > 0
|
|
img_names = tuple(img_names)
|
|
pid = dirname2pid[dirname]
|
|
tracklets.append((img_names, pid, 0))
|
|
num_imgs_per_tracklet.append(len(img_names))
|
|
|
|
if cam2:
|
|
person_dir = osp.join(self.cam_b_path, dirname)
|
|
img_names = glob.glob(osp.join(person_dir, '*.png'))
|
|
assert len(img_names) > 0
|
|
img_names = tuple(img_names)
|
|
pid = dirname2pid[dirname]
|
|
tracklets.append((img_names, pid, 1))
|
|
num_imgs_per_tracklet.append(len(img_names))
|
|
|
|
num_tracklets = len(tracklets)
|
|
num_pids = len(dirnames)
|
|
|
|
return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
|
|
|
|
class DukeMTMCVidReID(object):
|
|
"""
|
|
DukeMTMCVidReID
|
|
|
|
Reference:
|
|
Wu et al. Exploit the Unknown Gradually: One-Shot Video-Based Person
|
|
Re-Identification by Stepwise Learning. CVPR 2018.
|
|
|
|
URL: https://github.com/Yu-Wu/Exploit-Unknown-Gradually
|
|
|
|
Dataset statistics:
|
|
# identities: 702 (train) + 702 (test)
|
|
# tracklets: 2196 (train) + 2636 (test)
|
|
"""
|
|
dataset_dir = 'dukemtmc-vidreid'
|
|
|
|
def __init__(self, root='data', min_seq_len=0, **kwargs):
|
|
self.dataset_dir = osp.join(root, self.dataset_dir)
|
|
self.train_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/train_split')
|
|
self.query_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/query_split')
|
|
self.gallery_dir = osp.join(self.dataset_dir, 'dukemtmc_videoReID/gallery_split')
|
|
self.split_train_json_path = osp.join(self.dataset_dir, 'split_train.json')
|
|
self.split_query_json_path = osp.join(self.dataset_dir, 'split_query.json')
|
|
self.split_gallery_json_path = osp.join(self.dataset_dir, 'split_gallery.json')
|
|
|
|
self.min_seq_len = min_seq_len
|
|
self._check_before_run()
|
|
print("Note: if root path is changed, the previously generated json files need to be re-generated (so delete them first)")
|
|
|
|
train, num_train_tracklets, num_train_pids, num_imgs_train = \
|
|
self._process_dir(self.train_dir, self.split_train_json_path, relabel=True)
|
|
query, num_query_tracklets, num_query_pids, num_imgs_query = \
|
|
self._process_dir(self.query_dir, self.split_query_json_path, relabel=False)
|
|
gallery, num_gallery_tracklets, num_gallery_pids, num_imgs_gallery = \
|
|
self._process_dir(self.gallery_dir, self.split_gallery_json_path, relabel=False)
|
|
|
|
num_imgs_per_tracklet = num_imgs_train + num_imgs_query + num_imgs_gallery
|
|
min_num = np.min(num_imgs_per_tracklet)
|
|
max_num = np.max(num_imgs_per_tracklet)
|
|
avg_num = np.mean(num_imgs_per_tracklet)
|
|
|
|
num_total_pids = num_train_pids + num_query_pids
|
|
num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets
|
|
|
|
print("=> DukeMTMC-VideoReID loaded")
|
|
print("Dataset statistics:")
|
|
print(" ------------------------------")
|
|
print(" subset | # ids | # tracklets")
|
|
print(" ------------------------------")
|
|
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
|
|
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
|
|
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
|
|
print(" ------------------------------")
|
|
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
|
|
print(" number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
|
|
print(" ------------------------------")
|
|
|
|
self.train = train
|
|
self.query = query
|
|
self.gallery = gallery
|
|
|
|
self.num_train_pids = num_train_pids
|
|
self.num_query_pids = num_query_pids
|
|
self.num_gallery_pids = num_gallery_pids
|
|
|
|
def _check_before_run(self):
|
|
"""Check if all files are available before going deeper"""
|
|
if not osp.exists(self.dataset_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
|
|
if not osp.exists(self.train_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.train_dir))
|
|
if not osp.exists(self.query_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.query_dir))
|
|
if not osp.exists(self.gallery_dir):
|
|
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
|
|
|
|
def _process_dir(self, dir_path, json_path, relabel):
|
|
if osp.exists(json_path):
|
|
print("=> {} generated before, awesome!".format(json_path))
|
|
split = read_json(json_path)
|
|
return split['tracklets'], split['num_tracklets'], split['num_pids'], split['num_imgs_per_tracklet']
|
|
|
|
print("=> Automatically generating split (might take a while for the first time, have a coffe)")
|
|
pdirs = glob.glob(osp.join(dir_path, '*')) # avoid .DS_Store
|
|
print("Processing {} with {} person identities".format(dir_path, len(pdirs)))
|
|
|
|
pid_container = set()
|
|
for pdir in pdirs:
|
|
pid = int(osp.basename(pdir))
|
|
pid_container.add(pid)
|
|
pid2label = {pid:label for label, pid in enumerate(pid_container)}
|
|
|
|
tracklets = []
|
|
num_imgs_per_tracklet = []
|
|
for pdir in pdirs:
|
|
pid = int(osp.basename(pdir))
|
|
if relabel: pid = pid2label[pid]
|
|
tdirs = glob.glob(osp.join(pdir, '*'))
|
|
for tdir in tdirs:
|
|
raw_img_paths = glob.glob(osp.join(tdir, '*.jpg'))
|
|
num_imgs = len(raw_img_paths)
|
|
|
|
if num_imgs < self.min_seq_len:
|
|
continue
|
|
|
|
num_imgs_per_tracklet.append(num_imgs)
|
|
img_paths = []
|
|
for img_idx in range(num_imgs):
|
|
# some tracklet starts from 0002 instead of 0001
|
|
img_idx_name = 'F' + str(img_idx+1).zfill(4)
|
|
res = glob.glob(osp.join(tdir, '*' + img_idx_name + '*.jpg'))
|
|
if len(res) == 0:
|
|
print("Warn: index name {} in {} is missing, jump to next".format(img_idx_name, tdir))
|
|
continue
|
|
img_paths.append(res[0])
|
|
img_name = osp.basename(img_paths[0])
|
|
camid = int(img_name[5]) - 1 # index-0
|
|
img_paths = tuple(img_paths)
|
|
tracklets.append((img_paths, pid, camid))
|
|
|
|
num_pids = len(pid_container)
|
|
num_tracklets = len(tracklets)
|
|
|
|
print("Saving split to {}".format(json_path))
|
|
split_dict = {
|
|
'tracklets': tracklets,
|
|
'num_tracklets': num_tracklets,
|
|
'num_pids': num_pids,
|
|
'num_imgs_per_tracklet': num_imgs_per_tracklet,
|
|
}
|
|
write_json(split_dict, json_path)
|
|
|
|
return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
|
|
|
|
"""Create dataset"""
|
|
|
|
__img_factory = {
|
|
'market1501': Market1501,
|
|
'cuhk03': CUHK03,
|
|
'dukemtmcreid': DukeMTMCreID,
|
|
'msmt17': MSMT17,
|
|
'viper': VIPeR,
|
|
'grid': GRID,
|
|
'cuhk01': CUHK01,
|
|
'prid450s': PRID450S,
|
|
'ilids': iLIDS,
|
|
}
|
|
|
|
__vid_factory = {
|
|
'mars': Mars,
|
|
'ilidsvid': iLIDSVID,
|
|
'prid': PRID2011,
|
|
'dukemtmcvidreid': DukeMTMCVidReID,
|
|
}
|
|
|
|
def get_names():
|
|
return list(__img_factory.keys()) + list(__vid_factory.keys())
|
|
|
|
def init_img_dataset(name, **kwargs):
|
|
if name not in __img_factory.keys():
|
|
raise KeyError("Invalid dataset, got '{}', but expected to be one of {}".format(name, __img_factory.keys()))
|
|
return __img_factory[name](**kwargs)
|
|
|
|
def init_vid_dataset(name, **kwargs):
|
|
if name not in __vid_factory.keys():
|
|
raise KeyError("Invalid dataset, got '{}', but expected to be one of {}".format(name, __vid_factory.keys()))
|
|
return __vid_factory[name](**kwargs) |