read_input.py 文件源码

def init_file_path(directory):
    """
    Get the image file path array
    :param directory: the directory that store images
    :return: an array of image file path
    """
    paths = []

    if not debug:
        print "Throwing all gray space images now... this may takes some time.."

    for file_name in os.listdir(directory):
        # Skip files that is not jpg
        file_path = '%s/%s' % (directory, file_name)
        if not file_name.endswith('.jpg') or imghdr.what(file_path) is not 'jpeg':
            continue
        if debug:
            paths.append(file_path)
        else:
            # Throw gray space images, this takes long time if have many images
            # TODO: maybe can change to a fast way
            img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
            if len(img.shape) == 3 and img.shape[2] != 1:
                paths.append(file_path)
    return paths