在Google Colab中如何将TensorFlow模型推理输出的matplotlib图像保存到Google Drive？

1. 挂载Google Drive到Colab运行环境 #

先在你的Colab笔记本中新建一个代码单元格，运行以下代码完成云盘挂载：

from google.colab import drive
drive.mount('/content/drive')

运行后会弹出授权窗口，按照提示登录你的谷歌账号，复制授权码粘贴到输入框回车即可完成挂载。
你可以提前在谷歌云盘根目录创建名为rps检测结果的文件夹用来存图，也可以运行以下代码自动创建：

import os
# exist_ok=True表示如果文件夹已存在就不会重复创建报错
os.makedirs('/content/drive/MyDrive/rps检测结果', exist_ok=True)

2. 修改原有推理代码添加保存逻辑 #

找到你原有代码末尾遍历测试图像的for循环，在绘制完图像的部分添加保存代码即可，修改后的循环片段如下：

for image_path in TEST_IMAGE_PATHS:
    image = Image.open(image_path)
    print(image_path)
    # 以下为原有逻辑，保持不变
    image_np = load_image_into_numpy_array(image)
    image_np_expanded = np.expand_dims(image_np, axis=0)
    output_dict = run_inference_for_single_image(image_np, detection_graph)
    vis_util.visualize_boxes_and_labels_on_image_array(
        image_np,
        output_dict['detection_boxes'],
        output_dict['detection_classes'],
        output_dict['detection_scores'],
        category_index,
        instance_masks=output_dict.get('detection_masks'),
        use_normalized_coordinates=True,
        line_thickness=40)
    plt.figure(figsize=IMAGE_SIZE)
    plt.imshow(image_np)

    # 新增保存逻辑
    # 提取原图像文件名作为保存时的文件名，避免重名
    base_name = os.path.basename(image_path)
    save_full_path = f'/content/drive/MyDrive/rps检测结果/检测结果_{base_name}'
    # 保存图像，bbox_inches='tight'参数可以避免图像边缘内容被裁剪
    plt.savefig(save_full_path, bbox_inches='tight', dpi=300)

    plt.show()
    # 关闭当前画布释放内存，避免多张图叠加占用资源
    plt.close()

如果你不需要matplotlib的画布边框，也可以直接把绘制完检测框的numpy数组转成PIL图像直接保存，速度更快：

# 替换上述plt.savefig部分的代码即可
from PIL import Image
result_image = Image.fromarray(image_np)
result_image.save(save_full_path)

完成修改后重新运行代码，所有检测结果图就会自动保存到你的谷歌云盘对应的文件夹里了。

内容的提问来源于stack exchange，提问作者Ashley Afueh