如何在TensorFlow图像处理及Jupyter Notebook中用for循环实现cv.write自动存图

前置准备 #

首先在原有导入基础上，新增输出目录配置，避免每次手动创建文件夹：

import pathlib
import numpy as np
from PIL import Image

# 原有读取配置
PATH_TO_TEST_IMAGES_DIR = pathlib.Path('C:/Object_detection/models-master/research/object_detection/test_images')
TEST_IMAGE_PATHS = sorted(list(PATH_TO_TEST_IMAGES_DIR.glob("*.jpg")))

# 新增输出目录配置
PATH_TO_OUTPUT_DIR = pathlib.Path('C:/Object_detection/models-master/research/object_detection/output_images')
PATH_TO_OUTPUT_DIR.mkdir(exist_ok=True) # 自动创建不存在的输出目录

步骤1：修改可视化推理函数，新增保存逻辑 #

在原有show_inference函数中增加自动保存逻辑，无需手动右键另存：

def show_inference(model, image_path, output_dir):
  image_np = np.array(Image.open(image_path))
  output_dict = run_inference_for_single_image(model, image_np)

  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_reframed', None),
      use_normalized_coordinates=True,
      line_thickness=8)
  
  # 新增自动保存代码
  result_img = Image.fromarray(image_np)
  save_path = output_dir / image_path.name # 保留原文件名保存到输出目录
  result_img.save(save_path)
  
  # 如不需要在Jupyter中显示可注释下方代码
  # display(result_img)

步骤2：新增for循环批量处理所有图片 #

在你完成模型加载后，执行以下循环即可批量处理所有测试图片并自动保存：

for image_path in TEST_IMAGE_PATHS:
    show_inference(model, image_path, PATH_TO_OUTPUT_DIR)

可选调整说明 #

• 如需修改保存格式，可调整保存路径后缀，比如保存为png格式：save_path = output_dir / f"{image_path.stem}.png"
• 如需匹配更多图片格式，可修改读取时的glob规则，比如匹配所有jpg、jpeg、png格式：PATH_TO_TEST_IMAGES_DIR.glob("*.jp*g")、PATH_TO_TEST_IMAGES_DIR.glob("*.png")
• 输出目录路径可根据自身需求自由修改

内容的提问来源于stack exchange，提问作者Code Edge