YOLO数据集生成器代码

import os
import random
import sys
from PIL import Image

Noo = 6  # Number of objects
Momor = 5  # Multiples of magnification or reduction
Tffp = 797  # The first few pictures

# 请仔细看完代码在谨慎修改LsY与Ley参数
LsX = 0  # Limit the start point X-axis
LeX = 200  # Limit the end point X-axis
LsY = 0  # Limit the start point Y-axis
LeY = 1000  # Limit the end point Y-axis

LsXf = 0  # Limit the X-axis of the starting point of the foreground
LeXf = 200
LsYf = 1400
LeYf = 2600

Fd = 20  # Frame distance

GXd = 30  # Graph X-axis distance

foreground_folder = "E:\\document\\YOLO_data_set_genneration_tool\\foreground\\yellow"
background_folder = "E:\\document\\YOLO_data_set_genneration_tool\\background"
output_folder = "E:\\document\\YOLO_data_set_genneration_tool\\3"  # 确保此文件夹存在


class Rect:
    def __init__(self, x1, y1, x2, y2):
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2

    def cross_line(self, other):
        w1 = self.x2 - self.x1
        h1 = self.y2 - self.y1
        w2 = other.x2 - other.x1
        h2 = other.y2 - other.y2
        w = abs((self.x1 + self.x2) / 2 - (other.x1 + other.x2) / 2)
        h = abs((self.y1 + self.y2) / 2 - (other.y1 + other.y2) / 2)
        return w < (w1 + w2) / 2 and h < (h1 + h2) / 2


# 获取图片列表
foreground_images = [os.path.join(foreground_folder, file) for file in os.listdir(foreground_folder) if
                     file.endswith('.jpg') or file.endswith('.png')]
background_images = [os.path.join(background_folder, file) for file in os.listdir(background_folder) if
                     file.endswith('.jpg') or file.endswith('.png')]


def progress_bar(finish_tasks_number, tasks_number):
    percentage = round(finish_tasks_number / tasks_number * 100)
    print("\r进度: {}%: ".format(percentage), "▓" * (percentage // 2), end="")
    sys.stdout.flush()


def save_positions_txt(output_folder, index, positions):
    with open(os.path.join(output_folder, f"positions_random_{index}.txt"), 'w') as file:
        for pos in positions:
            file.write(' '.join(map(str, pos)) + '\n')


# 主处理逻辑
for i, bg_image_path in enumerate(background_images):
    background_image = Image.open(bg_image_path)
    width, height = background_image.size
    positions_rects = []  # 新增列表用于存储Rect对象
    positions = []  # 用于存储前景图片的位置
    labels = []  # 存储前景图片的标签
    blacklist = []  # 新增黑名单列表
    current_x = 0  # 初始化当前X坐标为0，用于顺序排列

    # 特殊处理第一张前景图
    fg_image_path = random.choice(foreground_images)
    foreground_image = Image.open(fg_image_path)
    foreground_resized = foreground_image.resize((foreground_image.width * Momor, foreground_image.height * Momor))
    foreground_width, foreground_height = foreground_resized.size

    label = int(fg_image_path.split('/')[-1].split('.')[-2][-1])
    # 第一张前景图的初始位置
    offset_x = random.randint(LsX, LeX)
    offset_y = random.randint(LsY, LeY)
    rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)

    # 确保第一张图片的位置不超出边界
    # offset_x = min(offset_x, width - foreground_width)  # 垃圾代码
    # offset_y = min(offset_y, height - foreground_height)  # 垃圾代码

    background_image.paste(foreground_resized, (offset_x, offset_y),
                           foreground_resized.convert("RGBA").getchannel("A"))
    positions.append((offset_x, offset_y, foreground_width, foreground_height))
    positions_rects.append(rect1)
    labels.append(label)
    blacklist.append(rect1)
    current_x = offset_x + foreground_width + GXd  # 更新当前X坐标

    # 处理剩余的前景图
    for _ in range(Noo - 1):  # 减1是因为第一张已经处理过了
        fg_image_path = random.choice(foreground_images)
        foreground_image = Image.open(fg_image_path)
        foreground_resized = foreground_image.resize((foreground_image.width * Momor, foreground_image.height * Momor))
        foreground_width, foreground_height = foreground_resized.size

        label = int(fg_image_path.split('/')[-1].split('.')[-2][-1])

        # 顺序放置，从当前X坐标开始
        offset_x = current_x
        offset_y = random.randint(0, height - foreground_height)  # Y坐标随机

        # 确保前景图不会超出背景图
        if offset_x + foreground_width > width:
            offset_x = random.randint(LsXf, LeXf)
            offset_y = random.randint(LsYf, LeYf)
            # 检查是否与黑名单中的矩形相交，如果相交则调整Y坐标
            rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)
            j = 0
            # 下面纯纯傻逼代码
            while any(rect1.cross_line(blacklisted_rect) for blacklisted_rect in blacklist):
                j += 1
                if j == 1:
                    offset_y = random.randint(2000, height - foreground_height)#这行代码随时修改
                    rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)
                elif j == 2:
                    offset_y = random.randint(0, height - foreground_height)
                    offset_x += 400
                    rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)
                elif j == 3:
                    offset_y = random.randint(2000, height - foreground_height)
                    offset_x += 400
                    rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)

        # 检查是否与黑名单中的矩形相交，如果相交则调整Y坐标
        rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)
        while any(rect1.cross_line(blacklisted_rect) for blacklisted_rect in blacklist):
            offset_y = random.randint(0, height - foreground_height)
            rect1 = Rect(offset_x, offset_y, offset_x + foreground_width, offset_y + foreground_height)

        background_image.paste(foreground_resized, (offset_x, offset_y),
                               foreground_resized.convert("RGBA").getchannel("A"))
        positions.append((offset_x, offset_y, foreground_width, foreground_height))
        positions_rects.append(rect1)
        labels.append(label)
        blacklist.append(rect1)
        current_x += foreground_width  # 更新当前X坐标，为下一张图片做准备

    # 保存新图片
    output_path = os.path.join(output_folder, f"merged_random_{i + Tffp}.png")
    background_image.save(output_path)

    # 保存所有前景图片在背景图片上的位置信息到一个.txt文件
    # 修改为YOLO格式labels_random_0
    with open(os.path.join(output_folder, f"merged_random_{i + Tffp}.txt"), 'w') as file:
        for pos, label in zip(positions, labels):
            x, y, w, h = pos
            x = x - Fd
            y = y - Fd
            w = w + 2 * Fd
            h = h + 2 * Fd
            cx = x + w / 2
            cy = y + h / 2
            cx_normalized = cx / width
            cy_normalized = cy / height
            w_normalized = w / width
            h_normalized = h / height
            cx_normalized_rounded = round(cx_normalized, 6)
            cy_normalized_rounded = round(cy_normalized, 6)
            w_normalized_rounded = round(w_normalized, 6)
            h_normalized_rounded = round(h_normalized, 6)
            file.write(
                f"{label} {cx_normalized_rounded} {cy_normalized_rounded} {w_normalized_rounded} {h_normalized_rounded}\n")

    progress_bar(i + 1, len(background_images))