mirror of https://github.com/JDAI-CV/fast-reid.git
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
# encoding: utf-8
|
|
"""
|
|
@author: liaoxingyu
|
|
@contact: liaoxingyu2@jd.com
|
|
"""
|
|
|
|
import glob
|
|
import os.path as osp
|
|
import re
|
|
import mat4py
|
|
import numpy as np
|
|
|
|
from fastreid.data.datasets import DATASET_REGISTRY
|
|
|
|
from .bases import Dataset
|
|
|
|
|
|
@DATASET_REGISTRY.register()
|
|
class DukeMTMCAttr(Dataset):
|
|
"""DukeMTMCAttr.
|
|
|
|
Reference:
|
|
Lin, Yutian, et al. "Improving person re-identification by attribute and identity learning."
|
|
Pattern Recognition 95 (2019): 151-161.
|
|
|
|
URL: `<https://github.com/vana77/DukeMTMC-attribute>`_
|
|
|
|
The folder structure should be:
|
|
DukeMTMC-reID/
|
|
bounding_box_train/ # images
|
|
bounding_box_test/ # images
|
|
duke_attribute.mat
|
|
"""
|
|
dataset_dir = 'DukeMTMC-reID'
|
|
dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'
|
|
dataset_name = "dukemtmc"
|
|
|
|
def __init__(self, root='datasets', **kwargs):
|
|
self.root = root
|
|
self.dataset_dir = osp.join(self.root, self.dataset_dir)
|
|
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
|
|
self.test_dir = osp.join(self.dataset_dir, 'bounding_box_test')
|
|
|
|
required_files = [
|
|
self.dataset_dir,
|
|
self.train_dir,
|
|
self.test_dir,
|
|
]
|
|
self.check_before_run(required_files)
|
|
|
|
duke_attr = mat4py.loadmat(osp.join(self.dataset_dir, 'duke_attribute.mat'))['duke_attribute']
|
|
sorted_attrs = sorted(duke_attr['train'].keys())
|
|
sorted_attrs.remove('image_index')
|
|
attr_dict = {i: str(attr) for i, attr in enumerate(sorted_attrs)}
|
|
|
|
train = self.process_dir(self.train_dir, duke_attr['train'], sorted_attrs)
|
|
test = val = self.process_dir(self.test_dir, duke_attr['test'], sorted_attrs)
|
|
|
|
super(DukeMTMCAttr, self).__init__(train, val, test, attr_dict=attr_dict, **kwargs)
|
|
|
|
def process_dir(self, dir_path, annotation, sorted_attrs):
|
|
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
|
|
pattern = re.compile(r'([-\d]+)_c(\d)')
|
|
|
|
data = []
|
|
for img_path in img_paths:
|
|
pid, camid = map(int, pattern.search(img_path).groups())
|
|
assert 1 <= camid <= 8
|
|
|
|
img_index = annotation['image_index'].index(str(pid).zfill(4))
|
|
attrs = np.array([int(annotation[i][img_index]) - 1 for i in sorted_attrs], dtype=np.float32)
|
|
data.append((img_path, attrs))
|
|
|
|
return data
|