[目标检测] 训练之前要做什么

背景：训练一个Yolo8模型，在训练之前，数据集的处理是影响效果的关键因素。

Step1 定义规则

什么是人/车，比如人的话可能是站着的人，如果是骑电动车/自行车就不算是人。

Step2 收集数据集

1. 自己标注。如果是自己标注，那么根据上述的规则进行清洗。

2. 采集他人的数据集，标注好的。那么最好是能可视化一下标签。下面是可视化代码。

import os
import cv2

def visualize_yolo_boxes(image_folder, label_folder, output_folder):
    # 确保输出文件夹存在
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 遍历图片文件夹
    for image_name in os.listdir(image_folder):
        # 获取图片路径
        image_path = os.path.join(image_folder, image_name)
        # 获取对应的标签路径
        label_path = os.path.join(label_folder, os.path.splitext(image_name)[0] + '.txt')

        # 如果标签文件不存在，跳过
        if not os.path.exists(label_path):
            print(f"标签文件不存在: {label_path}")
            continue

        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图片: {image_path}")
            continue

        # 获取图片的宽度和高度
        img_height, img_width = image.shape[:2]

        # 读取标签文件
        with open(label_path, 'r') as f:
            lines = f.readlines()

        # 遍历每个标签
        for line in lines:
            parts = line.strip().split()
            if len(parts) != 5:
                continue

            # 解析YOLO格式的标签
            class_id = int(parts[0])
            x_center = float(parts[1]) * img_width
            y_center = float(parts[2]) * img_height
            width = float(parts[3]) * img_width
            height = float(parts[4]) * img_height

            # 计算边界框的左上角和右下角坐标
            x1 = int(x_center - width / 2)
            y1 = int(y_center - height / 2)
            x2 = int(x_center + width / 2)
            y2 = int(y_center + height / 2)

            # 在图片上绘制边界框
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, f'Class {class_id}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        # 保存可视化结果
        output_path = os.path.join(output_folder, image_name)
        cv2.imwrite(output_path, image)
        print(f"已保存可视化结果: {output_path}")

# 示例用法
image_folder = r'.\train\images'
label_folder = r'.\train\labels'
output_folder = r'.\train\visualize'

visualize_yolo_boxes(image_folder, label_folder, output_folder)

Step3 修改标签

修改标签是指，比如要合并一些标签。已经标注好的数据集，比如把车子分为了truck，car，bus之类的，但是，我们都统称为car。所以要进行一些合并。下面是合并代码。

import os
import glob

def process_line(line):
    """Process a single line of text according to the mapping rules."""
    parts = line.strip().split()
    if not parts:  # Skip empty lines
        return None
        
    try:
        id_num = int(parts[0])
        # Keep only specified IDs
        if id_num not in [0, 1, 2, 6]:
            return None
            
        # Apply mapping: 0,1 -> 0 and 3,4,5,8 -> 1
        if id_num in [0, 1, 6]:
            new_id = 1
        else:  # id in [3,4,5,8]
            new_id = 0
            
        # Replace first number and keep rest of the line the same
        return f"{new_id} {' '.join(parts[1:])}\n"
        
    except ValueError:
        return None

def process_file(input_path, output_path):
    """Process a single text file and save to output directory."""
    try:
        with open(input_path, 'r') as infile:
            lines = infile.readlines()
            
        # Process lines and filter out None results
        processed_lines = [
            processed for line in lines
            if (processed := process_line(line)) is not None
        ]
        
        # Create output directory if it doesn't exist
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # Write processed lines to output file
        with open(output_path, 'w') as outfile:
            outfile.writelines(processed_lines)
            
        print(f"Processed {input_path} -> {output_path}")
        
    except Exception as e:
        print(f"Error processing {input_path}: {str(e)}")

def process_directory(input_dir, output_dir):
    """Process all .txt files in the input directory."""
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Find all .txt files in input directory
    txt_files = glob.glob(os.path.join(input_dir, "*.txt"))
    
    for input_path in txt_files:
        # Create corresponding output path
        relative_path = os.path.relpath(input_path, input_dir)
        output_path = os.path.join(output_dir, relative_path)
        
        # Process the file
        process_file(input_path, output_path)

# Example usage
if __name__ == "__main__":
    input_directory = r"F:\1-dataset\raw\add_vehicle_person\combine\vehicle_class\valid\labels"  # Replace with your input directory
    output_directory = r"F:\1-dataset\raw\add_vehicle_person\combine\vehicle_class\valid\labels_person_car"  # Replace with your output directory
    
    process_directory(input_directory, output_directory)