utils.py

import os
import torch
import shutil
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import math

cmap = plt.cm.viridis
counter = 0


def parse_command():
    data_names = ["nyudepthv2"]

    from dataloaders.dataloader import MyDataloader

    modality_names = MyDataloader.modality_names

    import argparse

    parser = argparse.ArgumentParser(description="FastDepth")
    parser.add_argument(
        "--data",
        metavar="DATA",
        default="nyudepthv2",
        choices=data_names,
        help="dataset: " + " | ".join(data_names) + " (default: nyudepthv2)",
    )
    parser.add_argument(
        "--modality",
        "-m",
        metavar="MODALITY",
        default="rgb",
        choices=modality_names,
        help="modality: " + " | ".join(modality_names) + " (default: rgb)",
    )
    parser.add_argument(
        "-j",
        "--workers",
        default=16,
        type=int,
        metavar="N",
        help="number of data loading workers (default: 16)",
    )
    parser.add_argument(
        "--print-freq",
        "-p",
        default=50,
        type=int,
        metavar="N",
        help="print frequency (default: 50)",
    )
    parser.add_argument(
        "--model",
        default="",
        type=str,
        metavar="PATH",
    )
    parser.add_argument(
        "-e",
        "--evaluate",
        default="",
        type=str,
        metavar="PATH",
    )
    parser.add_argument(
        "-r", "--run", required=False, action="store_true", help="Specify a run option"
    )
    parser.add_argument(
        "--rescale", required=False, action="store_true", help="Specify a save option"
    )
    parser.add_argument(
        "-i", "--image", type=str, required=False, help="Path to image JPG or PNG."
    )
    parser.add_argument(
        "-f",
        "--folder",
        type=str,
        required=False,
        help="Path to image folder JPG or PNG.",
    )
    parser.add_argument(
        "--txt",
        type=str,
        required=False,
        help="Path to list of images.",
    )
    parser.add_argument(
        "--cam",
        type=str,
        required=True,
        metavar="CAMERA_CONFIG",
        help="Path to YAML file with camera configuration",
    )

    parser.add_argument("--gpu", default="0", type=str, metavar="N", help="gpu id")
    parser.set_defaults(cuda=True)

    args = parser.parse_args()
    return args


def colored_depthmap(depth, d_min=None, d_max=None):
    if d_min is None:
        d_min = np.min(depth)
    if d_max is None:
        d_max = np.max(depth)
    depth_relative = (depth - d_min) / (d_max - d_min)
    return 255 * cmap(depth_relative)[:, :, :3]  # H, W, C


def merge_into_row(input, depth_target, depth_pred):
    rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1, 2, 0))  # H, W, C
    depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
    depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())

    d_min = min(np.min(depth_target_cpu), np.min(depth_pred_cpu))
    d_max = max(np.max(depth_target_cpu), np.max(depth_pred_cpu))
    depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
    depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)
    img_merge = np.hstack([rgb, depth_target_col, depth_pred_col])

    return img_merge


def merge_into_row_with_gt(input, depth_input, depth_target, depth_pred):
    rgb = 255 * np.transpose(np.squeeze(input.cpu().numpy()), (1, 2, 0))  # H, W, C
    depth_input_cpu = np.squeeze(depth_input.cpu().numpy())
    depth_target_cpu = np.squeeze(depth_target.cpu().numpy())
    depth_pred_cpu = np.squeeze(depth_pred.data.cpu().numpy())

    d_min = min(
        np.min(depth_input_cpu), np.min(depth_target_cpu), np.min(depth_pred_cpu)
    )
    d_max = max(
        np.max(depth_input_cpu), np.max(depth_target_cpu), np.max(depth_pred_cpu)
    )
    depth_input_col = colored_depthmap(depth_input_cpu, d_min, d_max)
    depth_target_col = colored_depthmap(depth_target_cpu, d_min, d_max)
    depth_pred_col = colored_depthmap(depth_pred_cpu, d_min, d_max)

    img_merge = np.hstack([rgb, depth_input_col, depth_target_col, depth_pred_col])

    return img_merge


def add_row(img_merge, row):
    return np.vstack([img_merge, row])


def save_image(img_merge, filename):
    img_merge = Image.fromarray(img_merge.astype("uint8"))
    img_merge.save(filename)