Yolo model.save() 函数将检测到的图像保存到单独的文件夹中

Question

您好，我正在尝试从文件中读取帧并检测每个帧。然后将它们存储在输出文件夹中。但模型所做的是，它为每个输出图像创建不同的文件夹。

import os
import torch
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # yolov5n - yolov5x6 official model
#                                            'custom', 'path/to/best.pt')  # custom model
folderPath = r"C:\yolo\deneme\miniframes"
imageFiles = [f for f in os.listdir(folderPath) if f.endswith(('.jpg', '.jpeg', '.png'))]

output_folder_path = r"C:\yolo\deneme\detectedframes"
#os.makedirs(output_folder_path, exist_ok=True)

output_txt_path = os.path.join(output_folder_path, "results.txt")

all_results = []

for frameNumber, image in enumerate(imageFiles):
    imagePath = os.path.join(folderPath, image)
    results = model(imagePath)
    results.save(save_dir = output_folder_path )

    df = results.pandas().xyxy[0]
    #print(df)

    for i in df['name']:  # name->labels
        all_results.append((frameNumber, i))

# Write results to text file
with open(output_txt_path, 'w') as f:
    for frameNumber, obj in all_results:
        f.write(f"Frame:{frameNumber} {obj}\n")

在这里您可以看到它如何将检测到的结果保存到不同的文件夹中：

这里还有chatgpt解决方案，但它不起作用：

 output_image_path = os.path.join(output_folder_path, f"detected_{image}")
 results.save(save_dir=output_folder_path, save_img=True, img_format='jpg', img_name=f"detected_{image}")

我也尝试使用 model.predict 函数，但如果不通过导入 ultralytics 直接使用 YOLO 模型，它就不起作用。我的 yolov5s.pt 位于我的项目目录中。

编辑：这是yolov5s端的保存功能（yolov5s/common/models）

def save(self, labels=True, save_dir="runs/detect/exp", exist_ok=False):
    save_dir = increment_path(save_dir, exist_ok, mkdir=True)  # increment save_dir
    self._run(save=True, labels=labels, save_dir=save_dir)  # save results

Answer 1

这不是最佳解决方案，但我已经通过将它们放入临时文件夹然后移动来解决它。最后删除临时文件夹。我仍在等待更好的解决方案。这是代码：

` 导入操作系统 进口舒蒂尔 进口火炬 从 PIL 导入 ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True

型号

model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # yolov5n - yolov5x6 官方模型

'custom', 'path/to/best.pt') # 自定义模型

folderPath = r"C:\yolo\deneme\miniframes" imageFiles = [f for f in os.listdir(folderPath) if f.endswith(('.jpg', '.jpeg', '.png'))]

output_folder_path = r“C：\ yolo \ deneme \ detectoredframes” #os.makedirs（output_folder_path，exist_ok = True）

output_txt_path = os.path.join(output_folder_path, "results.txt")

所有结果 = []

对于frameNumber，枚举中的图像（imageFiles）： imagePath = os.path.join(folderPath, image)

results = model(imagePath)

df = results.pandas().xyxy[0]

for i in df['name']:  # name->labels
    all_results.append((frameNumber, i))



output_image_path = os.path.join(output_folder_path, f"detected_{frameNumber}")

temp_save_dir = os.path.join(output_folder_path, "temp_save")
results.save(save_dir = temp_save_dir )
#results.save(save_dir=temp_save_dir, save_img=True, img_format='jpg', img_name=f"detected_{image}")

for file in os.listdir(temp_save_dir):
    file_path = os.path.join(temp_save_dir, file)
    shutil.move(file_path, output_folder_path)

os.rmdir(temp_save_dir)
#print(df)

将结果写入文本文件

打开（output_txt_path，'w'）作为f： 对于 all_results 中的frameNumber、obj： f.write(f"帧：{frameNumber} {obj} ”）

`