我有我的自定义训练模型（best.pt），它检测人和车头灯两件事。现在我想要根据这些条件输出

Question

你能帮我吗.............. 我有我的自定义训练模型（best.pt），它检测人和车头灯两件事。现在我想要根据以下条件输出： 1. 如果模型仅检测到车头灯返回 0, 2. 如果模型仅检测到人返回 1, 3. 如果模型检测到车头灯和人都返回 0.

import cv2
from ultralytics import YOLO

video_path = 'data/video1.mp4'
video_out_path = 'out.mp4'

cap = cv2.VideoCapture(video_path)

# Check if the video file is opened successfully
if not cap.isOpened():
    print("Error: Could not open the video file.")
    exit()

ret, frame = cap.read()

# Check if the first frame is read successfully
if not ret:
    print("Error: Could not read the first frame from the video.")
    exit()

cap_out = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'MP4V'), cap.get(cv2.CAP_PROP_FPS),
                          (int(cap.get(3)), int(cap.get(4))))  # Use cap.get(3) and cap.get(4) for width and height

model = YOLO("bestall5.pt")

detection_threshold = 0.5
while ret:
    results_list = model(frame)

    headlight_detected = False
    person_detected = False

    # Iterate through the list of results
    for results in results_list:
        # Check if the current result has the necessary attributes
        if hasattr(results, 'xyxy'):
            for result in results.xyxy:
                x1, y1, x2, y2, score, class_id = result.tolist()
                x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)

                # Assuming class_id is the index of the class in the model's class list
                class_name = model.names[class_id]

                if class_name == "headlight" and score > detection_threshold:
                    headlight_detected = True
                elif class_name == "person" and score > detection_threshold:
                    person_detected = True

    # Output based on the specified conditions
    if headlight_detected and person_detected:
        output = 0
    elif headlight_detected:
        output = 0
    elif person_detected:
        output = 1
    else:
        output = -1  # No person or headlight detected

    print("Output:", output)

    cap_out.write(frame)

    cv2.imshow('Object Detection', frame)
    
    # Break the loop if 'q' key is pressed
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

    ret, frame = cap.read()

cap.release()
cap_out.release()
cv2.destroyAllWindows()

我尝试了这个，但只得到 -1 作为输出，但我的视频既有车头灯又有人物

Answer 1

这里的条件

if hasattr(results, 'xyxy')

始终为负。 results 的

可用属性

有：

orig_img, orig_shape, boxes, masks, probs, keypoints, obb, speed, names, path

要获取xyxy框坐标、score和class_id，请参阅

results.boxes

。 boxes 的

可用属性

有：

xyxy, conf, cls, id, xywh, xyxyn, xywhn

他们都以

torch.Tensor

的形式回归。要获取行值，您可以执行以下操作：

# Iterate through the list of results
for results in results_list:
    # Check if the current result has the necessary attributes
    if hasattr(results, 'boxes'):
        for box in results.boxes:
            x1, y1, x2, y2 = box.xyxy.tolist()[0]
            x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)
            score = box.conf.item()
            class_id = int(box.cls.item())