YOLO实践

一. 环境安装

参考视频

1. Pytorch环境安装细节

• pytorch安装：一个单独的环境中，能使用pip就尽量使用pip，实在有问题的情况，例如没有合适的编译好的系统版本的安装包，再使用conda进行安装，不要来回混淆
• CUDA是否要安装：如果只需要训练、简单推理，则无需单独安装cUDA，直接安装pytorch；如果有部署需求，例如导出TensorRT模型，则需要进行CUDA安装
• Pytorch安装注意事项:必须使用官网的命令进行安装，否则安装的是cpu的版本

2. 正确使用windows终端：

• 使用cmd,而不是powershell（无法激活环境）。
• 在其他软件，如Pycharm,vscode中也要注意！

4. 可能出现的问题

• Arial.ttf字体文件无法下载
  • 手动下载，放到对应的位置，windows下的目录是：～/AppData/Roaming/Ultralytics
• 页面文件太小,无法完成操作
  • 调整训练参数中的workers，设置为1
  • 修改虚拟内存，将环境安装位置所在的盘，设置一个较大的参数
• 'Upsample' object has no attribute 'recompute _scale_factor'
  • pytorch版本过高导致，可以选择降版本，1.8.2目前是不会报错的版本
  • 如不想降低版本，可以修改pytorch源码，打开报错的unsampling.py，删除
    recompute_scale_factor这个参数

设置电脑虚拟内存：

二 .数据集制作

数据转换代码数据格式转化关键根据相应的文件格式解析处class以及bbox

1. voc2yolo

def convert_xml_to_txt_format(xml_file_path: str, class_list: list) -> list:
    """
    将单个xml文件转换为txt格式。

    参数:
    xml_file_path (str): XML文件的路径。
    class_list (list): 包含所有类别的列表。

    返回:
    list: 包含YOLO格式的标签列表。
    """
    # 根据class_list构建class_to_id_map
    class_to_id_map = {class_name: index for index, class_name in enumerate(class_list)}

    tree = ET.parse(xml_file_path)
    root = tree.getroot()

    width = int(root.find('.//size/width').text)
    height = int(root.find('.//size/height').text)

    txt_format = []
    for obj in root.findall('.//object'):
        class_name = obj.find('name').text
        if class_name not in class_to_id_map:
            continue
        class_id = class_to_id_map[class_name]
        bbox = obj.find('bndbox')
        xmin = float(bbox.find('xmin').text)
        ymin = float(bbox.find('ymin').text)
        xmax = float(bbox.find('xmax').text)
        ymax = float(bbox.find('ymax').text)

        x_center = (xmin + xmax) / 2 / width
        y_center = (ymin + ymax) / 2 / height
        bbox_width = (xmax - xmin) / width
        bbox_height = (ymax - ymin) / height

        txt_format.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")

    return txt_format



def write_yolo_format_labels_to_file(yolo_labels_dir: str, yolo_label_file: str, yolo_content: list):
    """
    将YOLO格式的标签写入到文件。

    参数:
    yolo_labels_dir (str): YOLO标签目录的路径。
    yolo_label_file (str): YOLO标签文件的名称。
    yolo_content (list): 包含YOLO格式标签的列表。
    """
    yolo_label_path = os.path.join(yolo_labels_dir, yolo_label_file)
    with open(yolo_label_path, 'w') as f:
        for line in yolo_content:
            f.write(line + '\n')


def convert_voc2yolo(voc_annotations_dir: str, yolo_labels_dir: str, class_list: list):
    """
    将VOC格式的XML注释转换为YOLO格式的注释。

    参数:
    voc_annotations_dir (str): VOC注释目录的路径。
    yolo_labels_dir (str): YOLO标签目录的路径。
    class_list (list): 类别名称的list。
    """
    if not os.path.exists(yolo_labels_dir):
        os.makedirs(yolo_labels_dir)

    for xml_file in os.listdir(voc_annotations_dir):
        if xml_file.endswith('.xml'):
            yolo_label_file = xml_file.replace('.xml', '.txt')
            xml_file_path = os.path.join(voc_annotations_dir, xml_file)

            yolo_content = convert_xml_to_txt_format(xml_file_path,class_list)
            write_yolo_format_labels_to_file(yolo_labels_dir, yolo_label_file, yolo_content)
    with open(os.path.join(yolo_labels_dir,"classes.txt"),"w",encoding='utf-8') as f:
        for c in class_list:
            f.write(c+'\n')
if __name__ == "__main__":
    dataset_dir = "./dataset"
    labels_voc_dir = "labels_voc"
    labels_yolo_dir = "labels_yolo"
    class_names = ["橘子", "香蕉", "草莓"]
    convert_voc2yolo(labels_voc_dir,labels_yolo_dir,class_names)

2. yolo2voc

def convert_txt_to_xml_format(txt_file_path: str, voc_file_path: str, class_names: list, images_dir: str):
    """
    将单个txt文件转换为xml格式。

    参数:
    txt_file_path (str): YOLO标签文件txt的路径。
    voc_file_path (str): XML文件的路径。
    class_names (list): 包含类别名称的列表。
    images_dir (str): 包含图像文件的目录路径。
    """
    with open(txt_file_path, 'r') as file:
        lines = file.readlines()

    # 获取txt文件的文件名，然后去掉后缀
    txt_filename = os.path.basename(txt_file_path)
    txt_filename_without_ext = os.path.splitext(txt_filename)[0]

    # 使用glob库在images_dir下搜索与txt文件同名的图片文件
    pattern = os.path.join(images_dir, f"{txt_filename_without_ext}.*")
    matching_files = glob.glob(pattern)

    if not matching_files:
        raise FileNotFoundError(f"No matching image found for {txt_filename_without_ext} in {images_dir}")

    img_file = matching_files[0]  # 假设找到的第一个匹配文件就是正确的图片

    # 从与YOLO标签同名的图像文件中读取宽度和高度
    img = Image.open(img_file)
    width, height = img.size

    root = ET.Element("annotation")

    folder = ET.SubElement(root, "folder")
    folder.text = os.path.basename(images_dir)

    filename = ET.SubElement(root, "filename")
    filename.text = os.path.basename(img_file)

    path = ET.SubElement(root, "path")
    path.text = os.path.abspath(img_file)

    source = ET.SubElement(root, "source")
    database = ET.SubElement(source, "database")
    database.text = "Unknown"

    size = ET.SubElement(root, "size")
    ET.SubElement(size, "width").text = str(width)
    ET.SubElement(size, "height").text = str(height)
    ET.SubElement(size, "depth").text = "3"

    segmented = ET.SubElement(root, "segmented")
    segmented.text = "0"

    for line in lines:
        values = line.strip().split()
        if len(values) < 5:
            continue

        class_id, x_center, y_center, bbox_width, bbox_height = map(float, values)

        object = ET.SubElement(root, "object")
        name = ET.SubElement(object, "name")
        name.text = class_names[int(class_id)]

        pose = ET.SubElement(object, "pose")
        pose.text = "Unspecified"

        truncated = ET.SubElement(object, "truncated")
        truncated.text = "0"

        difficult = ET.SubElement(object, "difficult")
        difficult.text = "0"

        bndbox = ET.SubElement(object, "bndbox")
        xmin = int((x_center - bbox_width / 2) * width)
        ymin = int((y_center - bbox_height / 2) * height)
        xmax = int((x_center + bbox_width / 2) * width)
        ymax = int((y_center + bbox_height / 2) * height)
        ET.SubElement(bndbox, "xmin").text = str(xmin)
        ET.SubElement(bndbox, "ymin").text = str(ymin)
        ET.SubElement(bndbox, "xmax").text = str(xmax)
        ET.SubElement(bndbox, "ymax").text = str(ymax)

    # 使用prettify函数来格式化输出
    pretty_xml_as_string = prettify(root)
    with open(voc_file_path, 'w', encoding='utf-8') as f:
        f.write(pretty_xml_as_string)


def prettify(elem):
    from xml.dom import minidom
    """Return a pretty-printed XML string for the Element."""
    rough_string = ET.tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="	")


def convert_yolo2voc(yolo_labels_dir: str, voc_labels_dir: str, images_dir: str, class_list: list):
    """
    将YOLO格式的标签文件夹转换为VOC格式的XML注释文件夹。

    参数:
    yolo_labels_dir (str): YOLO标签文件夹的路径。
    voc_labels_dir (str): 转换后的VOC XML文件夹的路径。
    images_dir (str): 包含图像文件的目录路径。
    class_names (list): 包含类别名称的列表。
    """
    # 确保voc_labels_dir存在
    if not os.path.exists(voc_labels_dir):
        os.makedirs(voc_labels_dir)

    # 遍历YOLO标签文件夹
    for label_file in os.listdir(yolo_labels_dir):
        if label_file.endswith(".txt"):
            yolo_file_path = os.path.join(yolo_labels_dir, label_file)
            voc_file_path = os.path.join(voc_labels_dir, label_file.replace(".txt", ".xml"))
            convert_txt_to_xml_format(yolo_file_path, voc_file_path, class_list, images_dir)
if __name__ == "__main__":
    dataset_dir = "./dataset"
    labels_voc_dir = "labels_voc"
    labels_yolo_dir = "labels_yolo"
    class_names = ["橘子", "香蕉", "草莓"]
    convert_yolo2voc(labels_yolo_dir,labels_voc_dir,"imgs",class_names)
    

4. Json转Yolo

提示可以采用 data_df = pd.read_json(json_path),然后利用pandas的语法进行提取bbox和class_id

5. 划分数据集

def split_dataset(images_origin_dir: str, labels_origin_dir: str, dataset_dir: str, train_ratio: float = 0.8):
    """
    将原始数据集分割为训练集和验证集，并将相应的文件复制到新的文件夹中。

    参数:
    images_origin (str): 存放所有图片文件夹路径
    labels_origin (str): 存放所有标签文件夹路径。
    dataset_dir (str): 新的数据集根目录。
    train_ratio (float, optional): 训练集的比例，默认为 0.8。
    """
    images_dir = os.path.join(dataset_dir, "images")
    labels_dir = os.path.join(dataset_dir, "labels")
    train_dir = os.path.join(images_dir, "train")
    val_dir = os.path.join(images_dir, "val")

    # 创建训练集和验证集的文件夹
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(os.path.join(labels_dir, "train"), exist_ok=True)
    os.makedirs(os.path.join(labels_dir, "val"), exist_ok=True)

    image_files=[]
    for file in os.listdir(images_origin_dir):
        if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg"):
            image_files.append(os.path.join(images_origin_dir, file))

    # 使用 random.sample 函数进行随机划分
    random.seed(42)  # 设置随机种子
    train_images = random.sample(image_files, k=int(len(image_files) * train_ratio))
    val_images = [f for f in image_files if f not in train_images]

    # 复制文件到训练集和验证集文件夹
    for subset, images in [('train', train_images), ('val', val_images)]:
        for image_file in images:
            # 复制图片文件
            target_image_path = os.path.join(images_dir, subset, os.path.basename(image_file))
            shutil.copyfile(image_file, target_image_path)

            # 复制标签文件
            label_file = os.path.splitext(os.path.basename(image_file))[0] + '.txt'
            label_path = os.path.join(labels_origin_dir, label_file)
            if os.path.exists(label_path):
                target_label_path = os.path.join(labels_dir, subset, label_file)
                shutil.copyfile(label_path, target_label_path)

    print(f"训练集大小: {len(train_images)}")
    print(f"验证集大小: {len(val_images)}")
 if __name__ == "__main__":
    dataset_dir = "./dataset"
    labels_voc_dir = "labels_voc"
    labels_yolo_dir = "labels_yolo"
    class_names = ["橘子", "香蕉", "草莓"]
    split_dataset("imgs",labels_yolo_dir,dataset_dir)

三.模型使用

1. YOLOV5 模型调用（本地）

# Model
model = torch.hub.load('.', 'custom', path=r'best.pt',source='local')

.: 本地的yolov5的根目录
source="local"：从本地调用该模型

2. YOLOV8 模型调用

# Load the YOLOv8 model
model = YOLO("yolov8n.pt")  # Make sure the model file is in the correct path

如果本地不存在yolov8n.pt，则从网上下载

四. 结果处理

.xyxy()：即为[x_min,y_min,x_max,y_max]

1. YOLOV5 参考博客

import torch
model = torch.hub.load('.', 'yolov5s', source='local') 
im = r'data\images\bus.jpg'  # file, Path, PIL.Image, OpenCV, nparray, list`
results = model(im)  # inference
# results.crop()  # or .show(), .save(), .crop(), .pandas(), etc.

results是一个Detections对象，主要有以下方法：

# results.crop()  # or .show(), .save(), .crop(), .pandas(), etc.
results.xyxy[0]   #Tensor类型
results.pandas().xyxy[0]  #转为pandas格式
results.print()  #控制台显示
frame = results.render()[0]  # 将模型的输出绘制回帧上

另外：想要个性化输出，限定类别

# 将检测结果转换为Pandas DataFrame格式
detections_df = results.pandas().xyxy[0]

# 筛选出标签名为'person'的所有检测结果
person_detections = detections_df[detections_df['name'] == 'person'].to_numpy()

# 遍历检测到的人
for detection in person_detections:
    label_name = detection[6]  # 获取标签名
    bbox = detection[:4].astype('int')  # 获取边界框坐标并转换为整数

    # 在图像上绘制边界框
    cv.rectangle(image_temp, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 3)

    # 在边界框左上角添加标签名
    cv.putText(image_temp, label_name, (bbox[0] - 10, bbox[1] - 10), cv.FONT_ITALIC, 1, (0, 255, 0), 2)

2. YOLOV8

建议：ultalytics源码安装，在ultralytics的根目录下pip install -e . -e指的是本地可编辑，下载该库时不会直接下载到site-packages下，而是以链接的方式链接到源码的根目录下

from ultralytics import YOLO
model=YOLO("yolov8n-seg.pt")
results=model(r"images\bus.jpg")

这里的results是list类型。

for idx, result in enumerate(results):
        boxes = result.boxes  # Boxes object for bounding box outputs
        masks = result.masks  # Masks object for segmentation masks outputs
        keypoints = result.keypoints  # Keypoints object for pose outputs
        probs = result.probs  # Probs object for classification outputs
        obb = result.obb  # Oriented boxes object for OBB outputs

        if len(boxes.cls) == 0:
            continue
        
        xyxy = boxes.xyxy.data.cpu().numpy().round()
        cls = boxes.cls.data.cpu().numpy().round()
        conf = boxes.conf.data.cpu().numpy()

另外:如果想要个性化输出，可以使用 patched_yolo_infer库（我只验证了YOLOV8,其它的YOLO系列不知道）

import cv2
from ultralytics import YOLO
from patched_yolo_infer import visualize_results_usual_yolo_inference


# Load the image
img_path = 'images/bus.jpg'
img = cv2.imread(img_path)

img=visualize_results_usual_yolo_inference(
        img,
        model,
        conf=0.4,
        iou=0.7,
        show_classes_list=[0], #Whether to perform instance segmentation. Default is False.
        segment=True,          #是否分割
        thickness=5,
        show_boxes=False,
        fill_mask=False,      
        alpha=0.7,			#The transparency of filled masks. Default is 0.3.
        show_class=False,
        delta_colors=25,     #The random seed offset for color variation. Default is 0.
        inference_extra_args={'retina_masks':True}, #increase the accuracy of the contours
        return_image_array=True
    )

附:Pyme组件

import cv2 as cv
from PIL import ImageTk, Image
class VideoPlayer():
    def __init__(self,elementName,video_source=0):
        super().__init__()
        self.video_source = video_source
        self.vid = cv.VideoCapture(self.video_source)
        self.label= Fun.GetElement(uiName,elementName)  #如果想要替换为tk，这里直接注入label对象
        self.delay = 15
        self.paused = False  # 暂停状态标志
        self.update()
        
    def update(self):
        # 检查是否成功读取帧
        if self.vid.isOpened():
            if not self.paused:
                ret, frame = self.vid.read()
                if ret:
                    frame = self.process_frame(frame)
                    # 转换为Image
                    self.photo = ImageTk.PhotoImage(image=Image.fromarray(frame))
                    self.label.config(image=self.photo)
                else:
                    self.vid.release()
                    return
        self.label.after(self.delay, self.update)
        
    def process_frame(self, frame):
        #convert BGR to RGB
        frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
        return frame
    
    def pause(self):
        self.paused = not self.paused
        return self.paused
    
    def stop(self):
        self.vid.release()