如何实现仅当YOLOv5模型检测到缺陷时自动捕获并保存图像？

你的代码中判断是否检测到缺陷的条件有误，导致即使没有检测结果时也会保存图像。原因是**results.pred[0]即使没有检测目标时也不是None，而是一个空的张量**，所以原来的if results.pred[0] is not None条件永远成立。

正确的判断方式是检查检测结果的数量，也就是判断results.pred[0]的长度是否大于0，只有当存在至少一个检测框时才保存图像。

修改后的完整代码如下：

import tkinter as tk
import torch
import cv2
import numpy as np
from PIL import Image, ImageTk

model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5/runs/train/exp2/weights/best.pt')

capture_count = 1

def update_feed(label):
    global capture_count
    ret, frame = cap.read()
    if ret:
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = model(frame_rgb)
        rendered_img = Image.fromarray(np.squeeze(results.render()))

        imgtk = ImageTk.PhotoImage(image=rendered_img)
        label.imgtk = imgtk
        label.configure(image=imgtk)
        label.after(10, update_feed, label)

        # 仅在检测到缺陷时保存图像
        if len(results.pred[0]) > 0:  # 判断是否存在检测结果
            save_path = f"D:/Projects/captured_frame_auto_{capture_count}.jpg"
            annotated_img = results.render()[0]
            annotated_img_rgb = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
            cv2.imwrite(save_path, annotated_img_rgb)
            capture_count += 1

cap = cv2.VideoCapture(0)

root = tk.Tk()
root.title("Defect Detection App")

video_feed = tk.Label(root)
video_feed.pack()

update_feed(video_feed)

def close_app():
    cap.release()
    root.destroy()

root.mainloop()

额外优化建议：

• 避免重复捕获同一缺陷：可以设置冷却时间，比如1秒内仅保存一次图像，防止短时间内连续保存相同缺陷帧
• 过滤低置信度误检：添加置信度阈值判断，比如if len(results.pred[0]) > 0 and results.pred[0][:, 4].max() > 0.5，只保留置信度高于0.5的检测结果

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