blanchon/Metric3D

1import atexit2import logging3import os4import sys5import time6import torch7from termcolor import colored8 9__all__ = ["setup_logger", ]10 11class _ColorfulFormatter(logging.Formatter):12    def __init__(self, *args, **kwargs):13        self._root_name = kwargs.pop("root_name") + "."14        self._abbrev_name = kwargs.pop("abbrev_name", "")15        if len(self._abbrev_name):16            self._abbrev_name = self._abbrev_name + "."17        super(_ColorfulFormatter, self).__init__(*args, **kwargs)18 19    def formatMessage(self, record):20        record.name = record.name.replace(self._root_name, self._abbrev_name)21        log = super(_ColorfulFormatter, self).formatMessage(record)22        if record.levelno == logging.WARNING:23            prefix = colored("WARNING", "red", attrs=["blink"])24        elif record.levelno == logging.ERROR or record.levelno == logging.CRITICAL:25            prefix = colored("ERROR", "red", attrs=["blink", "underline"])26        else:27            return log28        return prefix + " " + log29 30def setup_logger(31    output=None, distributed_rank=0, *, name='metricdepth', color=True, abbrev_name=None32):33    """34    Initialize the detectron2 logger and set its verbosity level to "DEBUG".35    Args:36        output (str): a file name or a directory to save log. If None, will not save log file.37            If ends with ".txt" or ".log", assumed to be a file name.38            Otherwise, logs will be saved to `output/log.txt`.39        abbrev_name (str): an abbreviation of the module, to avoid log names in logs.40            Set to "" not log the root module in logs.41            By default, will abbreviate "detectron2" to "d2" and leave other42            modules unchanged.43    Returns:44        logging.Logger: a logger45    """46    logger = logging.getLogger()47    logger.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG48    logger.propagate = False49 50    if abbrev_name is None:51        abbrev_name = "d2"52        53    plain_formatter = logging.Formatter(54        "[%(asctime)s] %(name)s %(levelname)s %(message)s ", datefmt="%m/%d %H:%M:%S"55    )56    # stdout logging: master  only57    if distributed_rank == 0:58        ch = logging.StreamHandler(stream=sys.stdout)59        ch.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG60        if color:61            formatter = _ColorfulFormatter(62                colored("[%(asctime)s %(name)s]: ", "green") + "%(message)s",63                datefmt="%m/%d %H:%M:%S",64                root_name=name,65                abbrev_name=str(abbrev_name),66            )67        else:68            formatter = plain_formatter69        ch.setFormatter(formatter)70        logger.addHandler(ch)71    72    # file logging: all workers73    if output is not None:74        if output.endswith(".txt") or output.endswith(".log"):75            filename = output76        else:77            filename = os.path.join(output, "log.txt")78        if distributed_rank > 0:79            filename = filename + ".rank{}".format(distributed_rank)80        os.makedirs(os.path.dirname(filename), exist_ok=True)81 82        fh = logging.StreamHandler(_cached_log_stream(filename))83        fh.setLevel(logging.INFO) # NOTE: if more detailed, change it to logging.DEBUG84        fh.setFormatter(plain_formatter)85        logger.addHandler(fh)86    87 88    return logger89 90from iopath.common.file_io import PathManager as PathManagerBase91 92 93PathManager = PathManagerBase()94 95# cache the opened file object, so that different calls to 'setup_logger96# with the same file name can safely write to the same file.97def _cached_log_stream(filename):98    # use 1K buffer if writting to cloud storage99    io = PathManager.open(filename, "a", buffering=1024 if "://" in filename else -1)100    atexit.register(io.close)101    return io    102