我有一个自定义训练的模型(best.pt),它可以检测两种对象:人与车灯。现在我希望根据以下条件输出结果

请问您能帮我吗………我有一个自定义训练的模型(best.pt),它可以检测两种对象:人与车灯。现在我希望根据以下条件输出结果:1. 如果模型只检测到车灯,返回0,2. 如果模型只检测到人,返回1,3. 如果模型同时检测到车灯和人,返回0。

import cv2from ultralytics import YOLOvideo_path = 'data/video1.mp4'video_out_path = 'out.mp4'cap = cv2.VideoCapture(video_path)# Check if the video file is opened successfullyif not cap.isOpened():    print("Error: Could not open the video file.")    exit()ret, frame = cap.read()# Check if the first frame is read successfullyif 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 heightmodel = YOLO("bestall5.pt")detection_threshold = 0.5while 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作为输出,而我的视频中同时有车灯和人


回答:

这里的条件if hasattr(results, 'xyxy')总是为负。根据文档results的可用属性是:

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

要获取xyxy框坐标、scoreclass_id,请参考results.boxes。根据文档boxes的可用属性是:

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

这些属性都以torch.Tensor的形式返回。要获取行值,您可以这样做:

# Iterate through the list of resultsfor 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())

Related Posts

神经网络反向传播代码不工作

我需要编写一个简单的由1个输出节点、1个包含3个节点的…

值错误:y 包含先前未见过的标签:

我使用了 决策树分类器,我想将我的 输入 作为 字符串…

使用不平衡数据集进行特征选择时遇到的问题

我正在使用不平衡数据集(54:38:7%)进行特征选择…

广义随机森林/因果森林在Python上的应用

我在寻找Python上的广义随机森林/因果森林算法,但…

如何用PyTorch仅用标量损失来训练神经网络?

假设我们有一个神经网络,我们希望它能根据输入预测三个值…

什么是RNN中间隐藏状态的良好用途?

我已经以三种不同的方式使用了RNN/LSTM: 多对多…

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注