How can images be automatically captured and saved to a directory only when detections are made by the model?

Question

I've trained a computer vision model using YOLOv5 to detect real-time defects in a tape production line. This model is used to create a database of images. What I need is to capture the frame only when the model detects a defect. Currently, the model captures images even when no detections occur.

this is my code

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  # Declare capture_count as global to keep track of the captured images
    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)

        # Perform automatic capture and save when detections are made
        if results.pred[0] is not None:  # Check if there are any predictions
            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  # Increment the capture count for the next image

 
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()