pytorch-semseg图像分割代码解读utils.py

"""
Misc Utility functions
"""
import os
import logging
import datetime
import numpy as np

from collections import OrderedDict


def recursive_glob(rootdir=".", suffix=""):
    """Performs recursive glob with given suffix and rootdir
        :param rootdir is the root directory
        :param suffix is the suffix to be searched
    """
    # 这个return的作用是返回一个列表，列表是所有符合条件的文件名
    # 中括号叫做列表解析，根据已有的rootdir列表搜索符合条件的值重新组成列表
    # 返回值是在for中所有的符合if条件的os.path.join，一个列表
    return [
        os.path.join(looproot, filename)
        for looproot, _, filenames in os.walk(rootdir)# os.walk返回的（起始目录，子目录，文件名）
        for filename in filenames
        if filename.endswith(suffix)
    ]


def alpha_blend(input_image, segmentation_mask, alpha=0.5):
    """Alpha Blending utility to overlay RGB masks on RBG images
        :param input_image is a np.ndarray with 3 channels
        :param segmentation_mask is a np.ndarray with 3 channels
        :param alpha is a float value
    """
    # 设置透明度
    blended = np.zeros(input_image.size, dtype=np.float32)
    blended = input_image * alpha + segmentation_mask * (1 - alpha)
    return blended


def convert_state_dict(state_dict):
    """Converts a state dict saved from a dataParallel module to normal
       module state_dict inplace
       :param state_dict is the loaded DataParallel model_state
    """
    # next中是false即读进来的有问题直接返回
    # 将原模型参数按有序字典存储并返回，其中name剃掉了'module.'几个字符
    if not next(iter(state_dict)).startswith("module."):
        return state_dict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]
        new_state_dict[name] = v
    return new_state_dict


# 保存日志
def get_logger(logdir):
    logger = logging.getLogger("ptsemseg")
    ts = str(datetime.datetime.now()).split(".")[0].replace(" ", "_")
    ts = ts.replace(":", "_").replace("-", "_")
    file_path = os.path.join(logdir, "run_{}.log".format(ts))
    hdlr = logging.FileHandler(file_path)
    formatter = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
    hdlr.setFormatter(formatter)
    logger.addHandler(hdlr)
    logger.setLevel(logging.INFO)
    return logger

需要注意的地方就是第一个函数的列表解析，值得学习，简洁高效