def evaluate(dataloader,
detector,
cfg_maskrcnn,
retrievalnet,
opt,
cfg_eval,
cfg_visualization,
is_cuda=False,
logger_prefix="detector-retrieval"):
logger = logging.getLogger(f"{logger_prefix}.evaluate")
dataset_name = dataloader.get_name()
dataset_scale = dataloader.get_eval_scale()
logger.info("Starting to eval on {0}, scale {1}".format(
dataset_name, dataset_scale))
t_start_eval = time.time()
detector.eval()
retrievalnet.eval()
## setup retrievalnet
# setting up the multi-scale parameters
ms = [1]
msp = 1
if opt.retrieval_multiscale:
ms = [1, 1. / math.sqrt(2), 1. / 2]
if retrievalnet.meta[
"pooling"] == "gem" and retrievalnet.whiten is None:
msp = retrievalnet.pool.p.data.tolist()[0]
#setup whitening
if opt.retrieval_whitening_path is not None:
logger.info("Whitening is precomputed, loading it from {0}".format(
opt.retrieval_whitening_path))
whitening_data = torch.load(opt.retrieval_whitening_path)
if ( (opt.retrieval_multiscale and "ms" in whitening_data) or \
(not opt.retrieval_multiscale and "ss" in whitening_data ) ):
if opt.retrieval_multiscale:
Lw = copy.deepcopy(whitening_data["ms"])
else:
Lw = copy.deepcopy(whitening_data["ss"])
else:
raise RuntimeError(
"Whitening should be precomputed with the network")
# convert whitening data to torch tensors
Lw["m"], Lw["P"] = torch.from_numpy(Lw["m"]), torch.from_numpy(Lw["P"])
if is_cuda:
Lw["m"], Lw["P"] = Lw["m"].cuda(), Lw["P"].cuda()
else:
Lw = None
with torch.no_grad(
): # do evaluation in forward mode only (for speed and memory)
# extract features from query images
query_images, _, _ = dataloader.get_all_class_images(do_resize=False)
if is_cuda:
query_images = [img.cuda() for img in query_images]
query_images = [img[0] for img in query_images
] # get rid of the batch dimension
query_images = [
resize_image_tensor(img, opt.retrieval_image_size)
for img in query_images
]
query_images = [dataloader.unnorm_image(img) for img in query_images]
query_images_with_aug = []
for im in query_images:
query_images_with_aug.append(im)
if not cfg_eval.class_image_augmentation:
num_class_views = 1
elif cfg_eval.class_image_augmentation == "rotation90":
im90 = im.rot90(1, [1, 2])
im180 = im90.rot90(1, [1, 2])
im270 = im180.rot90(1, [1, 2])
query_images_with_aug.append(im90)
query_images_with_aug.append(im180)
query_images_with_aug.append(im270)
num_class_views = 4
elif cfg_eval.class_image_augmentation == "horflip":
im_flipped = im.flip(2)
query_images_with_aug.append(im_flipped)
num_class_views = 2
else:
raise RuntimeError(
f"Unknown value of class_image_augmentation: {cfg_eval.class_image_augmentation}"
)
query_images = query_images_with_aug
query_vectors = extract_vectors_from_images(retrievalnet,
query_images,
ms=ms,
msp=msp)
# apply whitening if defined
if Lw is not None:
query_vectors = whitenapply(query_vectors, Lw["m"], Lw["P"])
query_vectors = torch.transpose(query_vectors, 0, 1)
# prepare looping over all iamges
iterator = make_iterator_extract_scores_from_images_batched(
dataloader,
detector,
cfg_maskrcnn,
logger,
image_batch_size=cfg_eval.batch_size,
is_cuda=is_cuda)
boxes, labels, scores = [], [], []
gt_boxes = []
image_ids = []
losses = OrderedDict()
# loop over all dataset images
num_evaluted_images = 0
for data in iterator:
image_id, boxes_one_image, image_pyramid, query_img_sizes, class_ids, initial_img_size = data
image_ids.append(image_id)
logger.info(f"Image {num_evaluted_images}: id {image_id}")
num_evaluted_images += 1
img_size_pyramid = [
FeatureMapSize(img=img) for img in image_pyramid
]
gt_boxes_one_image = dataloader.get_image_annotation_for_imageid(
image_id)
gt_boxes.append(gt_boxes_one_image)
# vizualize GT for debug
if cfg_visualization.show_gt_boxes:
visualizer.show_gt_boxes(image_id, gt_boxes_one_image,
class_ids, dataloader)
# decode image predictions
# merge boxes_one_image, labels_one_image, scores_one_image from different pyramid layers
boxes_one_image = cat_boxlist(boxes_one_image)
# do NMS
good_indices = nms(
boxes_one_image,
opt.nms_iou_threshold_detector_score,
nms_score_threshold=opt.nms_score_threshold_detector_score)
boxes_one_image = boxes_one_image[good_indices]
# extract feature vectors from the predictions
image_original = dataloader._transform_image(image_id,
do_augmentation=True,
hflip=False,
vflip=False)[0]
if is_cuda:
image_original = image_original.cuda()
image_patches = crop_resize_image_patches(
image_original,
boxes_one_image,
opt.retrieval_image_size,
logger,
unnorm_image=dataloader.unnorm_image,
is_cuda=is_cuda)
# filter out cases when failed to crop a box: outside of the image
good_indices = [
i for i, p in enumerate(image_patches) if p is not None
]
if good_indices:
# non empty
image_patches = [p for p in image_patches if p is not None]
boxes_one_image = boxes_one_image[good_indices]
image_vectors = extract_vectors_from_images(retrievalnet,
image_patches,
ms=ms,
msp=msp)
# compute class scores from image_vectors and query_vectors (already transposed)
if Lw is not None:
# apply whitening if defined
image_vectors = whitenapply(image_vectors, Lw["m"],
Lw["P"])
scores_retrieval = torch.mm(query_vectors, image_vectors)
num_queries = scores_retrieval.size(0)
num_detections = scores_retrieval.size(1)
list_of_active_label = torch.LongTensor(class_ids)
if cfg_eval.class_image_augmentation:
list_of_active_label = torch.stack(
[list_of_active_label] * num_class_views, 1).view(-1)
# take all labels for all boxes - will sort them by scores at eval
scores_one_image = scores_retrieval.view(-1)
boxes_one_image = cat_boxlist([boxes_one_image] * num_queries)
labels_one_image = torch.stack([list_of_active_label] *
num_detections,
1).contiguous().view(-1)
# add scores and labels: overwrite if existed
boxes_one_image.add_field("labels", labels_one_image)
boxes_one_image.add_field("scores", scores_one_image)
# NMS using the retrieval scores
good_indices = nms(
boxes_one_image,
cfg_eval.nms_iou_threshold,
nms_score_threshold=cfg_eval.nms_score_threshold,
do_separate_per_label=not cfg_eval.nms_across_classes)
boxes_one_image = boxes_one_image[good_indices]
else:
boxes_one_image.add_field(
"labels",
torch.zeros(0,
dtype=torch.long,
device=boxes_one_image.bbox_xyxy.device))
boxes_one_image.add_field(
"scores",
torch.zeros(0,
dtype=torch.float,
device=boxes_one_image.bbox_xyxy.device))
boxes.append(boxes_one_image.cpu())
if cfg_visualization.show_detections:
# do not pass class_ids - this is already taken care of
visualizer.show_detections(boxes_one_image,
image_id,
dataloader,
cfg_visualization,
class_ids=None)
# normalize by number of steps
for k in losses:
losses[k] /= num_evaluted_images
# Save detection if requested
if cfg_visualization.path_to_save_detections:
data = {
"image_ids": image_ids,
"boxes_xyxy": [bb.bbox_xyxy for bb in boxes],
"labels": [bb.get_field("labels") for bb in boxes],
"scores": [bb.get_field("scores") for bb in boxes],
"gt_boxes_xyxy": [bb.bbox_xyxy for bb in gt_boxes],
"gt_labels": [bb.get_field("labels") for bb in gt_boxes],
"gt_difficults": [bb.get_field("difficult") for bb in gt_boxes]
}
dataset_name = dataloader.get_name()
os.makedirs(cfg_visualization.path_to_save_detections, exist_ok=True)
save_path = os.path.join(cfg_visualization.path_to_save_detections,
dataset_name + "_detections.pth")
torch.save(data, save_path)
# compute mAP
for mAP_iou_threshold in cfg_eval.mAP_iou_thresholds:
logger.info("Evaluating at IoU th {:0.2f}".format(mAP_iou_threshold))
ap_data = do_voc_evaluation(boxes,
gt_boxes,
iou_thresh=mAP_iou_threshold,
use_07_metric=False)
losses["mAP@{:0.2f}".format(mAP_iou_threshold)] = ap_data["map"]
losses["mAPw@{:0.2f}".format(
mAP_iou_threshold)] = ap_data["map_weighted"]
losses["recall@{:0.2f}".format(mAP_iou_threshold)] = ap_data["recall"]
# per class AP information
for i_class, (ap, recall, n_pos) in enumerate(
zip(ap_data["ap_per_class"], ap_data["recall_per_class"],
ap_data["n_pos"])):
if not np.isnan(ap):
assert i_class in class_ids, "Could not find class_id in the list of ids"
logger.info(
"Class {0} (local {3}), AP {1:0.4f}, #obj {2}, recall {4:0.4f}"
.format(i_class, ap, n_pos, class_ids.index(i_class),
recall))
# save timing
losses["eval_time"] = (time.time() - t_start_eval)
logger.info("Evaluated on {0}, scale {1}".format(dataset_name,
dataset_scale))
print_meters(losses, logger)
return losses
def trainval_loop(dataloader_train, net, cfg, criterion, optimizer, dataloaders_eval=[]):
"""Main train+val loop
Args:
dataloader_train -dataloader to get training batches
net - the network to use
cfg - config with all the parameters
criterion - criterion to optimize
optimizer - optimization to use
dataloaders_eval - a list of dataloaders to use for validation, at each validation stage all of them will be used sequentially
Returns nothing
"""
# init plotting and logging
logger = logging.getLogger("OS2D.train")
t_start = time.time()
num_steps_for_logging, meters_running = 0, {}
full_log = init_log()
if cfg.train.optim.max_iter > 0 and cfg.train.do_training:
logger.info("Start training")
# setup the learning rate schedule
_, anneal_lr_func = setup_lr(optimizer, full_log, cfg.train.optim.anneal_lr, cfg.eval.iter)
# evaluate the initial model
meters_eval = evaluate_model(dataloaders_eval, net, cfg, criterion)
if cfg.output.best_model.do_get_best_model:
assert (cfg.output.best_model.dataset and cfg.output.best_model.dataset in meters_eval) \
or (len(cfg.eval.dataset_names) > 0 and cfg.eval.dataset_names[0] in meters_eval), \
"Cannot determine which dataset to use for the best model"
best_model_dataset_name = cfg.output.best_model.dataset if cfg.output.best_model.dataset else cfg.eval.dataset_names[0]
best_model_metric = meters_eval[best_model_dataset_name][cfg.output.best_model.metric]
logger.info(f"Init model is the current best on {best_model_dataset_name} by {cfg.output.best_model.metric}, value {best_model_metric:.4f}")
if cfg.output.path:
checkpoint_best_model_name = f"best_model_{best_model_dataset_name}_{cfg.output.best_model.metric}"
checkpoint_best_model_path = \
checkpoint_model(net, optimizer, cfg.output.path, cfg.is_cuda, model_name=checkpoint_best_model_name,
extra_fields={"criterion_dataset": best_model_dataset_name,
"criterion_metric": cfg.output.best_model.metric,
"criterion_mode": cfg.output.best_model.mode,
"criterion_value": best_model_metric,
"criterion_value_old": None})
else:
raise RuntimeError("cfg.output.best_model.do_get_best_model i set to True, but cfg.output.path is not provided, so cannot save best models")
if cfg.train.optim.anneal_lr.reload_best_model_after_anneal_lr and\
cfg.train.optim.anneal_lr.type != "none":
assert cfg.output.best_model.do_get_best_model, "cfg.train.optim.anneal_lr.reload_best_model_after_anneal_lr was set to True, but cfg.output.best_model.do_get_best_model is False, so there is no best model to reload from"
# add the initial point
log_meters(full_log, t_start, -1, cfg.output.path,
meters_eval=meters_eval,
anneal_lr=anneal_lr_func)
# save initial model
if cfg.output.path:
checkpoint_model(net, optimizer, cfg.output.path, cfg.is_cuda, i_iter=0)
# start training
i_epoch = 0
i_batch = len(dataloader_train) # to start a new epoch at the first iteration
for i_iter in range(cfg.train.optim.max_iter):
# restart dataloader if needed
if i_batch >= len(dataloader_train):
i_epoch += 1
i_batch = 0
# shuffle dataset
dataloader_train.shuffle()
# mine hard negative classes
if cfg.train.mining.do_mining and i_iter % cfg.train.mining.mine_hard_patches_iter == 0:
hardnegdata_per_imageid = mine_hard_patches(dataloader_train, net, cfg, criterion)
dataloader_train.set_hard_negative_data(hardnegdata_per_imageid)
# print iter info
logger.info(f"Iter {i_iter} ({cfg.train.optim.max_iter}), epoch {i_epoch}, time {time_since(t_start)}")
# get data for training
t_start_loading = time.time()
batch_data = dataloader_train.get_batch(i_batch)
t_data_loading = time.time() - t_start_loading
i_batch += 1
num_steps_for_logging += 1
# train on one batch
meters = train_one_batch(batch_data, net, cfg, criterion, optimizer, dataloader_train, logger)
meters["loading_time"] = t_data_loading
# print meters
if i_iter % cfg.output.print_iter == 0:
print_meters(meters, logger)
# update logs
add_to_meters_in_dict(meters, meters_running)
# evaluation
if (i_iter + 1) % cfg.eval.iter == 0:
meters_eval = evaluate_model(dataloaders_eval, net, cfg, criterion)
# checkpoint the best model
if cfg.output.best_model.do_get_best_model:
cur_metric = meters_eval[best_model_dataset_name][cfg.output.best_model.metric]
assert cfg.output.best_model.mode in ["max", "min"], f"cfg.output.best_model.mode should be 'max' or 'min', but have {cfg.output.best_model.mode}"
if (cfg.output.best_model.mode=="max" and cur_metric > best_model_metric) or \
(cfg.output.best_model.mode=="min" and cur_metric < best_model_metric):
# overwrite the best model
logger.info(f"New best model on {best_model_dataset_name} by {cfg.output.best_model.metric}, value {cur_metric:.4f}")
if cfg.output.path:
checkpoint_best_model_path = \
checkpoint_model(net, optimizer, cfg.output.path, cfg.is_cuda, model_name=checkpoint_best_model_name,
extra_fields={"criterion_dataset": best_model_dataset_name,
"criterion_metric": cfg.output.best_model.metric,
"criterion_mode": cfg.output.best_model.mode,
"criterion_value": cur_metric,
"criterion_value_old": best_model_metric})
best_model_metric = cur_metric
# normalize by number of steps
for k in meters_running:
meters_running[k] /= num_steps_for_logging
# anneal learning rate
meters_running["lr"] = get_learning_rate(optimizer)
if anneal_lr_func:
lr = anneal_lr_func(i_iter + 1, anneal_now=i_iter > cfg.train.optim.anneal_lr.initial_patience)
flag_changed_lr = lr != meters_running["lr"]
else:
lr = meters_running["lr"]
flag_changed_lr = False
# if lr was annealed load the best up to now model and set it up
if cfg.train.optim.anneal_lr.reload_best_model_after_anneal_lr and flag_changed_lr:
if cfg.output.best_model.do_get_best_model: # if have the best model at all
optimizer_state = net.init_model_from_file(checkpoint_best_model_path)
if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state)
set_learning_rate(optimizer, lr)
# eval and log
log_meters(full_log, t_start, i_iter, cfg.output.path,
meters_running=meters_running,
meters_eval=meters_eval)
# init for the next
num_steps_for_logging, meters_running = 0, {}
# save intermediate model
if cfg.output.path and cfg.output.save_iter and i_iter % cfg.output.save_iter == 0:
checkpoint_model(net, optimizer, cfg.output.path, cfg.is_cuda, i_iter=i_iter)
# evaluate the final model
logger.info("Final evaluation")
meters_eval = evaluate_model(dataloaders_eval, net, cfg, criterion, print_per_class_results=True)
# add the final point
if cfg.train.optim.max_iter > 0 and cfg.train.do_training:
log_meters(full_log, t_start, cfg.train.optim.max_iter, cfg.output.path,
meters_eval=meters_eval)
# save the final model
if cfg.output.path:
checkpoint_model(net, optimizer, cfg.output.path, cfg.is_cuda, i_iter=cfg.train.optim.max_iter)
def evaluate(dataloader, net, cfg, criterion=None, print_per_class_results=False):
"""
Evaluation of the provided model at one dataset
Args:
dataloader - the dataloader to get data
net - the network to use
cfg - config with all the parameters
criterion - criterion (usually the same one as used for training), can be None, will just not compute related metrics
print_per_class_results - flag showing whether to printout extra data (per class AP) - usually used at the final evaluation
Returns:
losses (OrderedDict) - all computed metrics, e.g., losses["[email protected]"] - mAP at IoU threshold 0.5
"""
logger = logging.getLogger("OS2D.evaluate")
dataset_name = dataloader.get_name()
dataset_scale = dataloader.get_eval_scale()
logger.info("Starting to eval on {0}, scale {1}".format(dataset_name, dataset_scale))
t_start_eval = time.time()
net.eval()
iterator = make_iterator_extract_scores_from_images_batched(dataloader, net, logger,
image_batch_size=cfg.eval.batch_size,
is_cuda=cfg.is_cuda,
class_image_augmentation=cfg.eval.class_image_augmentation)
boxes = []
gt_boxes = []
losses = OrderedDict()
image_ids = []
# loop over all dataset images
num_evaluted_images = 0
for data in iterator:
image_id, image_loc_scores_pyramid, image_class_scores_pyramid,\
image_pyramid, query_img_sizes, class_ids,\
box_reverse_transform, image_fm_sizes_p, transform_corners_pyramid\
= data
image_ids.append(image_id)
num_evaluted_images += 1
img_size_pyramid = [FeatureMapSize(img=img) for img in image_pyramid]
num_labels = len(class_ids)
gt_boxes_one_image = dataloader.get_image_annotation_for_imageid(image_id)
gt_boxes.append(gt_boxes_one_image)
# compute losses
if len(gt_boxes_one_image) > 0:
# there is some annotation for this image
gt_labels_one_image = gt_boxes_one_image.get_field("labels")
dataloader.update_box_labels_to_local(gt_boxes_one_image, class_ids)
loc_targets_pyramid, class_targets_pyramid = \
dataloader.box_coder.encode_pyramid(gt_boxes_one_image,
img_size_pyramid, num_labels,
default_box_transform_pyramid=box_reverse_transform)
# return the original labels back
gt_boxes_one_image.add_field("labels", gt_labels_one_image)
# vizualize GT for debug
if cfg.visualization.eval.show_gt_boxes:
visualizer.show_gt_boxes(image_id, gt_boxes_one_image, class_ids, dataloader)
if cfg.is_cuda:
loc_targets_pyramid = [loc_targets.cuda() for loc_targets in loc_targets_pyramid]
class_targets_pyramid = [class_targets.cuda() for class_targets in class_targets_pyramid]
transform_corners_pyramid = [transform_corners.cuda() for transform_corners in transform_corners_pyramid]
add_batch_dim = lambda list_of_tensors: [t.unsqueeze(0) for t in list_of_tensors]
if criterion is not None:
# if criterion is provided, use it to compute all metrics it can
losses_iter = criterion(add_batch_dim(image_loc_scores_pyramid) if image_loc_scores_pyramid[0] is not None else None,
add_batch_dim(loc_targets_pyramid),
add_batch_dim(image_class_scores_pyramid),
add_batch_dim(class_targets_pyramid)
)
# convert to floats
for l in losses_iter:
losses_iter[l] = losses_iter[l].mean().item()
# printing
print_meters(losses_iter, logger)
# update logs
add_to_meters_in_dict(losses_iter, losses)
# decode image predictions
boxes_one_image = \
dataloader.box_coder.decode_pyramid(image_loc_scores_pyramid, image_class_scores_pyramid,
img_size_pyramid, class_ids,
nms_iou_threshold=cfg.eval.nms_iou_threshold,
nms_score_threshold=cfg.eval.nms_score_threshold,
inverse_box_transforms=box_reverse_transform,
transform_corners_pyramid=transform_corners_pyramid)
boxes.append(boxes_one_image.cpu())
if cfg.visualization.eval.show_detections:
visualizer.show_detection_from_dataloader(boxes_one_image, image_id, dataloader, cfg.visualization.eval, class_ids=None)
if cfg.visualization.eval.show_class_heatmaps:
visualizer.show_class_heatmaps(image_id, class_ids, image_fm_sizes_p, class_targets_pyramid, image_class_scores_pyramid,
cfg_local=cfg.visualization.eval,
class_image_augmentation=cfg.eval.class_image_augmentation)
if cfg.is_cuda:
torch.cuda.empty_cache()
# normalize by number of steps
for k in losses:
losses[k] /= num_evaluted_images
# Save detection if requested
path_to_save_detections = cfg.visualization.eval.path_to_save_detections
if path_to_save_detections:
data = {"image_ids" : image_ids,
"boxes_xyxy" : [bb.bbox_xyxy for bb in boxes],
"labels" : [bb.get_field("labels") for bb in boxes],
"scores" : [bb.get_field("scores") for bb in boxes],
"gt_boxes_xyxy" : [bb.bbox_xyxy for bb in gt_boxes],
"gt_labels" : [bb.get_field("labels") for bb in gt_boxes],
"gt_difficults" : [bb.get_field("difficult") for bb in gt_boxes]
}
dataset_name = dataloader.get_name()
os.makedirs(path_to_save_detections, exist_ok=True)
save_path = os.path.join(path_to_save_detections, dataset_name + "_detections.pth")
torch.save(data, save_path)
# compute mAP
for mAP_iou_threshold in cfg.eval.mAP_iou_thresholds:
logger.info("Evaluating at IoU th {:0.2f}".format(mAP_iou_threshold))
ap_data = do_voc_evaluation(boxes, gt_boxes, iou_thresh=mAP_iou_threshold, use_07_metric=False)
losses["mAP@{:0.2f}".format(mAP_iou_threshold)] = ap_data["map"]
losses["mAPw@{:0.2f}".format(mAP_iou_threshold)] = ap_data["map_weighted"]
losses["recall@{:0.2f}".format(mAP_iou_threshold)] = ap_data["recall"]
if print_per_class_results:
# per class AP information
for i_class, (ap, recall, n_pos) in enumerate(zip(ap_data["ap_per_class"], ap_data["recall_per_class"], ap_data["n_pos"])):
if not np.isnan(ap):
assert i_class in class_ids, "Could not find class_id in the list of ids"
logger.info("Class {0} (local {3}), AP {1:0.4f}, #obj {2}, recall {4:0.4f}".format(i_class,
ap,
n_pos,
class_ids.index(i_class),
recall))
# save timing
losses["eval_time"] = (time.time() - t_start_eval)
logger.info("Evaluated on {0}, scale {1}".format(dataset_name, dataset_scale))
print_meters(losses, logger)
return losses
def mine_hard_patches(dataloader, net, cfg, criterion):
"""Mine patches that are hard: classification false positives and negative, localization errors
At each level of sampled image pyramid, we need to cut out a piece of size appropriate for training
(levels are defined by cfg.train.mining.num_random_pyramid_scales, cfg.train.mining.num_random_negative_classes)
Args:
dataloader - dataloader to use (often the same as the one for training)
net - the network to use
cfg - config with all the parameters
criterion - criterion (usually the same one as used for training)
Returns:
hardnegdata_per_imageid (OrderedDict) - mined data, keys are the image ids;
further used in dataloader.set_hard_negative_data(hardnegdata_per_imageid) when preparing batches
"""
logger = logging.getLogger("OS2D.mining_hard_patches")
logger.info("Starting to mine hard patches")
t_start_mining = time.time()
net.eval()
num_batches = len(dataloader)
hardnegdata_per_imageid = OrderedDict()
iterator = make_iterator_extract_scores_from_images_batched(dataloader, net, logger,
image_batch_size=cfg.eval.batch_size,
is_cuda=cfg.is_cuda,
num_random_pyramid_scales=cfg.train.mining.num_random_pyramid_scales,
num_random_negative_labels=cfg.train.mining.num_random_negative_classes)
boxes = []
gt_boxes = []
losses = OrderedDict()
# loop over all dataset images
for data in iterator:
t_item_start = time.time()
image_id, image_loc_scores_pyramid, image_class_scores_pyramid, \
image_pyramid, query_img_sizes, \
batch_class_ids, box_reverse_transform_pyramid, image_fm_sizes_p, transform_corners_pyramid \
= data
img_size_pyramid = [FeatureMapSize(img=image) for image in image_pyramid]
gt_boxes_one_image = dataloader.get_image_annotation_for_imageid(image_id)
gt_boxes.append(gt_boxes_one_image)
# compute losses
# change labels to the ones local to the current image
dataloader.update_box_labels_to_local(gt_boxes_one_image, batch_class_ids)
num_labels = len(batch_class_ids)
loc_targets_pyramid, class_targets_pyramid = \
dataloader.box_coder.encode_pyramid(gt_boxes_one_image, img_size_pyramid, num_labels,
default_box_transform_pyramid=box_reverse_transform_pyramid)
# vizualize GT for debug
if cfg.visualization.mining.show_gt_boxes:
visualizer.show_gt_boxes(image_id, gt_boxes_one_image, batch_class_ids, dataloader)
# compute losses
if cfg.is_cuda:
loc_targets_pyramid = [loc_targets.cuda() for loc_targets in loc_targets_pyramid]
class_targets_pyramid = [class_targets.cuda() for class_targets in class_targets_pyramid]
add_batch_dim = lambda list_of_tensors: [t.unsqueeze(0) for t in list_of_tensors]
loc_scores_pyramid = add_batch_dim(image_loc_scores_pyramid)
cls_targets_remapped_pyramid = []
for loc_scores, img_size, box_reverse_transform in zip(loc_scores_pyramid, img_size_pyramid, box_reverse_transform_pyramid):
# loop over the pyramid levels
cls_targets_remapped, ious_anchor, ious_anchor_corrected = \
dataloader.box_coder.remap_anchor_targets(loc_scores, [img_size], query_img_sizes, [gt_boxes_one_image],
box_reverse_transform=[box_reverse_transform])
cls_targets_remapped_pyramid.append(cls_targets_remapped)
losses_iter, losses_per_anchor = criterion(loc_scores_pyramid,
add_batch_dim(loc_targets_pyramid),
add_batch_dim(image_class_scores_pyramid),
add_batch_dim(class_targets_pyramid),
cls_targets_remapped=cls_targets_remapped_pyramid,
patch_mining_mode=True)
if cfg.visualization.mining.show_class_heatmaps:
visualizer.show_class_heatmaps(image_id, batch_class_ids, image_fm_sizes_p, class_targets_pyramid, image_class_scores_pyramid,
cfg_local=cfg.visualization.mining)
assert dataloader.data_augmentation is not None, "Can mine hard patches only through data augmentation"
crop_size = dataloader.data_augmentation.random_crop_size
# convert to floats
for l in losses_iter:
losses_iter[l] = losses_iter[l].mean().item()
# printing
print_meters(losses_iter, logger)
# update logs
add_to_meters_in_dict(losses_iter, losses)
# construct crop boxes for all the anchors and NMS them - NMS pos ang neg anchors separately
query_fm_sizes = [dataloader.box_coder._get_feature_map_size_per_image_size(sz) for sz in query_img_sizes]
crops = []
achors = []
labels_of_anchors = []
pyramid_level_of_anchors = []
losses_of_anchors = []
corners_of_anchors = []
losses_loc_of_anchors = []
pos_mask_of_anchors = []
pos_loc_mask_of_anchors = []
neg_mask_of_anchors = []
anchor_indices = []
i_image_in_batch = 0 # only one image comes here
for i_p, img_size in enumerate(img_size_pyramid):
for i_label, query_fm_size in enumerate(query_fm_sizes):
crop_position, anchor_position, anchor_index = \
dataloader.box_coder.output_box_grid_generator.get_box_to_cut_anchor(img_size,
crop_size,
image_fm_sizes_p[i_p],
box_reverse_transform_pyramid[i_p])
cur_corners = transform_corners_pyramid[i_p][i_label].transpose(0,1)
cur_corners = dataloader.box_coder.apply_transform_to_corners(cur_corners, box_reverse_transform_pyramid[i_p], img_size)
if cfg.is_cuda:
crop_position, anchor_position = crop_position.cuda(), anchor_position.cuda()
crops.append(crop_position)
achors.append(anchor_position)
device = crop_position.bbox_xyxy.device
losses_of_anchors.append(losses_per_anchor["cls_loss"][i_p][i_image_in_batch, i_label].to(crop_position.bbox_xyxy))
pos_mask_of_anchors.append(losses_per_anchor["pos_mask"][i_p][i_image_in_batch, i_label].to(device=device))
neg_mask_of_anchors.append(losses_per_anchor["neg_mask"][i_p][i_image_in_batch, i_label].to(device=device))
losses_loc_of_anchors.append(losses_per_anchor["loc_loss"][i_p][i_image_in_batch, i_label].to(crop_position.bbox_xyxy))
pos_loc_mask_of_anchors.append(losses_per_anchor["pos_for_regression"][i_p][i_image_in_batch, i_label].to(device=device))
corners_of_anchors.append(cur_corners.to(crop_position.bbox_xyxy))
num_anchors = len(crop_position)
labels_of_anchors.append(torch.full([num_anchors], i_label, dtype=torch.long))
pyramid_level_of_anchors.append(torch.full([num_anchors], i_p, dtype=torch.long))
anchor_indices.append(anchor_index)
# stack all
crops = cat_boxlist(crops)
achors = cat_boxlist(achors)
labels_of_anchors = torch.cat(labels_of_anchors, 0)
pyramid_level_of_anchors = torch.cat(pyramid_level_of_anchors, 0)
losses_of_anchors = torch.cat(losses_of_anchors, 0)
losses_loc_of_anchors = torch.cat(losses_loc_of_anchors, 0)
pos_mask_of_anchors = torch.cat(pos_mask_of_anchors, 0)
pos_loc_mask_of_anchors = torch.cat(pos_loc_mask_of_anchors, 0)
neg_mask_of_anchors = torch.cat(neg_mask_of_anchors, 0)
anchor_indices = torch.cat(anchor_indices, 0)
corners_of_anchors = torch.cat(corners_of_anchors, 0)
def nms_masked_and_collect_data(mask, crops_xyxy, scores, nms_iou_threshold_in_mining, max_etries=None):
mask_ids = torch.nonzero(mask).squeeze(1)
boxes_selected = copy.deepcopy(crops_xyxy[mask])
boxes_selected.add_field("scores", scores[mask])
remaining_boxes = nms(boxes_selected, nms_iou_threshold_in_mining)
remaining_boxes = mask_ids[remaining_boxes]
# sort and take the topk, because NMS is not sorting by default
ids = torch.argsort(scores[remaining_boxes], descending=True)
if max_etries is not None:
ids = ids[:max_etries]
remaining_boxes = remaining_boxes[ids]
return remaining_boxes
nms_iou_threshold_in_mining = cfg.train.mining.nms_iou_threshold_in_mining
num_hard_patches_per_image = cfg.train.mining.num_hard_patches_per_image
# hard negatives
hard_negs = nms_masked_and_collect_data(neg_mask_of_anchors, crops, losses_of_anchors,
nms_iou_threshold_in_mining,
num_hard_patches_per_image)
# hard positives for classification
hard_pos = nms_masked_and_collect_data(pos_mask_of_anchors, crops, losses_of_anchors,
nms_iou_threshold_in_mining,
num_hard_patches_per_image)
# hard positives for localization
hard_pos_loc = nms_masked_and_collect_data(pos_loc_mask_of_anchors, crops, losses_loc_of_anchors,
nms_iou_threshold_in_mining,
num_hard_patches_per_image)
# merge all together
def standardize(v):
return v.item() if type(v) == torch.Tensor else v
def add_item(data, role, pyramid_level, label_local, anchor_index, crop_position_xyxy, anchor_position_xyxy, transform_corners):
new_item = OrderedDict()
new_item["pyramid_level"] = standardize(pyramid_level)
new_item["label_local"] = standardize(label_local)
new_item["anchor_index"] = standardize(anchor_index)
new_item["role"] = role
new_item["crop_position_xyxy"] = crop_position_xyxy
new_item["anchor_position_xyxy"] = anchor_position_xyxy
new_item["transform_corners"] = transform_corners
data.append(new_item)
hardnegdata = []
for i in hard_negs:
add_item(hardnegdata, "neg", pyramid_level_of_anchors[i],
labels_of_anchors[i], anchor_indices[i],
crops[i].cpu(), achors[i].cpu(), corners_of_anchors[i].cpu())
for i in hard_pos:
add_item(hardnegdata, "pos", pyramid_level_of_anchors[i],
labels_of_anchors[i], anchor_indices[i],
crops[i].cpu(), achors[i].cpu(), corners_of_anchors[i].cpu())
for i in hard_pos_loc:
add_item(hardnegdata, "pos_loc", pyramid_level_of_anchors[i],
labels_of_anchors[i], anchor_indices[i],
crops[i].cpu(), achors[i].cpu(), corners_of_anchors[i].cpu())
# extract loss values and compute the box positions to crop
for a in hardnegdata:
a["label_global"] = standardize(batch_class_ids[ a["label_local"] ])
a["loss"] = standardize(losses_per_anchor["cls_loss"][a["pyramid_level"]][i_image_in_batch, a["label_local"], a["anchor_index"]])
a["loss_loc"] = standardize(losses_per_anchor["loc_loss"][a["pyramid_level"]][i_image_in_batch, a["label_local"], a["anchor_index"]])
a["score"] = standardize(image_class_scores_pyramid[a["pyramid_level"]][a["label_local"], a["anchor_index"]])
a["image_id"] = standardize(image_id)
hardnegdata_per_imageid[image_id] = hardnegdata
if cfg.visualization.mining.show_mined_patches:
visualizer.show_mined_patches(image_id, batch_class_ids, dataloader, hardnegdata)
logger.info("Item time: {0}, since mining start: {1}".format(time_since(t_item_start), time_since(t_start_mining)))
logger.info("Hard negative mining finished in {0}".format(time_since(t_start_mining)))
return hardnegdata_per_imageid