load_imagenet.py 文件源码

def load_pickle(pickle_path, dataset_path):
    if not os.path.exists(pickle_path):

        import magic

        image_files = []
        for dir, _, _, in os.walk(dataset_path):
            filenames = glob.glob( os.path.join(dir, '*.JPEG'))  # may be JPEG, depending on your image files
            image_files.append(filenames)

            ## use magic to perform a simple check of the images
            # import magic
            # for filename in filenames:
            #   if magic.from_file(filename, mime=True) == 'image/jpeg':
            #       image_files.append(filename)
            #   else:
            #       print '%s is not a jpeg!' % filename
            #       print magic.from_file(filename)

        if len(image_files) > 0:
            image_files = np.hstack(image_files)

        dataset_filenames = {'image_path':image_files}
        pickle.dump( dataset_filenames, open( pickle_path, "wb" ) )
    else:
        dataset_filenames = pickle.load( open( pickle_path, "rb" ) )
    return dataset_filenames


# return a pd object