imcui/third_party/SOLD2/sold2/train.py

"""
This file implements the training process and all the summaries
"""
import os
import numpy as np
import cv2
import torch
from torch.nn.functional import pixel_shuffle, softmax
from torch.utils.data import DataLoader
import torch.utils.data.dataloader as torch_loader
from tensorboardX import SummaryWriter

from .dataset.dataset_util import get_dataset
from .model.model_util import get_model
from .model.loss import TotalLoss, get_loss_and_weights
from .model.metrics import AverageMeter, Metrics, super_nms
from .model.lr_scheduler import get_lr_scheduler
from .misc.train_utils import (convert_image, get_latest_checkpoint,
                               remove_old_checkpoints)


def customized_collate_fn(batch):
    """ Customized collate_fn. """
    batch_keys = ["image", "junction_map", "heatmap", "valid_mask"]
    list_keys = ["junctions", "line_map"]

    outputs = {}
    for key in batch_keys:
        outputs[key] = torch_loader.default_collate([b[key] for b in batch])
    for key in list_keys:
        outputs[key] = [b[key] for b in batch]

    return outputs


def restore_weights(model, state_dict, strict=True):
    """ Restore weights in compatible mode. """
    # Try to directly load state dict
    try:
        model.load_state_dict(state_dict, strict=strict)
    # Deal with some version compatibility issue (catch version incompatible)
    except:
        err = model.load_state_dict(state_dict, strict=False)
        
        # missing keys are those in model but not in state_dict
        missing_keys = err.missing_keys
        # Unexpected keys are those in state_dict but not in model
        unexpected_keys = err.unexpected_keys

        # Load mismatched keys manually
        model_dict = model.state_dict()
        for idx, key in enumerate(missing_keys):
            dict_keys = [_ for _ in unexpected_keys if not "tracked" in _]
            model_dict[key] = state_dict[dict_keys[idx]]
        model.load_state_dict(model_dict)
    
    return model


def train_net(args, dataset_cfg, model_cfg, output_path):
    """ Main training function. """
    # Add some version compatibility check
    if model_cfg.get("weighting_policy") is None:
        # Default to static
        model_cfg["weighting_policy"] = "static"

    # Get the train, val, test config
    train_cfg = model_cfg["train"]
    test_cfg = model_cfg["test"]

    # Create train and test dataset
    print("\t Initializing dataset...")
    train_dataset, train_collate_fn = get_dataset("train", dataset_cfg)
    test_dataset, test_collate_fn = get_dataset("test", dataset_cfg)

    # Create the dataloader
    train_loader = DataLoader(train_dataset,
                              batch_size=train_cfg["batch_size"],
                              num_workers=8,
                              shuffle=True, pin_memory=True,
                              collate_fn=train_collate_fn)
    test_loader = DataLoader(test_dataset,
                             batch_size=test_cfg.get("batch_size", 1),
                             num_workers=test_cfg.get("num_workers", 1),
                             shuffle=False, pin_memory=False,
                             collate_fn=test_collate_fn)
    print("\t Successfully intialized dataloaders.")


    # Get the loss function and weight first
    loss_funcs, loss_weights = get_loss_and_weights(model_cfg)

    # If resume.
    if args.resume:
        # Create model and load the state dict
        checkpoint = get_latest_checkpoint(args.resume_path,
                                           args.checkpoint_name)
        model = get_model(model_cfg, loss_weights)
        model = restore_weights(model, checkpoint["model_state_dict"])
        model = model.cuda()
        optimizer = torch.optim.Adam(
            [{"params": model.parameters(),
              "initial_lr": model_cfg["learning_rate"]}], 
            model_cfg["learning_rate"], 
            amsgrad=True)
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        # Optionally get the learning rate scheduler
        scheduler = get_lr_scheduler(
            lr_decay=model_cfg.get("lr_decay", False),
            lr_decay_cfg=model_cfg.get("lr_decay_cfg", None),
            optimizer=optimizer)
        # If we start to use learning rate scheduler from the middle
        if ((scheduler is not None)
            and (checkpoint.get("scheduler_state_dict", None) is not None)):
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        start_epoch = checkpoint["epoch"] + 1
    # Initialize all the components.
    else:
        # Create model and optimizer
        model = get_model(model_cfg, loss_weights)
        # Optionally get the pretrained wieghts
        if args.pretrained:
            print("\t [Debug] Loading pretrained weights...")
            checkpoint = get_latest_checkpoint(args.pretrained_path,
                                               args.checkpoint_name)
            # If auto weighting restore from non-auto weighting
            model = restore_weights(model, checkpoint["model_state_dict"],
                                    strict=False)
            print("\t [Debug] Finished loading pretrained weights!")
        
        model = model.cuda()
        optimizer = torch.optim.Adam(
            [{"params": model.parameters(),
              "initial_lr": model_cfg["learning_rate"]}], 
            model_cfg["learning_rate"], 
            amsgrad=True)
        # Optionally get the learning rate scheduler
        scheduler = get_lr_scheduler(
            lr_decay=model_cfg.get("lr_decay", False),
            lr_decay_cfg=model_cfg.get("lr_decay_cfg", None),
            optimizer=optimizer)
        start_epoch = 0
    
    print("\t Successfully initialized model")

    # Define the total loss
    policy = model_cfg.get("weighting_policy", "static")
    loss_func = TotalLoss(loss_funcs, loss_weights, policy).cuda()
    if "descriptor_decoder" in model_cfg:
        metric_func = Metrics(model_cfg["detection_thresh"],
                              model_cfg["prob_thresh"],
                              model_cfg["descriptor_loss_cfg"]["grid_size"],
                              desc_metric_lst='all')
    else:
        metric_func = Metrics(model_cfg["detection_thresh"],
                              model_cfg["prob_thresh"],
                              model_cfg["grid_size"])

    # Define the summary writer
    logdir = os.path.join(output_path, "log")
    writer = SummaryWriter(logdir=logdir)

    # Start the training loop
    for epoch in range(start_epoch, model_cfg["epochs"]):
        # Record the learning rate
        current_lr = optimizer.state_dict()["param_groups"][0]["lr"]
        writer.add_scalar("LR/lr", current_lr, epoch)

        # Train for one epochs
        print("\n\n================== Training ====================")
        train_single_epoch(
            model=model,
            model_cfg=model_cfg,
            optimizer=optimizer,
            loss_func=loss_func,
            metric_func=metric_func,
            train_loader=train_loader,
            writer=writer,
            epoch=epoch)

        # Do the validation
        print("\n\n================== Validation ==================")
        validate(
            model=model,
            model_cfg=model_cfg,
            loss_func=loss_func,
            metric_func=metric_func,
            val_loader=test_loader,
            writer=writer,
            epoch=epoch)

        # Update the scheduler
        if scheduler is not None:
            scheduler.step()

        # Save checkpoints
        file_name = os.path.join(output_path,
                                 "checkpoint-epoch%03d-end.tar"%(epoch))
        print("[Info] Saving checkpoint %s ..." % file_name)
        save_dict = {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "model_cfg": model_cfg}
        if scheduler is not None:
            save_dict.update({"scheduler_state_dict": scheduler.state_dict()})
        torch.save(save_dict, file_name)

        # Remove the outdated checkpoints
        remove_old_checkpoints(output_path, model_cfg.get("max_ckpt", 15))


def train_single_epoch(model, model_cfg, optimizer, loss_func, metric_func,
                       train_loader, writer, epoch):
    """ Train for one epoch. """
    # Switch the model to training mode
    model.train()

    # Initialize the average meter
    compute_descriptors = loss_func.compute_descriptors
    if compute_descriptors:
        average_meter = AverageMeter(is_training=True, desc_metric_lst='all')
    else:
        average_meter = AverageMeter(is_training=True)

    # The training loop
    for idx, data in enumerate(train_loader):
        if compute_descriptors:
            junc_map = data["ref_junction_map"].cuda()
            junc_map2 = data["target_junction_map"].cuda()
            heatmap = data["ref_heatmap"].cuda()
            heatmap2 = data["target_heatmap"].cuda()
            line_points = data["ref_line_points"].cuda()
            line_points2 = data["target_line_points"].cuda()
            line_indices = data["ref_line_indices"].cuda()
            valid_mask = data["ref_valid_mask"].cuda()
            valid_mask2 = data["target_valid_mask"].cuda()
            input_images = data["ref_image"].cuda()
            input_images2 = data["target_image"].cuda()

            # Run the forward pass
            outputs = model(input_images)
            outputs2 = model(input_images2)

            # Compute losses
            losses = loss_func.forward_descriptors(
                outputs["junctions"], outputs2["junctions"],
                junc_map, junc_map2, outputs["heatmap"], outputs2["heatmap"],
                heatmap, heatmap2, line_points, line_points2,
                line_indices, outputs['descriptors'], outputs2['descriptors'],
                epoch, valid_mask, valid_mask2)
        else:
            junc_map = data["junction_map"].cuda()
            heatmap = data["heatmap"].cuda()
            valid_mask = data["valid_mask"].cuda()
            input_images = data["image"].cuda()

            # Run the forward pass
            outputs = model(input_images)

            # Compute losses
            losses = loss_func(
                outputs["junctions"], junc_map,
                outputs["heatmap"], heatmap,
                valid_mask)
        
        total_loss = losses["total_loss"]

        # Update the model
        optimizer.zero_grad()
        total_loss.backward()                     
        optimizer.step()

        # Compute the global step
        global_step = epoch * len(train_loader) + idx
        ############## Measure the metric error #########################
        # Only do this when needed
        if (((idx % model_cfg["disp_freq"]) == 0)
            or ((idx % model_cfg["summary_freq"]) == 0)):
            junc_np = convert_junc_predictions(
                outputs["junctions"], model_cfg["grid_size"],
                model_cfg["detection_thresh"], 300)
            junc_map_np = junc_map.cpu().numpy().transpose(0, 2, 3, 1)

            # Always fetch only one channel (compatible with L1, L2, and CE)
            if outputs["heatmap"].shape[1] == 2:
                heatmap_np = softmax(outputs["heatmap"].detach(),
                                     dim=1).cpu().numpy()
                heatmap_np = heatmap_np.transpose(0, 2, 3, 1)[:, :, :, 1:]
            else:
                heatmap_np = torch.sigmoid(outputs["heatmap"].detach())
                heatmap_np = heatmap_np.cpu().numpy().transpose(0, 2, 3, 1)
            
            heatmap_gt_np = heatmap.cpu().numpy().transpose(0, 2, 3, 1)
            valid_mask_np = valid_mask.cpu().numpy().transpose(0, 2, 3, 1)

            # Evaluate metric results
            if compute_descriptors:
                metric_func.evaluate(
                    junc_np["junc_pred"], junc_np["junc_pred_nms"],
                    junc_map_np, heatmap_np, heatmap_gt_np, valid_mask_np,
                    line_points, line_points2, outputs["descriptors"],
                    outputs2["descriptors"], line_indices)
            else:
                metric_func.evaluate(
                    junc_np["junc_pred"], junc_np["junc_pred_nms"],
                    junc_map_np, heatmap_np, heatmap_gt_np, valid_mask_np)
            # Update average meter
            junc_loss = losses["junc_loss"].item()
            heatmap_loss = losses["heatmap_loss"].item()
            loss_dict = {
                "junc_loss": junc_loss,
                "heatmap_loss": heatmap_loss,
                "total_loss": total_loss.item()}
            if compute_descriptors:
                descriptor_loss = losses["descriptor_loss"].item()
                loss_dict["descriptor_loss"] = losses["descriptor_loss"].item()

            average_meter.update(metric_func, loss_dict, num_samples=junc_map.shape[0])

        # Display the progress
        if (idx % model_cfg["disp_freq"]) == 0:
            results = metric_func.metric_results
            average = average_meter.average()
            # Get gpu memory usage in GB
            gpu_mem_usage = torch.cuda.max_memory_allocated() / (1024 ** 3)
            if compute_descriptors:
                print("Epoch [%d / %d] Iter [%d / %d] loss=%.4f (%.4f), junc_loss=%.4f (%.4f), heatmap_loss=%.4f (%.4f), descriptor_loss=%.4f (%.4f), gpu_mem=%.4fGB"
                      % (epoch, model_cfg["epochs"], idx, len(train_loader),
                         total_loss.item(), average["total_loss"], junc_loss,
                         average["junc_loss"], heatmap_loss,
                         average["heatmap_loss"], descriptor_loss,
                         average["descriptor_loss"], gpu_mem_usage))
            else:
                print("Epoch [%d / %d] Iter [%d / %d] loss=%.4f (%.4f), junc_loss=%.4f (%.4f), heatmap_loss=%.4f (%.4f), gpu_mem=%.4fGB"
                      % (epoch, model_cfg["epochs"], idx, len(train_loader),
                         total_loss.item(), average["total_loss"],
                         junc_loss, average["junc_loss"], heatmap_loss,
                         average["heatmap_loss"], gpu_mem_usage))
            print("\t Junction     precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  % (results["junc_precision"], average["junc_precision"],
                     results["junc_recall"], average["junc_recall"]))
            print("\t Junction nms precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  % (results["junc_precision_nms"],
                     average["junc_precision_nms"],
                     results["junc_recall_nms"], average["junc_recall_nms"]))
            print("\t Heatmap      precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  %(results["heatmap_precision"],
                    average["heatmap_precision"],
                    results["heatmap_recall"], average["heatmap_recall"]))
            if compute_descriptors:
                print("\t Descriptors  matching score=%.4f (%.4f)"
                      %(results["matching_score"], average["matching_score"]))

        # Record summaries
        if (idx % model_cfg["summary_freq"]) == 0:
            results = metric_func.metric_results
            average = average_meter.average()
            # Add the shared losses
            scalar_summaries = {
                "junc_loss": junc_loss,
                "heatmap_loss": heatmap_loss,
                "total_loss": total_loss.detach().cpu().numpy(),
                "metrics": results,
                "average": average}
            # Add descriptor terms
            if compute_descriptors:
                scalar_summaries["descriptor_loss"] = descriptor_loss
                scalar_summaries["w_desc"] = losses["w_desc"]

            # Add weighting terms (even for static terms)
            scalar_summaries["w_junc"] = losses["w_junc"]
            scalar_summaries["w_heatmap"] = losses["w_heatmap"]
            scalar_summaries["reg_loss"] = losses["reg_loss"].item()

            num_images = 3
            junc_pred_binary = (junc_np["junc_pred"][:num_images, ...]
                                > model_cfg["detection_thresh"])
            junc_pred_nms_binary = (junc_np["junc_pred_nms"][:num_images, ...]
                                    > model_cfg["detection_thresh"])
            image_summaries = {
                "image": input_images.cpu().numpy()[:num_images, ...],
                "valid_mask": valid_mask_np[:num_images, ...],
                "junc_map_pred": junc_pred_binary,
                "junc_map_pred_nms": junc_pred_nms_binary,
                "junc_map_gt": junc_map_np[:num_images, ...],
                "junc_prob_map": junc_np["junc_prob"][:num_images, ...],
                "heatmap_pred": heatmap_np[:num_images, ...],
                "heatmap_gt": heatmap_gt_np[:num_images, ...]}
            # Record the training summary
            record_train_summaries(
                writer, global_step, scalars=scalar_summaries,
                images=image_summaries)


def validate(model, model_cfg, loss_func, metric_func, val_loader, writer, epoch):
    """ Validation. """
    # Switch the model to eval mode
    model.eval()

    # Initialize the average meter
    compute_descriptors = loss_func.compute_descriptors
    if compute_descriptors:
        average_meter = AverageMeter(is_training=True, desc_metric_lst='all')
    else:
        average_meter = AverageMeter(is_training=True)

    # The validation loop
    for idx, data in enumerate(val_loader):
        if compute_descriptors:
            junc_map = data["ref_junction_map"].cuda()
            junc_map2 = data["target_junction_map"].cuda()
            heatmap = data["ref_heatmap"].cuda()
            heatmap2 = data["target_heatmap"].cuda()
            line_points = data["ref_line_points"].cuda()
            line_points2 = data["target_line_points"].cuda()
            line_indices = data["ref_line_indices"].cuda()
            valid_mask = data["ref_valid_mask"].cuda()
            valid_mask2 = data["target_valid_mask"].cuda()
            input_images = data["ref_image"].cuda()
            input_images2 = data["target_image"].cuda()

            # Run the forward pass
            with torch.no_grad():
                outputs = model(input_images)
                outputs2 = model(input_images2)

                # Compute losses
                losses = loss_func.forward_descriptors(
                    outputs["junctions"], outputs2["junctions"],
                    junc_map, junc_map2, outputs["heatmap"],
                    outputs2["heatmap"], heatmap, heatmap2, line_points,
                    line_points2, line_indices, outputs['descriptors'],
                    outputs2['descriptors'], epoch, valid_mask, valid_mask2)
        else:
            junc_map = data["junction_map"].cuda()
            heatmap = data["heatmap"].cuda()
            valid_mask = data["valid_mask"].cuda()
            input_images = data["image"].cuda()

            # Run the forward pass
            with torch.no_grad():
                outputs = model(input_images)

                # Compute losses
                losses = loss_func(
                    outputs["junctions"], junc_map,
                    outputs["heatmap"], heatmap,
                    valid_mask)
        total_loss = losses["total_loss"]

        ############## Measure the metric error #########################
        junc_np = convert_junc_predictions(
            outputs["junctions"], model_cfg["grid_size"],
            model_cfg["detection_thresh"], 300)
        junc_map_np = junc_map.cpu().numpy().transpose(0, 2, 3, 1)
        # Always fetch only one channel (compatible with L1, L2, and CE)
        if outputs["heatmap"].shape[1] == 2:
            heatmap_np = softmax(outputs["heatmap"].detach(),
                                 dim=1).cpu().numpy().transpose(0, 2, 3, 1)
            heatmap_np = heatmap_np[:, :, :, 1:]
        else:
            heatmap_np = torch.sigmoid(outputs["heatmap"].detach())
            heatmap_np = heatmap_np.cpu().numpy().transpose(0, 2, 3, 1)


        heatmap_gt_np = heatmap.cpu().numpy().transpose(0, 2, 3, 1)
        valid_mask_np = valid_mask.cpu().numpy().transpose(0, 2, 3, 1)

        # Evaluate metric results
        if compute_descriptors:
            metric_func.evaluate(
                junc_np["junc_pred"], junc_np["junc_pred_nms"],
                junc_map_np, heatmap_np, heatmap_gt_np, valid_mask_np,
                line_points, line_points2, outputs["descriptors"],
                outputs2["descriptors"], line_indices)
        else:
            metric_func.evaluate(
                junc_np["junc_pred"], junc_np["junc_pred_nms"], junc_map_np,
                heatmap_np, heatmap_gt_np, valid_mask_np)
        # Update average meter
        junc_loss = losses["junc_loss"].item()
        heatmap_loss = losses["heatmap_loss"].item()
        loss_dict = {
            "junc_loss": junc_loss,
            "heatmap_loss": heatmap_loss,
            "total_loss": total_loss.item()}
        if compute_descriptors:
            descriptor_loss = losses["descriptor_loss"].item()
            loss_dict["descriptor_loss"] = losses["descriptor_loss"].item()
        average_meter.update(metric_func, loss_dict, num_samples=junc_map.shape[0])

        # Display the progress
        if (idx % model_cfg["disp_freq"]) == 0:
            results = metric_func.metric_results
            average = average_meter.average()
            if compute_descriptors:
                print("Iter [%d / %d] loss=%.4f (%.4f), junc_loss=%.4f (%.4f), heatmap_loss=%.4f (%.4f), descriptor_loss=%.4f (%.4f)"
                      % (idx, len(val_loader),
                         total_loss.item(), average["total_loss"],
                         junc_loss, average["junc_loss"],
                         heatmap_loss, average["heatmap_loss"],
                         descriptor_loss, average["descriptor_loss"]))
            else:
                print("Iter [%d / %d] loss=%.4f (%.4f), junc_loss=%.4f (%.4f), heatmap_loss=%.4f (%.4f)"
                      % (idx, len(val_loader),
                         total_loss.item(), average["total_loss"],
                         junc_loss, average["junc_loss"],
                         heatmap_loss, average["heatmap_loss"]))
            print("\t Junction     precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  % (results["junc_precision"], average["junc_precision"],
                     results["junc_recall"], average["junc_recall"]))
            print("\t Junction nms precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  % (results["junc_precision_nms"],
                     average["junc_precision_nms"],
                     results["junc_recall_nms"], average["junc_recall_nms"]))
            print("\t Heatmap      precision=%.4f (%.4f) / recall=%.4f (%.4f)"
                  % (results["heatmap_precision"],
                     average["heatmap_precision"],
                     results["heatmap_recall"], average["heatmap_recall"]))
            if compute_descriptors:
                print("\t Descriptors  matching score=%.4f (%.4f)"
                      %(results["matching_score"], average["matching_score"]))

    # Record summaries
    average = average_meter.average()
    scalar_summaries = {"average": average}
    # Record the training summary
    record_test_summaries(writer, epoch, scalar_summaries)


def convert_junc_predictions(predictions, grid_size,
                             detect_thresh=1/65, topk=300):
    """ Convert torch predictions to numpy arrays for evaluation. """
    # Convert to probability outputs first
    junc_prob = softmax(predictions.detach(), dim=1).cpu()
    junc_pred = junc_prob[:, :-1, :, :]

    junc_prob_np = junc_prob.numpy().transpose(0, 2, 3, 1)[:, :, :, :-1]
    junc_prob_np = np.sum(junc_prob_np, axis=-1)
    junc_pred_np = pixel_shuffle(
        junc_pred, grid_size).cpu().numpy().transpose(0, 2, 3, 1)
    junc_pred_np_nms = super_nms(junc_pred_np, grid_size, detect_thresh, topk)
    junc_pred_np = junc_pred_np.squeeze(-1)

    return {"junc_pred": junc_pred_np, "junc_pred_nms": junc_pred_np_nms,
            "junc_prob": junc_prob_np}


def record_train_summaries(writer, global_step, scalars, images):
    """ Record training summaries. """
    # Record the scalar summaries
    results = scalars["metrics"]
    average = scalars["average"]

    # GPU memory part
    # Get gpu memory usage in GB
    gpu_mem_usage = torch.cuda.max_memory_allocated() / (1024 ** 3)
    writer.add_scalar("GPU/GPU_memory_usage", gpu_mem_usage, global_step)

    # Loss part
    writer.add_scalar("Train_loss/junc_loss", scalars["junc_loss"],
                      global_step)
    writer.add_scalar("Train_loss/heatmap_loss", scalars["heatmap_loss"],
                      global_step)
    writer.add_scalar("Train_loss/total_loss", scalars["total_loss"],
                      global_step)
    # Add regularization loss
    if "reg_loss" in scalars.keys():
        writer.add_scalar("Train_loss/reg_loss", scalars["reg_loss"],
                          global_step)
    # Add descriptor loss
    if "descriptor_loss" in scalars.keys():
        key = "descriptor_loss"
        writer.add_scalar("Train_loss/%s"%(key), scalars[key], global_step)
        writer.add_scalar("Train_loss_average/%s"%(key), average[key],
                          global_step)
    
    # Record weighting
    for key in scalars.keys():
        if "w_" in key:
            writer.add_scalar("Train_weight/%s"%(key), scalars[key],
                              global_step)
    
    # Smoothed loss
    writer.add_scalar("Train_loss_average/junc_loss", average["junc_loss"],
                      global_step)
    writer.add_scalar("Train_loss_average/heatmap_loss",
                      average["heatmap_loss"], global_step)
    writer.add_scalar("Train_loss_average/total_loss", average["total_loss"],
                      global_step)
    # Add smoothed descriptor loss
    if "descriptor_loss" in average.keys():
        writer.add_scalar("Train_loss_average/descriptor_loss",
                          average["descriptor_loss"], global_step)

    # Metrics part
    writer.add_scalar("Train_metrics/junc_precision",
                      results["junc_precision"], global_step)
    writer.add_scalar("Train_metrics/junc_precision_nms",
                      results["junc_precision_nms"], global_step)
    writer.add_scalar("Train_metrics/junc_recall",
                      results["junc_recall"], global_step)
    writer.add_scalar("Train_metrics/junc_recall_nms",
                      results["junc_recall_nms"], global_step)
    writer.add_scalar("Train_metrics/heatmap_precision",
                      results["heatmap_precision"], global_step)
    writer.add_scalar("Train_metrics/heatmap_recall",
                      results["heatmap_recall"], global_step)
    # Add descriptor metric
    if "matching_score" in results.keys():
        writer.add_scalar("Train_metrics/matching_score",
                          results["matching_score"], global_step)

    # Average part
    writer.add_scalar("Train_metrics_average/junc_precision",
                      average["junc_precision"], global_step)
    writer.add_scalar("Train_metrics_average/junc_precision_nms",
                      average["junc_precision_nms"], global_step)
    writer.add_scalar("Train_metrics_average/junc_recall",
                      average["junc_recall"], global_step)
    writer.add_scalar("Train_metrics_average/junc_recall_nms",
                      average["junc_recall_nms"], global_step)
    writer.add_scalar("Train_metrics_average/heatmap_precision",
                      average["heatmap_precision"], global_step)
    writer.add_scalar("Train_metrics_average/heatmap_recall",
                      average["heatmap_recall"], global_step)
    # Add smoothed descriptor metric
    if "matching_score" in average.keys():
        writer.add_scalar("Train_metrics_average/matching_score",
                          average["matching_score"], global_step)

    # Record the image summary
    # Image part
    image_tensor = convert_image(images["image"], 1)
    valid_masks = convert_image(images["valid_mask"], -1)
    writer.add_images("Train/images", image_tensor, global_step,
                      dataformats="NCHW")
    writer.add_images("Train/valid_map", valid_masks, global_step,
                      dataformats="NHWC")

    # Heatmap part
    writer.add_images("Train/heatmap_gt",
                      convert_image(images["heatmap_gt"], -1), global_step,
                      dataformats="NHWC")
    writer.add_images("Train/heatmap_pred",
                      convert_image(images["heatmap_pred"], -1), global_step,
                      dataformats="NHWC")

    # Junction prediction part
    junc_plots = plot_junction_detection(
        image_tensor, images["junc_map_pred"],
        images["junc_map_pred_nms"], images["junc_map_gt"])
    writer.add_images("Train/junc_gt", junc_plots["junc_gt_plot"] / 255.,
                      global_step, dataformats="NHWC")
    writer.add_images("Train/junc_pred", junc_plots["junc_pred_plot"] / 255.,
                      global_step, dataformats="NHWC")
    writer.add_images("Train/junc_pred_nms",
                      junc_plots["junc_pred_nms_plot"] / 255., global_step,
                      dataformats="NHWC")
    writer.add_images(
        "Train/junc_prob_map",
        convert_image(images["junc_prob_map"][..., None], axis=-1),
        global_step, dataformats="NHWC")


def record_test_summaries(writer, epoch, scalars):
    """ Record testing summaries. """
    average = scalars["average"]

    # Average loss
    writer.add_scalar("Val_loss/junc_loss", average["junc_loss"], epoch)
    writer.add_scalar("Val_loss/heatmap_loss", average["heatmap_loss"], epoch)
    writer.add_scalar("Val_loss/total_loss", average["total_loss"], epoch)
    # Add descriptor loss
    if "descriptor_loss" in average.keys():
        key = "descriptor_loss"
        writer.add_scalar("Val_loss/%s"%(key), average[key], epoch)

    # Average metrics
    writer.add_scalar("Val_metrics/junc_precision", average["junc_precision"],
                      epoch)
    writer.add_scalar("Val_metrics/junc_precision_nms",
                      average["junc_precision_nms"], epoch)
    writer.add_scalar("Val_metrics/junc_recall",
                      average["junc_recall"], epoch)
    writer.add_scalar("Val_metrics/junc_recall_nms",
                      average["junc_recall_nms"], epoch)
    writer.add_scalar("Val_metrics/heatmap_precision",
                      average["heatmap_precision"], epoch)
    writer.add_scalar("Val_metrics/heatmap_recall",
                      average["heatmap_recall"], epoch)
    # Add descriptor metric
    if "matching_score" in average.keys():
        writer.add_scalar("Val_metrics/matching_score",
                          average["matching_score"], epoch)


def plot_junction_detection(image_tensor, junc_pred_tensor,
                            junc_pred_nms_tensor, junc_gt_tensor):
    """ Plot the junction points on images. """
    # Get the batch_size
    batch_size = image_tensor.shape[0]

    # Process through batch dimension
    junc_pred_lst = []
    junc_pred_nms_lst = []
    junc_gt_lst = []
    for i in range(batch_size):
        # Convert image to 255 uint8
        image = (image_tensor[i, :, :, :]
                 * 255.).astype(np.uint8).transpose(1,2,0)

        # Plot groundtruth onto image
        junc_gt = junc_gt_tensor[i, ...]
        coord_gt = np.where(junc_gt.squeeze() > 0)
        points_gt = np.concatenate((coord_gt[0][..., None],
                                    coord_gt[1][..., None]),
                                    axis=1)
        plot_gt = image.copy()
        for id in range(points_gt.shape[0]):
            cv2.circle(plot_gt, tuple(np.flip(points_gt[id, :])), 3,
                       color=(255, 0, 0), thickness=2)
        junc_gt_lst.append(plot_gt[None, ...])

        # Plot junc_pred
        junc_pred = junc_pred_tensor[i, ...]
        coord_pred = np.where(junc_pred > 0)
        points_pred = np.concatenate((coord_pred[0][..., None],
                                      coord_pred[1][..., None]),
                                      axis=1)
        plot_pred = image.copy()
        for id in range(points_pred.shape[0]):
            cv2.circle(plot_pred, tuple(np.flip(points_pred[id, :])), 3,
                       color=(0, 255, 0), thickness=2)
        junc_pred_lst.append(plot_pred[None, ...])

        # Plot junc_pred_nms
        junc_pred_nms = junc_pred_nms_tensor[i, ...]
        coord_pred_nms = np.where(junc_pred_nms > 0)
        points_pred_nms = np.concatenate((coord_pred_nms[0][..., None],
                                          coord_pred_nms[1][..., None]),
                                          axis=1)
        plot_pred_nms = image.copy()
        for id in range(points_pred_nms.shape[0]):
            cv2.circle(plot_pred_nms, tuple(np.flip(points_pred_nms[id, :])),
                       3, color=(0, 255, 0), thickness=2)
        junc_pred_nms_lst.append(plot_pred_nms[None, ...])

    return {"junc_gt_plot": np.concatenate(junc_gt_lst, axis=0),
            "junc_pred_plot": np.concatenate(junc_pred_lst, axis=0),
            "junc_pred_nms_plot": np.concatenate(junc_pred_nms_lst, axis=0)}