process.py 文件源码

def process(input_dir, output_dir, model_dir, resizing_size, gpu):
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3, visible_device_list=gpu)
    with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)).as_default():
        m = loader.LoadedModel(model_dir)

    os.makedirs(output_dir, exist_ok=True)

    input_filenames = glob(os.path.join(input_dir, '*.jpg')) + \
                      glob(os.path.join(input_dir, '*.png')) + \
                      glob(os.path.join(input_dir, '*.tif')) + \
                      glob(os.path.join(input_dir, '*.jp2'))

    for path in tqdm(input_filenames):
        img = Image.open(path).resize(resizing_size)
        mat = np.asarray(img)
        if len(mat.shape) == 2:
            mat = np.stack([mat, mat, mat], axis=2)
        predictions = m.predict(mat[None], prediction_key='labels')[0]
        plt.imsave(os.path.join(output_dir, os.path.relpath(path, input_dir)), predictions)