Roboflow YOLO：如何批量处理图片并保存对应检测结果JSON文件

以下是整合了两个需求的完整代码，包含Drive挂载、批量处理图片、自动保存对应JSON结果：

!pip install roboflow

from google.colab import drive
import os
import json

# 挂载Google Drive
drive.mount('/gdrive')

# 初始化Roboflow模型
rf = Roboflow(api_key="YOUR_PRIVATE_API_KEY")
workspace = rf.workspace("workspace-id")
project = workspace.project("model-id")
version = project.version("version-number")
model = version.model  # 修正原代码冗余写法

# 配置路径：替换为你Drive中的图片文件夹路径
INPUT_FOLDER = "/gdrive/MyDrive/your-image-folder"
# 结果保存路径：自动在原文件夹下创建子目录
OUTPUT_JSON_FOLDER = os.path.join(INPUT_FOLDER, "json_results")
OUTPUT_PLOT_FOLDER = os.path.join(INPUT_FOLDER, "plot_results")

# 创建结果目录（不存在则自动创建）
os.makedirs(OUTPUT_JSON_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_PLOT_FOLDER, exist_ok=True)

# 定义需要处理的图片格式
ALLOWED_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp')

# 批量处理文件夹内所有图片
for filename in os.listdir(INPUT_FOLDER):
    # 过滤非图片文件
    if not filename.lower().endswith(ALLOWED_EXTENSIONS):
        continue
    
    # 拼接完整图片路径
    img_path = os.path.join(INPUT_FOLDER, filename)
    
    # 运行模型预测
    prediction = model.predict(img_path)
    
    # 保存JSON结果：文件名与原图片一致，替换后缀为.json
    json_filename = os.path.splitext(filename)[0] + ".json"
    json_save_path = os.path.join(OUTPUT_JSON_FOLDER, json_filename)
    with open(json_save_path, 'w', encoding='utf-8') as f:
        json.dump(prediction.json(), f, indent=2, ensure_ascii=False)
    
    # 保存可视化预测图（可选）
    plot_filename = os.path.splitext(filename)[0] + "_plot.jpg"
    plot_save_path = os.path.join(OUTPUT_PLOT_FOLDER, plot_filename)
    prediction.save(plot_save_path)
    
    print(f"已处理：{filename}，结果已保存至 {json_save_path} 和 {plot_save_path}")

问题1：保存与输入图片同名的JSON文件 #

核心解决逻辑：

• 用os.path.splitext(filename)分离文件名和后缀，将后缀替换为.json生成结果文件名
• 调用Python内置json库，将prediction.json()返回的字典数据写入文件，指定UTF-8编码避免乱码

关键代码片段：

json_filename = os.path.splitext(filename)[0] + ".json"
json_save_path = os.path.join(OUTPUT_JSON_FOLDER, json_filename)
with open(json_save_path, 'w', encoding='utf-8') as f:
    json.dump(prediction.json(), f, indent=2, ensure_ascii=False)

问题2：批量处理Drive文件夹内所有图片 #

核心解决逻辑：

• 先通过drive.mount挂载Drive，获取目标图片文件夹的绝对路径
• 用os.listdir遍历文件夹，通过后缀过滤仅处理图片文件
• 循环遍历每个图片，调用模型预测并自动保存结果

注意事项：

• 替换INPUT_FOLDER为你Drive中实际的图片文件夹路径（例如/gdrive/MyDrive/TestImages）
• ALLOWED_EXTENSIONS可根据需求添加或移除图片格式
• 代码自动创建结果子目录，无需手动操作

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