290 lines
11 KiB
Python
290 lines
11 KiB
Python
|
from __future__ import absolute_import
|
||
|
|
from __future__ import print_function
|
||
|
|
from __future__ import division
|
||
|
|
|
||
|
|
import sys
|
||
|
|
import os
|
||
|
|
import os.path as osp
|
||
|
|
import numpy as np
|
||
|
|
import tarfile
|
||
|
|
import zipfile
|
||
|
|
import copy
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from torchreid.utils import read_image, mkdir_if_missing, download_url
|
||
|
|
|
||
|
|
|
||
|
|
class Dataset(object):
|
||
|
|
|
||
|
|
def __init__(self, train, query, gallery, transform=None, mode='train',
|
||
|
|
combineall=False, verbose=True, **kwargs):
|
||
|
|
self.train = train
|
||
|
|
self.query = query
|
||
|
|
self.gallery = gallery
|
||
|
|
self.transform = transform
|
||
|
|
self.mode = mode
|
||
|
|
self.combineall = combineall
|
||
|
|
self.verbose = verbose
|
||
|
|
|
||
|
|
self.num_train_pids = self.get_num_pids(self.train)
|
||
|
|
self.num_train_cams = self.get_num_cams(self.train)
|
||
|
|
|
||
|
|
if self.combineall:
|
||
|
|
self.combine_all()
|
||
|
|
|
||
|
|
if self.mode == 'train':
|
||
|
|
self.data = self.train
|
||
|
|
elif self.mode == 'query':
|
||
|
|
self.data = self.query
|
||
|
|
elif self.mode == 'gallery':
|
||
|
|
self.data = self.gallery
|
||
|
|
else:
|
||
|
|
raise ValueError('Invalid mode. Got {}, but expected to be '
|
||
|
|
'one of [train | query | gallery]'.format(self.mode))
|
||
|
|
|
||
|
|
if self.verbose:
|
||
|
|
self.show_summary()
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
raise NotImplementedError
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
return len(self.data)
|
||
|
|
|
||
|
|
def __add__(self, other):
|
||
|
|
"""Adds two datasets together (only the train set)"""
|
||
|
|
train = copy.deepcopy(self.train)
|
||
|
|
|
||
|
|
for img_path, pid, camid in other.train:
|
||
|
|
pid += self.num_train_pids
|
||
|
|
camid += self.num_train_cams
|
||
|
|
train.append((img_path, pid, camid))
|
||
|
|
|
||
|
|
return Dataset(train, self.query, self.gallery, transform=self.transform,
|
||
|
|
mode=self.mode, combineall=self.combineall, verbose=self.verbose)
|
||
|
|
|
||
|
|
def __radd__(self, other):
|
||
|
|
"""Supports sum([dataset1, dataset2, dataset3])"""
|
||
|
|
if other == 0:
|
||
|
|
return self
|
||
|
|
else:
|
||
|
|
return self.__add__(other)
|
||
|
|
|
||
|
|
def parse_data(self, data):
|
||
|
|
"""Parses data
|
||
|
|
|
||
|
|
:param data: data list containing tuples of (img_path(s), pid, camid).
|
||
|
|
:type data: list
|
||
|
|
|
||
|
|
:return num_pids: number of person identities
|
||
|
|
:rtype num_pids: int
|
||
|
|
:return num_cams: number of cameras
|
||
|
|
:rtype num_cams: int
|
||
|
|
"""
|
||
|
|
pids = set()
|
||
|
|
cams = set()
|
||
|
|
for _, pid, camid in data:
|
||
|
|
pids.add(pid)
|
||
|
|
cams.add(camid)
|
||
|
|
return len(pids), len(cams)
|
||
|
|
|
||
|
|
def get_num_pids(self, data):
|
||
|
|
return self.parse_data(data)[0]
|
||
|
|
|
||
|
|
def get_num_cams(self, data):
|
||
|
|
return self.parse_data(data)[1]
|
||
|
|
|
||
|
|
def show_summary(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def combine_all(self):
|
||
|
|
"""Combines train, query and gallery"""
|
||
|
|
combined = copy.deepcopy(self.train)
|
||
|
|
|
||
|
|
# relabel pids in gallery
|
||
|
|
g_pids = set()
|
||
|
|
for _, pid, _ in self.gallery:
|
||
|
|
if pid==0 or pid==-1:
|
||
|
|
continue
|
||
|
|
g_pids.add(pid)
|
||
|
|
pid2label = {pid: i for i, pid in enumerate(g_pids)}
|
||
|
|
|
||
|
|
def _combine_data(data):
|
||
|
|
for img_path, pid, camid in data:
|
||
|
|
if pid==0 or pid==-1:
|
||
|
|
continue
|
||
|
|
pid = pid2label[pid] + self.num_train_pids
|
||
|
|
combined.append((img_path, pid, camid))
|
||
|
|
|
||
|
|
_combine_data(self.query)
|
||
|
|
_combine_data(self.gallery)
|
||
|
|
|
||
|
|
self.train = combined
|
||
|
|
self.num_train_pids = self.get_num_pids(self.train)
|
||
|
|
|
||
|
|
def download_dataset(self, dataset_dir, dataset_url):
|
||
|
|
"""Downloads and extracts dataset
|
||
|
|
|
||
|
|
:param dataset_dir: dataset directory
|
||
|
|
:type dataset_dir: str
|
||
|
|
:param dataset_url: url to download dataset
|
||
|
|
:type dataset_url: str
|
||
|
|
"""
|
||
|
|
if osp.exists(dataset_dir):
|
||
|
|
return
|
||
|
|
|
||
|
|
if dataset_url is None:
|
||
|
|
raise RuntimeError('{} dataset needs to be manually '
|
||
|
|
'prepared, please follow the '
|
||
|
|
'document to prepare this dataset'.format(self.__class__.__name__))
|
||
|
|
|
||
|
|
print('Creating directory "{}"'.format(dataset_dir))
|
||
|
|
mkdir_if_missing(dataset_dir)
|
||
|
|
fpath = osp.join(dataset_dir, osp.basename(dataset_url))
|
||
|
|
|
||
|
|
print('Downloading {} dataset to "{}"'.format(self.__class__.__name__, dataset_dir))
|
||
|
|
download_url(dataset_url, fpath)
|
||
|
|
|
||
|
|
print('Extracting "{}"'.format(fpath))
|
||
|
|
extension = osp.basename(fpath).split('.')[-1]
|
||
|
|
try:
|
||
|
|
tar = tarfile.open(fpath)
|
||
|
|
tar.extractall(path=dataset_dir)
|
||
|
|
tar.close()
|
||
|
|
except:
|
||
|
|
zip_ref = zipfile.ZipFile(fpath, 'r')
|
||
|
|
zip_ref.extractall(dataset_dir)
|
||
|
|
zip_ref.close()
|
||
|
|
|
||
|
|
print('{} dataset is ready'.format(self.__class__.__name__))
|
||
|
|
|
||
|
|
def check_before_run(self, required_files):
|
||
|
|
"""Checks if required files exist before going deeper
|
||
|
|
|
||
|
|
:param required_files: string name(s) of file(s)
|
||
|
|
:type required_files: str or list
|
||
|
|
"""
|
||
|
|
if isinstance(required_files, str):
|
||
|
|
required_files = [required_files]
|
||
|
|
|
||
|
|
for fpath in required_files:
|
||
|
|
if not osp.exists(fpath):
|
||
|
|
raise RuntimeError('"{}" is not found'.format(fpath))
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
num_train_pids, num_train_cams = self.parse_data(self.train)
|
||
|
|
num_query_pids, num_query_cams = self.parse_data(self.query)
|
||
|
|
num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)
|
||
|
|
|
||
|
|
msg = ' ----------------------------------------\n' \
|
||
|
|
' subset | # ids | # items | # cameras\n' \
|
||
|
|
' ----------------------------------------\n' \
|
||
|
|
' train | {:5d} | {:7d} | {:9d}\n' \
|
||
|
|
' query | {:5d} | {:7d} | {:9d}\n' \
|
||
|
|
' gallery | {:5d} | {:7d} | {:9d}\n' \
|
||
|
|
' ----------------------------------------\n' \
|
||
|
|
' items: images/tracklets for image/video dataset\n'.format(
|
||
|
|
num_train_pids, len(self.train), num_train_cams,
|
||
|
|
num_query_pids, len(self.query), num_query_cams,
|
||
|
|
num_gallery_pids, len(self.gallery), num_gallery_cams
|
||
|
|
)
|
||
|
|
|
||
|
|
return msg
|
||
|
|
|
||
|
|
|
||
|
|
class ImageDataset(Dataset):
|
||
|
|
|
||
|
|
def __init__(self, train, query, gallery, **kwargs):
|
||
|
|
super(ImageDataset, self).__init__(train, query, gallery, **kwargs)
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
img_path, pid, camid = self.data[index]
|
||
|
|
img = read_image(img_path)
|
||
|
|
if self.transform is not None:
|
||
|
|
img = self.transform(img)
|
||
|
|
return img, pid, camid, img_path
|
||
|
|
|
||
|
|
def show_summary(self):
|
||
|
|
num_train_pids, num_train_cams = self.parse_data(self.train)
|
||
|
|
num_query_pids, num_query_cams = self.parse_data(self.query)
|
||
|
|
num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)
|
||
|
|
|
||
|
|
print('=> Loaded {}'.format(self.__class__.__name__))
|
||
|
|
print(' ----------------------------------------')
|
||
|
|
print(' subset | # ids | # images | # cameras')
|
||
|
|
print(' ----------------------------------------')
|
||
|
|
print(' train | {:5d} | {:8d} | {:9d}'.format(num_train_pids, len(self.train), num_train_cams))
|
||
|
|
print(' query | {:5d} | {:8d} | {:9d}'.format(num_query_pids, len(self.query), num_query_cams))
|
||
|
|
print(' gallery | {:5d} | {:8d} | {:9d}'.format(num_gallery_pids, len(self.gallery), num_gallery_cams))
|
||
|
|
print(' ----------------------------------------')
|
||
|
|
|
||
|
|
|
||
|
|
class VideoDataset(Dataset):
|
||
|
|
|
||
|
|
def __init__(self, train, query, gallery, seq_len=15, sample_method='evenly', **kwargs):
|
||
|
|
super(VideoDataset, self).__init__(train, query, gallery, **kwargs)
|
||
|
|
self.seq_len = seq_len
|
||
|
|
self.sample_method = sample_method
|
||
|
|
|
||
|
|
if self.transform is None:
|
||
|
|
raise RuntimeError('transform must not be None')
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
img_paths, pid, camid = self.data[index]
|
||
|
|
num_imgs = len(img_paths)
|
||
|
|
|
||
|
|
if self.sample_method == 'random':
|
||
|
|
# Randomly samples seq_len images from a tracklet of length num_imgs,
|
||
|
|
# if num_imgs is smaller than seq_len, then replicates images
|
||
|
|
indices = np.arange(num_imgs)
|
||
|
|
replace = False if num_imgs>=self.seq_len else True
|
||
|
|
indices = np.random.choice(indices, size=self.seq_len, replace=replace)
|
||
|
|
# sort indices to keep temporal order (comment it to be order-agnostic)
|
||
|
|
indices = np.sort(indices)
|
||
|
|
|
||
|
|
elif self.sample_method == 'evenly':
|
||
|
|
# Evenly samples seq_len images from a tracklet
|
||
|
|
if num_imgs >= self.seq_len:
|
||
|
|
num_imgs -= num_imgs % self.seq_len
|
||
|
|
indices = np.arange(0, num_imgs, num_imgs/self.seq_len)
|
||
|
|
else:
|
||
|
|
# if num_imgs is smaller than seq_len, simply replicate the last image
|
||
|
|
# until the seq_len requirement is satisfied
|
||
|
|
indices = np.arange(0, num_imgs)
|
||
|
|
num_pads = self.seq_len - num_imgs
|
||
|
|
indices = np.concatenate([indices, np.ones(num_pads).astype(np.int32)*(num_imgs-1)])
|
||
|
|
assert len(indices) == self.seq_len
|
||
|
|
|
||
|
|
elif self.sample_method == 'all':
|
||
|
|
# Samples all images in a tracklet. batch_size must be set to 1
|
||
|
|
indices = np.arange(num_imgs)
|
||
|
|
|
||
|
|
else:
|
||
|
|
raise ValueError('Unknown sample method: {}'.format(self.sample_method))
|
||
|
|
|
||
|
|
imgs = []
|
||
|
|
for index in indices:
|
||
|
|
img_path = img_paths[int(index)]
|
||
|
|
img = read_image(img_path)
|
||
|
|
if self.transform is not None:
|
||
|
|
img = self.transform(img)
|
||
|
|
img = img.unsqueeze(0) # img must be torch.Tensor
|
||
|
|
imgs.append(img)
|
||
|
|
imgs = torch.cat(imgs, dim=0)
|
||
|
|
|
||
|
|
return imgs, pid, camid
|
||
|
|
|
||
|
|
def show_summary(self):
|
||
|
|
num_train_pids, num_train_cams = self.parse_data(self.train)
|
||
|
|
num_query_pids, num_query_cams = self.parse_data(self.query)
|
||
|
|
num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)
|
||
|
|
|
||
|
|
print('=> Loaded {}'.format(self.__class__.__name__))
|
||
|
|
print(' -------------------------------------------')
|
||
|
|
print(' subset | # ids | # tracklets | # cameras')
|
||
|
|
print(' -------------------------------------------')
|
||
|
|
print(' train | {:5d} | {:11d} | {:9d}'.format(num_train_pids, len(self.train), num_train_cams))
|
||
|
|
print(' query | {:5d} | {:11d} | {:9d}'.format(num_query_pids, len(self.query), num_query_cams))
|
||
|
|
print(' gallery | {:5d} | {:11d} | {:9d}'.format(num_gallery_pids, len(self.gallery), num_gallery_cams))
|
||
|
|
print(' -------------------------------------------')
|