iterator.py 文件源码

def _find_filenames(directory, extensions, class_indices, follow_links=False):
    def _recursive_list(subpath):
        return sorted(os.walk(subpath, followlinks=follow_links), key=lambda tpl: tpl[0])

    classes = []

    filenames = []

    subdir = os.path.basename(directory)

    basedir = os.path.dirname(directory)

    for root, _, files in _recursive_list(directory):
        for fname in files:
            is_valid = False

            for extension in extensions:
                if fname.lower().endswith('.' + extension):
                    is_valid = True
                    break

            if is_valid:
                classes.append(class_indices[subdir])

                # add filename relative to directory
                absolute_path = os.path.join(root, fname)

                filenames.append(os.path.relpath(absolute_path, basedir))

    return classes, filenames