目标检测txt标签绘制在图片上(yolo)

        为了满足论文对比原始标注的需求,我特地从现有代码中提炼出一个简明脚本,旨在实现标签的可视化呈现。如下:


from ultralytics.utils.plotting import  colors  #  需要引用代码原来的模块
names = ['bridge', 'ship', 'airplane', 'oil tank']  # 换成自己的类别
import os
import cv2
import torch

def find_txt_by_image_name(image_name, labels_folder):
    """根据图片名称在指定的标签文件夹中查找对应的txt文件路径"""
    txt_name = os.path.splitext(image_name)[0] + '.txt'  # 移除图片扩展名并加上.txt
    return os.path.join(labels_folder, txt_name)


def convert_to_absolute_coords(image_path,  labels_folder):
    """
    读取YOLov8的.txt标签,并将相对坐标转换为绝对坐标。
    假设图片与txt文件名相同(不包括扩展名)。
    """
    # 获取图片的宽度和高度
    img = cv2.imread(image_path)
    height, width, _ = img.shape

    # 确保对应的txt文件存在并读取坐标信息
    txt_path = find_txt_by_image_name(os.path.basename(image_path), labels_folder)
    if not os.path.exists(txt_path):
        print(f"Warning: TXT file not found for image {os.path.basename(image_path)}")
        return []

    with open(txt_path, 'r') as file:
        lines = file.readlines()

    absolute_boxes = []
    for line in lines:
        data = line.strip().split()
        class_id = int(data[0])
        center_x_rel = float(data[1])
        center_y_rel = float(data[2])
        width_rel = float(data[3])
        height_rel = float(data[4])

        center_x_abs = center_x_rel * width
        center_y_abs = center_y_rel * height
        x1_abs = center_x_abs - (width_rel * width / 2)
        y1_abs = center_y_abs - (height_rel * height / 2)
        x2_abs = center_x_abs + (width_rel * width / 2)
        y2_abs = center_y_abs + (height_rel * height / 2)

        absolute_boxes.append([(class_id, x1_abs, y1_abs, x2_abs, y2_abs)])

    return absolute_boxes


# 图片和标签文件夹路径
images_folder = 'C:XXXX'  #  你图片的路径
output_folder = 'C:/XXXX'  # 标签的路径
labels_folder = 'C:XXXX'  #  绘制好图片保存的路径




def box_label(img, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False):
    """Add one xyxy box to image with label."""
    if isinstance(box, torch.Tensor):
        box = box.tolist()
    lw = max(round(sum(img.shape) / 2 * 0.003), 2)
    tf = max(lw - 1, 1)  # font thickness
    sf = lw / 3  # font scale
    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    cv2.rectangle(img, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA)
    w, h = cv2.getTextSize(label, 0, fontScale=sf, thickness=tf)[0]  # text width, height
    outside = p1[1] - h >= 3
    p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
    cv2.rectangle(img, p1, p2, color, -1, cv2.LINE_AA)  # filled
    cv2.putText(
        img,
        label,
        (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
        0,
        sf,
        txt_color,
        thickness=tf,
        lineType=cv2.LINE_AA,
    )



# 遍历图片文件夹
for filename in os.listdir(images_folder):
    if filename.endswith('.jpg') or filename.endswith('.png'):
        image_path = os.path.join(images_folder, filename)

        boxes = convert_to_absolute_coords(image_path, labels_folder)
        print(f"Image: {filename}, Boxes: {boxes}")
        img = cv2.imread(image_path)


        for d in reversed(boxes):
            for i in d:
                c= int(i[0])
                name =  names[c]
                label = f"{name} "
                box = i[1:5]
                box = torch.tensor(box, dtype=torch.float32)
                box_label(img,box, label, color=colors(c, True), rotated=False)
        save_path = os.path.join(output_folder, filename)
        cv2.imwrite(save_path, img)

只需要更改自己类别的映射和对应文件的位置即可
生成的图片如下
    

  • 9
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值