背景：训练一个Yolo8模型，在训练之前，数据集的处理是影响效果的关键因素。
Step1 定义规则
什么是人/车，比如人的话可能是站着的人，如果是骑电动车/自行车就不算是人。
Step2 收集数据集
1. 自己标注。如果是自己标注，那么根据上述的规则进行清洗。
2. 采集他人的数据集，标注好的。那么最好是能可视化一下标签。下面是可视化代码。
import os
import cv2def 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_widthy_center = float(parts[2]) * img_heightwidth = float(parts[3]) * img_widthheight = 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 globdef process_line(line):"""Process a single line of text according to the mapping rules."""parts = line.strip().split()if not parts: # Skip empty linesreturn Nonetry:id_num = int(parts[0])# Keep only specified IDsif id_num not in [0, 1, 2, 6]:return None# Apply mapping: 0,1 -> 0 and 3,4,5,8 -> 1if id_num in [0, 1, 6]:new_id = 1else: # id in [3,4,5,8]new_id = 0# Replace first number and keep rest of the line the samereturn f"{new_id} {' '.join(parts[1:])}\n"except ValueError:return Nonedef 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 resultsprocessed_lines = [processed for line in linesif (processed := process_line(line)) is not None]# Create output directory if it doesn't existos.makedirs(os.path.dirname(output_path), exist_ok=True)# Write processed lines to output filewith 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 existos.makedirs(output_dir, exist_ok=True)# Find all .txt files in input directorytxt_files = glob.glob(os.path.join(input_dir, "*.txt"))for input_path in txt_files:# Create corresponding output pathrelative_path = os.path.relpath(input_path, input_dir)output_path = os.path.join(output_dir, relative_path)# Process the fileprocess_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 directoryoutput_directory = r"F:\1-dataset\raw\add_vehicle_person\combine\vehicle_class\valid\labels_person_car" # Replace with your output directoryprocess_directory(input_directory, output_directory)
