YOLO-x代码修改，输出检测框坐标和物体类别及其中心坐标

最新推荐文章于 2024-04-15 14:31:21 发布

努力的BigJiang

最新推荐文章于 2024-04-15 14:31:21 发布

阅读量504

点赞数

文章标签： YOLO python 计算机视觉

本文链接：https://blog.csdn.net/m0_62948300/article/details/132941840

版权

yolo-x的demo能够直接输出可视化的结果，但我还想得到物体类别和坐标等信息，直接在demo上更改就行了。

修改def image_demo：

def image_demo(predictor, vis_folder, path, current_time, save_result):
    if os.path.isdir(path):
        files = get_image_list(path)
    else:
        files = [path]
    files.sort()
    for image_name in files:
        outputs, img_info = predictor.inference(image_name)
        result_image = predictor.visual(outputs[0], img_info, predictor.confthre)
       #以下为修改内容
        for bbox in outputs[0]:
            label = int(bbox[6])  # 类别索引
            confidence = bbox[4] * bbox[5]  # 置信度
            class_name = predictor.cls_names[label]  # 类别名称

            # 计算中心坐标
            center_x = (bbox[0] + bbox[2]) / 2.0
            center_y = (bbox[1] + bbox[3]) / 2.0

            # 输出四个角的坐标
            x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]

            print(f"Detected object: {class_name}, Confidence: {confidence:.2f}")
            print(f"  Center coordinates: ({center_x:.2f}, {center_y:.2f})")
            print(f"  Top-left corner: ({x1:.2f}, {y1:.2f})")
            print(f"  Top-right corner: ({x2:.2f}, {y1:.2f})")
            print(f"  Bottom-lef corner: ({x1:.2f}, {y2:.2f})")
            print(f"  Bottom-right corner: ({x2:.2f}, {y2:.2f})")
#修改到此结束

        if save_result:
            save_folder = os.path.join(
                vis_folder, time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
            )
            os.makedirs(save_folder, exist_ok=True)
            save_file_name = os.path.join(save_folder, os.path.basename(image_name))
            logger.info("Saving detection result in {}".format(save_file_name))
            cv2.imwrite(save_file_name, result_image)
        ch = cv2.waitKey(0)
        if ch == 27 or ch == ord("q") or ch == ord("Q"):
            break

看看运行结果:

2023-09-17 13:28:08.605 | INFO     | __main__:main:278 - Args: Namespace(camid=0, ckpt='preModels/yolox_l.pth', conf=0.25, demo='image', device='gpu', exp_file='exps/default/yolox_l.py', experiment_name='yolox_l', fp16=False, fuse=False, legacy=False, name=None, nms=0.45, path='assets/7.jpg', save_result=True, trt=False, tsize=640)
2023-09-17 13:28:09.057 | INFO     | __main__:main:288 - Model Summary: Params: 54.21M, Gflops: 155.65
2023-09-17 13:28:09.925 | INFO     | __main__:main:301 - loading checkpoint
2023-09-17 13:28:11.197 | INFO     | __main__:main:305 - loaded checkpoint done.
2023-09-17 13:28:16.930 | INFO     | __main__:inference:165 - Infer time: 5.7199s
Detected object: bus, Confidence: 0.66
  Center coordinates: (335.09, 276.93)
  Top-left corner: (269.66, 222.62)
  Top-right corner: (400.52, 222.62)
  Bottom-lef corner: (269.66, 331.24)
  Bottom-right corner: (400.52, 331.24)
2023-09-17 13:28:16.934 | INFO     | __main__:image_demo:221 - Saving detection result in ./YOLOX_outputs/yolox_l/vis_res/2023_09_17_13_28_11/7.jpg

成功！