utils.py

# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import hashlib
import os
import zipfile
from tempfile import NamedTemporaryFile

from torch.hub import download_url_to_file


def readlines(filename):
    """Read all the lines in a text file and return as a list
    """
    with open(filename, 'r') as f:
        lines = f.read().splitlines()
    return lines


def normalize_image(x):
    """Rescale image pixels to span range [0, 1]
    """
    ma = float(x.max().cpu().data)
    mi = float(x.min().cpu().data)
    d = ma - mi if ma != mi else 1e5
    return (x - mi) / d


def sec_to_hm(t):
    """Convert time in seconds to time in hours, minutes and seconds
    e.g. 10239 -> (2, 50, 39)
    """
    t = int(t)
    s = t % 60
    t //= 60
    m = t % 60
    t //= 60
    return t, m, s


def sec_to_hm_str(t):
    """Convert time in seconds to a nice string
    e.g. 10239 -> '02h50m39s'
    """
    h, m, s = sec_to_hm(t)
    return "{:02d}h{:02d}m{:02d}s".format(h, m, s)


# values are tuples of (<google cloud URL>, <md5 checksum>)
pretrained_models = {
    "mono_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_640x192.zip",
         "a964b8356e08a02d009609d9e3928f7c"),
    "stereo_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_640x192.zip",
         "3dfb76bcff0786e4ec07ac00f658dd07"),
    "mono+stereo_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_640x192.zip",
         "c024d69012485ed05d7eaa9617a96b81"),
    "mono_no_pt_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_no_pt_640x192.zip",
         "9c2f071e35027c895a4728358ffc913a"),
    "stereo_no_pt_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_no_pt_640x192.zip",
         "41ec2de112905f85541ac33a854742d1"),
    "mono+stereo_no_pt_640x192":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_no_pt_640x192.zip",
         "46c3b824f541d143a45c37df65fbab0a"),
    "mono_1024x320":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_1024x320.zip",
         "0ab0766efdfeea89a0d9ea8ba90e1e63"),
    "stereo_1024x320":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_1024x320.zip",
         "afc2f2126d70cf3fdf26b550898b501a"),
    "mono+stereo_1024x320":
        ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_1024x320.zip",
         "cdc5fc9b23513c07d5b19235d9ef08f7"),
}

_state_files = ("depth.pth", "encoder.pth", "pose.pth", "pose_encoder.pth")


def torch_model_dir():
    from torch.hub import _get_torch_home
    torch_home = _get_torch_home()
    return os.path.join(torch_home, 'checkpoints')


def get_state_file(model_name, state_file):
    if model_name not in pretrained_models:
        raise ValueError("Unexpected model name requested: {}. "
                         "Must be one of {}".format(model_name, list(pretrained_models)))
    if state_file not in _state_files:
        raise ValueError("Unexpected state file requested: {}".format(state_file))

    return os.path.join(torch_model_dir(), "-".join(["monodepth2", model_name, state_file]))


def download_model(model_name):
    def check_file_matches_md5(checksum, fpath):
        if not os.path.exists(fpath):
            return False
        with open(fpath, 'rb') as f:
            current_md5checksum = hashlib.md5(f.read()).hexdigest()
        return current_md5checksum == checksum

    model_url, required_md5checksum = pretrained_models[model_name]
    zip_path = NamedTemporaryFile(suffix="monodepth2-{}.zip".format(model_name)).name
    print("-> Downloading pretrained model {}".format(model_name))
    try:
        download_url_to_file(model_url, zip_path)
        if not check_file_matches_md5(required_md5checksum, zip_path):
            raise RuntimeError("Failed to download a file which matches the checksum - quitting")
        model_dir = torch_model_dir()
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        with zipfile.ZipFile(zip_path, 'r') as zf:
            for state_file in _state_files:
                path = get_state_file(model_name, state_file)
                with open(path, 'wb') as f:
                    f.write(zf.read(state_file))
    finally:
        if os.path.exists(zip_path):
            os.unlink(zip_path)