python扩充yolov8数据集（图像和标签同时扩充）

'''扩充单张图像'''
import os
import cv2

def flip_yolo_labels(labels, img_width, img_height):
    flipped_labels = []
    for label in labels:
        class_id, x_center, y_center, width, height = label.split()
        x_center = 1 - float(x_center)  # 水平翻转中心点 x 坐标
        flipped_labels.append(f"{class_id} {x_center} {y_center} {width} {height}\n")
    return flipped_labels


def flip_image_and_labels(image_folder, label_folder, output_folder):
    for filename in os.listdir(image_folder):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            image_path = os.path.join(image_folder, filename)
            label_path = os.path.join(label_folder, os.path.splitext(filename)[0] + '.txt')

            # 读取并翻转图像
            img = cv2.imread(image_path)
            flipped_img = cv2.flip(img, 1)  # 水平翻转

            # 读取并翻转标签
            with open(label_path, 'r') as file:
                labels = file.readlines()
            flipped_labels = flip_yolo_labels(labels, img.shape[1], img.shape[0])

            # 保存翻转后的图像和标签
            cv2.imwrite(os.path.join(output_folder, 'flipped_' + filename), flipped_img)
            with open(os.path.join(output_folder, 'flipped_' + os.path.splitext(filename)[0] + '.txt'), 'w') as file:
                file.writelines(flipped_labels)


# 设置文件夹路径
image_folder = 'D:\WORK\work1\labelss\\allimage'
label_folder = 'D:\WORK\work1\labelss\\alllabels'
output_folder = 'D:\WORK\work1\labelss\kuochong-images'

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

# 执行翻转操作
flip_image_and_labels(image_folder, label_folder, output_folder)


'''每张图像可以设置扩充多少张'''
# import os
# import cv2
#
# def flip_yolo_labels(labels, flip_mode, img_width, img_height):
#     flipped_labels = []
#     for label in labels:
#         parts = label.split()
#         class_id = parts[0]
#         x_center, y_center, width, height = map(float, parts[1:])
#
#         if flip_mode == 1 or flip_mode == -1:  # 水平翻转
#             x_center = 1 - x_center
#         if flip_mode == 0 or flip_mode == -1:  # 垂直翻转
#             y_center = 1 - y_center
#
#         flipped_labels.append(f"{class_id} {x_center} {y_center} {width} {height}\n")
#
#     return flipped_labels
#
# def flip_image_and_labels(image_path, label_path, output_image_path, output_label_path, flip_mode):
#     # 读取图像
#     img = cv2.imread(image_path)
#     height, width = img.shape[:2]
#
#     # 翻转图像
#     flipped_img = cv2.flip(img, flip_mode)
#
#     # 读取并翻转标签
#     with open(label_path, 'r') as file:
#         labels = file.readlines()
#     flipped_labels = flip_yolo_labels(labels, flip_mode, width, height)
#
#     # 保存翻转后的图像和标签
#     cv2.imwrite(output_image_path, flipped_img)
#     with open(output_label_path, 'w') as file:
#         file.writelines(flipped_labels)
#
# def augment_dataset(image_folder, label_folder, output_folder, num_copies=3):#num_copies设置扩充多少张
#     if not os.path.exists(output_folder):
#         os.makedirs(output_folder)
#
#     for filename in os.listdir(image_folder):
#         if filename.endswith('.jpg') or filename.endswith('.png'):
#             image_path = os.path.join(image_folder, filename)
#             label_path = os.path.join(label_folder, os.path.splitext(filename)[0] + '.txt')
#
#             # 为每个图像生成指定数量的扩充版本
#             for i in range(num_copies):
#                 output_image_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}_flip_{i}.jpg")
#                 output_label_path = os.path.join(output_folder, f"{os.path.splitext(filename)[0]}_flip_{i}.txt")
#                 flip_mode = i - 1  # -1: 水平+垂直翻转, 0: 垂直翻转, 1: 水平翻转
#                 flip_image_and_labels(image_path, label_path, output_image_path, output_label_path, flip_mode)
#
# # 设置文件夹路径
# image_folder = 'path_to_your_image_folder'
# label_folder = 'path_to_your_label_folder'
# output_folder = 'path_to_output_folder'
#
# # 执行数据集扩充
# augment_dataset(image_folder, label_folder, output_folder)