yolov5输出的边界框坐标如何带图片名

Question

博主你好，帖子“YOLOv5输出检测到的边界框的坐标、类别以及置信度”当中可以在文件中再写入一个识别的图片名吗，比方说识别的文件中，图片名为1.jpg,2.jpg,3.jpg。再txt输出文件的每一行坐标后面，加上图片名：1，2，3

Answer 1

以下回复参考：皆我百晓生、券券喵儿等免费微信小程序作答:

可以在YOLOv5的输出文件中添加图片名。首先，你需要在运行YOLOv5检测时获取图片的文件名，然后在输出结果时将文件名添加到每一行的坐标后面。以下是一个简单的示例：

import os
from pathlib import Path

# 假设你的图片文件夹路径为 images_folder
images_folder = "path/to/your/images"
image_files = os.listdir(images_folder)

# 遍历图片文件夹中的每一张图片
for image_file in image_files:
    # 获取图片文件名（不包括扩展名）
    image_name = Path(image_file).stem

    # 运行YOLOv5检测并将结果保存到 output_folder
    os.system(f"python detect.py --source {os.path.join(images_folder, image_file)} --output {output_folder}")

    # 读取YOLOv5输出的txt文件
    txt_file = os.path.join(output_folder, f"{image_name}.txt")
    with open(txt_file, "r") as f:
        lines = f.readlines()

    # 修改txt文件内容，将图片名添加到每一行的坐标后面
    with open(txt_file, "w") as f:
        for line in lines:
            # 假设每行的格式为：类别 x1 y1 x2 y2 confidence
            parts = line.strip().split(" ")
            new_line = f"{parts[0]} {parts[1]} {parts[2]} {parts[3]} {parts[4]} {parts[5]} {image_name}"
            f.write(new_line + "
")

这段代码首先遍历图片文件夹中的所有图片，然后对每一张图片运行YOLOv5检测并将结果保存到指定的输出文件夹。接下来，它会读取YOLOv5输出的txt文件，并在每一行的坐标后面添加图片名。最后，将修改后的内容写回txt文件。