utils.py

import glob
import os
import random
import zipfile

import cv2
import numpy as np
import torch
import yaml


def setup_runtime(args):
    """Load configs, initialize CUDA, CuDNN and the random seeds."""

    # Setup CUDA
    cuda_device_id = args.gpu
    if cuda_device_id is not None:
        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
        os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_device_id)
    if torch.cuda.is_available():
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = True

    # Setup random seeds for reproducibility
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    # Load config
    cfgs = {}
    if args.config is not None and os.path.isfile(args.config):
        cfgs = load_yaml(args.config)

    cfgs["config"] = args.config
    cfgs["seed"] = args.seed
    cfgs["num_workers"] = args.num_workers
    cfgs["device"] = "cuda:0" if torch.cuda.is_available() and cuda_device_id is not None else "cpu"

    print(f"Environment: GPU {cuda_device_id} seed {args.seed} number of workers {args.num_workers}")
    return cfgs


def load_yaml(path):
    print(f"Loading configs from {path}")
    with open(path, "r") as f:
        return yaml.safe_load(f)


def dump_yaml(path, cfgs):
    print(f"Saving configs to {path}")
    xmkdir(os.path.dirname(path))
    with open(path, "w") as f:
        return yaml.safe_dump(cfgs, f)


def xmkdir(path):
    """Create directory PATH recursively if it does not exist."""
    os.makedirs(path, exist_ok=True)


def clean_checkpoint(checkpoint_dir, keep_num=2):
    if keep_num > 0:
        names = list(sorted(glob.glob(os.path.join(checkpoint_dir, "checkpoint*.pth"))))
        names_pose = list(sorted(glob.glob(os.path.join(checkpoint_dir, "pose_checkpoint*.pth"))))
        if len(names) > keep_num:
            for name in names[:-keep_num]:
                print(f"Deleting obslete checkpoint file {name}")
                os.remove(name)

        if len(names_pose) > keep_num:
            for name in names_pose[:-keep_num]:
                print(f"Deleting obslete checkpoint file {name}")
                os.remove(name)


def archive_code(arc_path, filetypes=[".py", ".yml"]):
    print(f"Archiving code to {arc_path}")
    xmkdir(os.path.dirname(arc_path))
    zipf = zipfile.ZipFile(arc_path, "w", zipfile.ZIP_DEFLATED)
    cur_dir = os.getcwd()
    flist = []
    for ftype in filetypes:
        flist.extend(glob.glob(os.path.join(cur_dir, "**", "*" + ftype), recursive=True))
    [zipf.write(f, arcname=f.replace(cur_dir, "archived_code", 1)) for f in flist]
    zipf.close()


def get_model_device(model):
    return next(model.parameters()).device


def set_requires_grad(nets, requires_grad=False):
    if not isinstance(nets, list):
        nets = [nets]
    for net in nets:
        if net is not None:
            for param in net.parameters():
                param.requires_grad = requires_grad


def draw_bbox(im, size):
    b, c, h, w = im.shape
    h2, w2 = (h - size) // 2, (w - size) // 2
    marker = np.tile(np.array([[1.0], [0.0], [0.0]]), (1, size))
    marker = torch.FloatTensor(marker)
    im[:, :, h2, w2 : w2 + size] = marker
    im[:, :, h2 + size, w2 : w2 + size] = marker
    im[:, :, h2 : h2 + size, w2] = marker
    im[:, :, h2 : h2 + size, w2 + size] = marker
    return im


def save_videos(out_fold, imgs, prefix="", suffix="", sep_folder=True, ext=".mp4", cycle=False):
    if sep_folder:
        out_fold = os.path.join(out_fold, suffix)
    xmkdir(out_fold)
    prefix = prefix + "_" if prefix else ""
    suffix = "_" + suffix if suffix else ""
    offset = len(glob.glob(os.path.join(out_fold, prefix + "*" + suffix + ext))) + 1

    imgs = imgs.transpose(0, 1, 3, 4, 2)  # BxTxCxHxW -> BxTxHxWxC
    for i, fs in enumerate(imgs):
        if cycle:
            fs = np.concatenate([fs, fs[::-1]], 0)
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        # fourcc = cv2.VideoWriter_fourcc(*'avc1')
        vid = cv2.VideoWriter(
            os.path.join(out_fold, prefix + "%05d" % (i + offset) + suffix + ext),
            fourcc,
            5,
            (fs.shape[2], fs.shape[1]),
        )
        [vid.write(np.uint8(f[..., ::-1] * 255.0)) for f in fs]
        vid.release()


def save_images(out_fold, imgs, prefix="", suffix="", sep_folder=True, ext=".png"):
    if sep_folder:
        out_fold = os.path.join(out_fold, suffix)
    xmkdir(out_fold)
    prefix = prefix + "_" if prefix else ""
    suffix = "_" + suffix if suffix else ""
    offset = len(glob.glob(os.path.join(out_fold, prefix + "*" + suffix + ext))) + 1

    imgs = imgs.transpose(0, 2, 3, 1)
    for i, img in enumerate(imgs):
        if "depth" in suffix:
            im_out = np.uint16(img[..., ::-1] * 65535.0)
        else:
            im_out = np.uint8(img[..., ::-1] * 255.0)
        cv2.imwrite(os.path.join(out_fold, prefix + "%05d" % (i + offset) + suffix + ext), im_out)


def save_txt(out_fold, data, prefix="", suffix="", sep_folder=True, ext=".txt"):
    if sep_folder:
        out_fold = os.path.join(out_fold, suffix)
    xmkdir(out_fold)
    prefix = prefix + "_" if prefix else ""
    suffix = "_" + suffix if suffix else ""
    offset = len(glob.glob(os.path.join(out_fold, prefix + "*" + suffix + ext))) + 1

    [
        np.savetxt(
            os.path.join(out_fold, prefix + "%05d" % (i + offset) + suffix + ext), d, fmt="%.6f", delimiter=", "
        )
        for i, d in enumerate(data)
    ]


def compute_sc_inv_err(d_pred, d_gt, mask=None):
    b = d_pred.size(0)
    diff = d_pred - d_gt
    if mask is not None:
        diff = diff * mask
        avg = diff.view(b, -1).sum(1) / (mask.view(b, -1).sum(1))
        score = (diff - avg.view(b, 1, 1)) ** 2 * mask
    else:
        avg = diff.view(b, -1).mean(1)
        score = (diff - avg.view(b, 1, 1)) ** 2
    return score  # masked error maps


def compute_angular_distance(n1, n2, mask=None):
    dist = (n1 * n2).sum(3).clamp(-1, 1).acos() / np.pi * 180
    return dist * mask if mask is not None else dist


def save_scores(out_path, scores, header=""):
    print("Saving scores to %s" % out_path)
    np.savetxt(out_path, scores, fmt="%.8f", delimiter=",\t", header=header)