make_prediction_layouts.py 文件源码

def find_source_dirs(path_root_dir):
    """Find source directories to make layouts, going at most 1 layer deep.

    Returns : list of source directories
    """
    def is_source_dir(path):
        if not os.path.isdir(path):
            return False
        has_signal, has_target, has_prediction = False, False, False
        for entry in [i.path for i in os.scandir(path) if i.is_file()]:
            if any(tag in entry for tag in TAGS_SIGNAL):
                has_signal = True
            if any(tag in entry for tag in TAGS_TARGET):
                has_target = True
            if any(tag in entry for tag in TAGS_PREDICTION):
                has_prediction = True
        return has_signal and has_target and has_prediction

    if is_source_dir(path_root_dir):
        return [path_root_dir]
    results = []
    for entry in os.scandir(path_root_dir):
        if is_source_dir(entry.path):
            results.append(entry.path)
    return results