def _cp_img_to(src, dst, rank, prefix):
"""
Args:
src: image path or tuple (for vidreid)
dst: target directory
rank: int, denoting ranked position, starting from 1
prefix: string
"""
if isinstance(src, tuple) or isinstance(src, list):
dst = osp.join(dst, prefix + '_top' + str(rank).zfill(3))
make_dirs(dst)
for img_path in src:
shutil.copy(img_path, dst)
else:
dst = osp.join(
dst, prefix + '_top' + str(rank).zfill(3) + '_name_' +
osp.basename(src))
shutil.copy(src, dst)
def __init__(self, config, loader):
self.config = config
self.loader = loader
# Model Config
self.mode = config.mode
self.cnnbackbone = config.cnnbackbone
self.pid_num = config.pid_num
self.t_margin = config.t_margin
# Logger Configuration
self.max_save_model_num = config.max_save_model_num
self.output_path = config.output_path
self.output_dirs_dict = {
'logs': os.path.join(self.output_path, 'logs/'),
'models': os.path.join(self.output_path, 'models/'),
'images': os.path.join(self.output_path, 'images/'),
'features': os.path.join(self.output_path, 'features/')
}
# make directions
for current_dir in self.output_dirs_dict.values():
make_dirs(current_dir)
# Train Configuration
# resume_train_dir
self.resume_train_dir = config.resume_train_dir
# init
self._init_device()
self._init_model()
self._init_criterion()
self._init_optimizer()
if config.output_featuremaps:
self._init_fixed_values()
if config.fp_16:
self.amp = amp
self.model_list = list(self.model_dict.values())
self.optimizer_list = list(self.optimizer_dict.values())
self.model_list, self.optimizer_list = self.amp.initialize(
self.model_list, self.optimizer_list, opt_level="O1")
else:
self.amp = None
self.model_list = None
self.optimizer_list = None
def visualize_ranked_results2(distmat, dataset, save_dir='', topk=20):
"""Visualizes ranked results.
Supports both image-reid and video-reid.
Args:
distmat (numpy.ndarray): distance matrix of shape (num_query, num_gallery).
dataset (tuple): a 2-tuple containing (query, gallery), each of which contains
tuples of (img_path(s), pid, camid).
save_dir (str): directory to save output images.
topk (int, optional): denoting top-k images in the rank list to be visualized.
"""
num_q, num_g = distmat.shape
print('Visualizing top-{} ranks'.format(topk))
print('# query: {}\n# gallery {}'.format(num_q, num_g))
print('Saving images to "{}"'.format(save_dir))
query, gallery = dataset
assert num_q == len(query)
assert num_g == len(gallery)
indices = np.argsort(distmat, axis=1)
make_dirs(save_dir)
def _cp_img_to(src, dst, rank, prefix):
"""
Args:
src: image path or tuple (for vidreid)
dst: target directory
rank: int, denoting ranked position, starting from 1
prefix: string
"""
if isinstance(src, tuple) or isinstance(src, list):
dst = osp.join(dst, prefix + '_top' + str(rank).zfill(3))
make_dirs(dst)
for img_path in src:
shutil.copy(img_path, dst)
else:
dst = osp.join(
dst, prefix + '_top' + str(rank).zfill(3) + '_name_' +
osp.basename(src))
shutil.copy(src, dst)
for q_idx in range(num_q):
qimg_path, qpid, qcamid = query[q_idx]
if isinstance(qimg_path, tuple) or isinstance(qimg_path, list):
qdir = osp.join(save_dir, osp.basename(qimg_path[0]))
else:
qdir = osp.join(save_dir, osp.basename(qimg_path))
make_dirs(qdir)
_cp_img_to(qimg_path, qdir, rank=0, prefix='query')
rank_idx = 1
for g_idx in indices[q_idx, :]:
gimg_path, gpid, gcamid = gallery[g_idx]
invalid = (qpid == gpid) & (qcamid == gcamid)
if not invalid:
_cp_img_to(gimg_path, qdir, rank=rank_idx, prefix='gallery')
rank_idx += 1
if rank_idx > topk:
break
print("Done")
def visualize_ranked_results(distmat,
dataset,
save_dir='',
topk=20,
sort='descend',
mode='inter-camera',
only_show=None):
"""Visualizes ranked results.
Args:
dismat (numpy.ndarray): distance matrix of shape (nq, ng)
dataset (tupple): a 2-tuple including (query,gallery), each of which contains
tuples of (img_paths, pids, camids)
save_dir (str): directory to save output images.
topk (int, optional): denoting top-k images in the rank list to be visualized.
sort (string): ascend means small value is similar, otherwise descend
mode (string): intra-camera/inter-camera/all
intra-camera only visualize results in the same camera with the query
inter-camera only visualize results in the different camera with the query
all visualize all results
"""
num_q, num_g = distmat.shape
print('Visualizing top-{} ranks'.format(topk))
print('# query: {}\n# gallery {}'.format(num_q, num_g))
print('Saving images to "{}"'.format(save_dir))
query, gallery = dataset
assert num_q == len(query)
assert num_g == len(gallery)
assert sort in ['descend', 'ascend']
assert mode in ['intra-camera', 'inter-camera', 'all']
if sort is 'ascend':
indices = np.argsort(distmat, axis=1)
elif sort is 'descend':
indices = np.argsort(distmat, axis=1)[:, ::-1]
make_dirs(save_dir)
def cat_imgs_to(image_list, hit_list, text_list, target_dir):
images = []
for img, hit, text in zip(image_list, hit_list, text_list):
img = Image.open(img).resize((64, 128))
d = ImageDraw.Draw(img)
d.text((3, 1), "{:.3}".format(text), fill=(255, 255, 0))
if hit:
img = ImageOps.expand(img, border=4, fill='green')
else:
img = ImageOps.expand(img, border=4, fill='red')
images.append(img)
widths, heights = zip(*(i.size for i in images))
total_width = sum(widths)
max_height = max(heights)
new_im = Image.new('RGB', (total_width, max_height))
x_offset = 0
for im in images:
new_im.paste(im, (x_offset, 0))
x_offset += im.size[0]
new_im.save(target_dir)
counts = 0
for q_idx in range(num_q):
image_list = []
hit_list = []
text_list = []
# query image
qimg_path, qpid, qcamid = query[q_idx]
image_list.append(qimg_path)
hit_list.append(True)
text_list.append(0.0)
# target dir
if isinstance(qimg_path, tuple) or isinstance(qimg_path, list):
qdir = osp.join(save_dir, osp.basename(qimg_path[0]))
else:
qdir = osp.join(save_dir, osp.basename(qimg_path))
# matched images
rank_idx = 1
for ii, g_idx in enumerate(indices[q_idx, :]):
gimg_path, gpid, gcamid = gallery[g_idx]
if mode == 'intra-camera':
valid = qcamid == gcamid
elif mode == 'inter-camera':
valid = (qpid != gpid
and qcamid == gcamid) or (qcamid != gcamid)
elif mode == 'all':
valid = True
if valid:
if only_show == 'pos' and qpid != gpid: continue
if only_show == 'neg' and qpid == gpid: continue
image_list.append(gimg_path)
hit_list.append(qpid == gpid)
text_list.append(distmat[q_idx, g_idx])
rank_idx += 1
if rank_idx > topk:
break
counts += 1
cat_imgs_to(image_list, hit_list, text_list, qdir)
print(counts, qdir)
def featuremaps2heatmaps(self,
original_images,
featuremaps,
image_paths,
current_epoch,
if_save=False,
if_fixed=False,
if_fake=False):
height = original_images.size(2)
width = original_images.size(3)
imgs = original_images
outputs = featuremaps.sum(1)
# outputs = (outputs ** 2).sum(1)
#b, h, w = outputs.size()
#outputs = outputs.view(b, h * w)
# outputs = F.normalize(outputs, p=2, dim=1)
#outputs = outputs.view(b, h, w)
imgs, outputs = imgs.cpu(), outputs.cpu()
grid_img_tensor = []
if if_save:
save_dir = osp.join(self.output_dirs_dict['images'],
str(current_epoch))
make_dirs(save_dir)
if if_fixed:
if if_fake:
save_dir = osp.join(self.output_dirs_dict['images'],
str(current_epoch), 'fake')
make_dirs(save_dir)
else:
save_dir = osp.join(self.output_dirs_dict['images'],
str(current_epoch), 'true')
make_dirs(save_dir)
else:
save_dir = osp.join(self.output_dirs_dict['images'],
str(current_epoch))
for j in range(outputs.size(0)):
# get image name
path = image_paths[j]
imname = osp.basename(osp.splitext(path)[0])
# RGB image
img = imgs[j, ...]
for t, m, s in zip(img, IMAGENET_MEAN, IMAGENET_STD):
t.mul_(s).add_(m).clamp_(0, 1)
img_np = np.uint8(np.floor(img.numpy() * 255))
img_np = img_np.transpose((1, 2, 0)) # (c, h, w) -> (h, w, c)
# activation map
am = outputs[j, ...].numpy()
am = cv2.resize(am, (width, height))
am = 255 * (am - np.min(am)) / (np.max(am) - np.min(am) + 1e-12)
am = np.uint8(np.floor(am))
am = cv2.applyColorMap(am, cv2.COLORMAP_JET)
# overlapped
overlapped = img_np * 0.3 + am * 0.7
overlapped[overlapped > 255] = 255
overlapped = overlapped.astype(np.uint8)
# save images in a single figure (add white spacing between images)
# from left to right: original image, activation map, overlapped image
grid_img = 255 * np.ones(
(height, 3 * width + 2 * GRID_SPACING, 3), dtype=np.uint8)
grid_img[:, :width, :] = img_np[:, :, ::-1]
grid_img[:, width + GRID_SPACING:2 * width + GRID_SPACING, :] = am
grid_img[:, 2 * width + 2 * GRID_SPACING:, :] = overlapped
grid_img_tensor.append(grid_img)
if if_save:
cv2.imwrite(osp.join(save_dir, imname + '.jpg'), grid_img)
grid_img_tensor = np.transpose(np.stack(grid_img_tensor, axis=0),
(0, 3, 1, 2))
return torch.from_numpy(grid_img_tensor)
def save_model(self, save_step, save_epoch):
'''save model as save_epoch'''
# save model
models_steps_path = os.path.join(self.output_dirs_dict['models'],
str(save_step))
if not osp.exists(models_steps_path):
make_dirs(models_steps_path)
for module_name, module in self.model_dict.items():
torch.save(
module.state_dict(),
os.path.join(models_steps_path,
f'model_{module_name}_{save_epoch}.pkl'))
for optimizer_name, optimizer in self.optimizer_dict.items():
torch.save(
optimizer.state_dict(),
os.path.join(models_steps_path,
f'optimizer_{optimizer_name}_{save_epoch}.pkl'))
if self.config.fp_16 and self.amp:
torch.save(
self.amp.state_dict(),
os.path.join(models_steps_path, f'amp_{save_epoch}.pkl'))
# if saved model is more than max num, delete the model with smallest epoch
if self.max_save_model_num > 0:
root, _, files = os_walk(models_steps_path)
# get indexes of saved models
indexes = []
for file in files:
indexes.append(int(file.replace('.pkl', '').split('_')[-1]))
# remove the bad-case and get available indexes
model_num = len(self.model_dict)
optimizer_num = len(self.optimizer_dict)
available_indexes = copy.deepcopy(indexes)
for element in indexes:
if indexes.count(element) < model_num + optimizer_num:
available_indexes.remove(element)
available_indexes = sorted(list(set(available_indexes)),
reverse=True)
unavailable_indexes = list(
set(indexes).difference(set(available_indexes)))
# delete all unavailable models
for unavailable_index in unavailable_indexes:
try:
# os.system('find . -name "{}*_{}.pth" | xargs rm -rf'.format(self.config.save_models_path, unavailable_index))
for module_name in self.model_dict.keys():
os.remove(
os.path.join(
root,
f'model_{module_name}_{unavailable_index}.pkl')
)
for optimizer_name in self.optimizer_dict.keys():
os.remove(
os.path.join(
root,
f'optimizer_{optimizer_name}_{unavailable_index}.pkl'
))
if self.config.fp_16 and self.amp:
os.remove(
os.path.join(root, f'amp_{unavailable_index}.pkl'))
except:
pass
# delete extra models
if len(available_indexes) >= self.max_save_model_num:
for extra_available_index in available_indexes[
self.max_save_model_num:]:
# os.system('find . -name "{}*_{}.pth" | xargs rm -rf'.format(self.config.save_models_path, extra_available_index))
for mudule_name, mudule in self.model_dict.items():
os.remove(
os.path.join(
root,
f'model_{mudule_name}_{extra_available_index}.pkl'
))
for optimizer_name, optimizer in self.optimizer_dict.items(
):
os.remove(
os.path.join(
root,
f'optimizer_{optimizer_name}_{extra_available_index}.pkl'
))
if self.config.fp_16 and self.amp:
os.remove(
os.path.join(root,
f'amp_{extra_available_index}.pkl'))