YOLOV5标注信息显示到原图

 YOLO格式数据集

import os
import cv2
import numpy as np
from pathlib import Path
from ultralytics.utils.plotting import Annotator  # 用于绘制边界框和标签
from utils.general import xywh2xyxy  # YOLOv5 提供的工具函数，用于坐标转换
from utils.plots import colors  # YOLOv5 提供的工具类，用于获取颜色


def read_yolo_labels(label_path):
    """
    读取 YOLO 格式的标注文件 (.txt)，并返回标注信息。
    
    Args:
        label_path (str): YOLO 标注文件路径 (.txt)。
    
    Returns:
        list: 标注信息，每个元素为 `[class_id, x_center, y_center, width, height]`。
    """
    annotations = []
    with open(label_path, "r") as f:
        for line in f.readlines():
            parts = line.strip().split()
            if len(parts) == 5:  # 确保格式正确
                annotations.append([int(parts[0])] + [float(x) for x in parts[1:]])
    return annotations


def yolo_to_xyxy(annotations, img_width, img_height):
    """
    将 YOLO 的归一化坐标转换为像素坐标。
    
    Args:
        annotations (list): YOLO 格式标注信息，每个元素为 `[class_id, x_center, y_center, width, height]`。
        img_width (int): 图片宽度。
        img_height (int): 图片高度。
    
    Returns:
        list: 转换后的标注信息，每个元素为 `[class_id, x1, y1, x2, y2]`。
    """
    converted_annotations = []
    for ann in annotations:
        class_id, x_center, y_center, width, height = ann
        x1 = int((x_center - width / 2) * img_width)
        y1 = int((y_center - height / 2) * img_height)
        x2 = int((x_center + width / 2) * img_width)
        y2 = int((y_center + height / 2) * img_height)
        converted_annotations.append([class_id, x1, y1, x2, y2])
    return converted_annotations


def draw_yolo_annotations(image, annotations, class_names, save_path=None):
    """
    将 YOLO 格式的标注绘制到图片上。
    
    Args:
        image (numpy.ndarray): 输入图片 (BGR 格式)。
        annotations (list): 标注信息，每个元素为 `[class_id, x1, y1, x2, y2]`。
        class_names (list): 类别名称列表。
        save_path (str): 保存路径，若为 None 则不保存。
    
    Returns:
        numpy.ndarray: 绘制了标注信息的图片。
    """
    annotator = Annotator(image, line_width=2, font_size=16, pil=True)  # 初始化绘制工具
    for ann in annotations:
        class_id, x1, y1, x2, y2 = ann
        label = f"{class_names[class_id]}"  # 显示类别名称
        color = colors(class_id)  # 根据类别 ID 获取颜色
        annotator.box_label([x1, y1, x2, y2], label, color=color)  # 绘制边界框和标签
    
    if save_path:
        annotator.im.save(save_path)  # 保存图片到本地

    # 将 PIL 图像转换为 NumPy 格式 (BGR 格式以兼容 OpenCV)
    annotated_image = np.array(annotator.im)[:, :, :]  # 从 RGB 转换为 BGR

    return annotated_image



def visualize_yolo_annotations(image_path, label_path, class_names, save_path=None):
    """
    主流程：读取图片和 YOLO 标注信息，并绘制标注。
    
    Args:
        image_path (str): 图片路径。
        label_path (str): YOLO 标注文件路径。
        class_names (list): 类别名称列表。
        save_path (str): 保存路径，若为 None 则不保存。
    """
    # 加载图片
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"图像未找到: {image_path}")
    img_height, img_width = image.shape[:2]

    # 读取 YOLO 标注信息
    annotations = read_yolo_labels(label_path)

    # 转换 YOLO 格式为像素坐标
    pixel_annotations = yolo_to_xyxy(annotations, img_width, img_height)

    # 绘制标注信息
    annotated_image = draw_yolo_annotations(image, pixel_annotations, class_names, save_path)

    # 显示图片（可选）
    cv2.imshow("Annotated Image", annotated_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

VOC 格式数据集

import os
import cv2
import xml.etree.ElementTree as ET
from pathlib import Path
from ultralytics.utils.plotting import Annotator  # YOLOv5 的绘图工具
from utils.plots import colors  # YOLOv5 的绘图颜色工具


def parse_voc_xml(xml_path):
    """
    解析 Pascal VOC 格式的 XML 文件，提取类别和边界框信息。

    Args:
        xml_path (str): Pascal VOC 格式的 XML 文件路径。

    Returns:
        dict: 包含图片路径、图片尺寸和标注信息的字典。
    """
    tree = ET.parse(xml_path)
    root = tree.getroot()

    # 提取图片相关信息
    image_info = {
        "filename": root.find("filename").text,
        "size": {
            "width": int(root.find("size/width").text),
            "height": int(root.find("size/height").text),
            "depth": int(root.find("size/depth").text),
        },
        "objects": [],
    }

    # 提取目标对象信息
    for obj in root.findall("object"):
        obj_data = {
            "name": obj.find("name").text,
            "bndbox": {
                "xmin": int(obj.find("bndbox/xmin").text),
                "ymin": int(obj.find("bndbox/ymin").text),
                "xmax": int(obj.find("bndbox/xmax").text),
                "ymax": int(obj.find("bndbox/ymax").text),
            },
        }
        image_info["objects"].append(obj_data)

    return image_info


def draw_voc_annotations(image, annotations, save_path=None):
    """
    在图片上绘制 Pascal VOC 格式的标注。

    Args:
        image (numpy.ndarray): 图片 (BGR 格式)。
        annotations (list): 标注信息，每个元素包含类别名称和边界框坐标。
        save_path (str): 保存路径，若为 None 则不保存。

    Returns:
        numpy.ndarray: 绘制了标注信息的图片。
    """
    annotator = Annotator(image, line_width=2, font_size=16, pil=True)  # 初始化绘制工具
    for ann in annotations:
        label = ann["name"]  # 类别名称
        box = [ann["bndbox"]["xmin"], ann["bndbox"]["ymin"], ann["bndbox"]["xmax"], ann["bndbox"]["ymax"]]
        color = colors(hash(label) % len(colors.palette))  # 根据类别名称哈希获取颜色
        annotator.box_label(box, label, color=color)  # 绘制边界框和类别标签

    if save_path:
        annotator.im.save(save_path)  # 保存图片到本地

    # 将 PIL 图像转换为 NumPy 格式 (BGR 格式以兼容 OpenCV)
    annotated_image = np.array(annotator.im)[:, :, :]  # 从 RGB 转换为 BGR

    return annotated_image

def visualize_voc_annotations(image_dir, xml_path, save_path=None):
    """
    主流程：解析 Pascal VOC 标注文件并绘制到图片上。

    Args:
        image_dir (str): 图片存储路径。
        xml_path (str): Pascal VOC 格式的 XML 文件路径。
        save_path (str): 保存路径，若为 None 则不保存。
    """
    # 解析 XML 文件
    annotation_data = parse_voc_xml(xml_path)

    # 加载图片
    image_path = os.path.join(image_dir, annotation_data["filename"])
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"图像未找到: {image_path}")

    # 提取标注信息
    annotations = annotation_data["objects"]

    # 绘制标注到图片上
    annotated_image = draw_voc_annotations(image, annotations, save_path)

    # 显示图片（可选）
    cv2.imshow("Annotated Image", annotated_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 文件路径
image_path = "data/images/example.jpg"
label_path = "data/labels/example.txt"
save_path = "output/annotated_example.jpg"

# 类别名称
class_names = ["person", "car", "cat"]

# 可视化 YOLO 标注
visualize_yolo_annotations(image_path, label_path, class_names, save_path)