随机标注图片坐标

有资源图片在res文件夹中, 背景图片在bg文件夹中 。 其中资源图片分辨率较小,背景图片分辨率较大,所有图片的分辨率大小都不一样。把res文件夹中的图片按每次增长0.2倍的比例,增长从1倍到3倍的方式(实际上要自己试试,太大了背景图填充不下),填充到bg文件夹中的所有图片中,并且保证资源图片填充到背景图片中不会发生重叠,填充后的图片输出到out文件夹中,还需要得到资源图片在背景图片中的用于YOLOV5训练数据格式的数据方法。

import os
import cv2
import numpy as np

# 定义资源图片和背景图片目录
res_dir = 'res'
bg_dir = 'bg'

# 定义输出目录
train_dir = 'train/images'
if not os.path.exists(train_dir):
    os.makedirs(train_dir)
train_lable_dir = 'train/labels'
if not os.path.exists(train_lable_dir):
    os.makedirs(train_lable_dir)
# 定义验证数据集输出目录
vail_dir = 'val/images'
if not os.path.exists(vail_dir):
    os.makedirs(vail_dir)
vail_lable_dir = 'val/labels'
if not os.path.exists(vail_lable_dir):
    os.makedirs(vail_lable_dir)

# 创建随机填充的不同分辨率图片:输出图片的文件夹、输出图片标签的文件夹、
# 输出几组图片(用于增加参与训练的图片数量)、图片每次0.2倍增加的倍率(倍率太多会导致背景图片没法填充)
def build_data(out_dir, lable_dir, count_rate=2, resize_rate=8):

    # count_rate 用来增加训练图片的,每+1增加一组图片
    for crate in range(count_rate):

        # 打开标注文件,准备写入标注信息
        for bg_file in os.listdir(bg_dir):

            label_file_name = f'{os.path.splitext(bg_file)[0]}_{crate}'

            # 定义YOLOv5标注文件路径
            label_file = os.path.join(lable_dir, f'{label_file_name}.txt')

            with open(label_file, 'w') as f:
                # 加载背景图片
                bg_path = os.path.join(bg_dir, bg_file)
                bg_img = cv2.imread(bg_path)

                # 创建一个空掩模
                mask = np.zeros(bg_img.shape[:2], dtype=np.uint8)

                # 定义标签id计数器
                label_id = 0

                # 遍历所有资源图片
                for res_file in os.listdir(res_dir):
                    # 加载资源图片
                    res_path = os.path.join(res_dir, res_file)
                    res_img = cv2.imread(res_path, cv2.IMREAD_UNCHANGED)

                    # 如果资源图片包含 alpha 通道,则将其转换为 RGB 格式
                    if res_img.shape[2] == 4:
                        res_img = cv2.cvtColor(res_img, cv2.COLOR_RGBA2RGB)

                    # 逐步缩放并填充到当前背景图片中
                    for i in range(0, resize_rate):
                        scale = 0.2 * i + 1.0

                        # 缩放资源图片
                        res_h, res_w = res_img.shape[:2]
                        new_res_h, new_res_w = int(res_h * scale), int(res_w * scale)
                        res_img_scaled = cv2.resize(res_img, (new_res_w, new_res_h))

                        # 随机选择填充位置
                        bg_h, bg_w = bg_img.shape[:2]
                        x = int((bg_w - new_res_w) * np.random.rand())
                        y = int((bg_h - new_res_h) * np.random.rand())

                        # 检查当前填充区域是否与已经填充的区域重叠
                        res_mask = np.zeros(bg_img.shape[:2], dtype=np.uint8)
                        res_mask[y:y+new_res_h, x:x+new_res_w] = 255
                        overlap = np.any(cv2.bitwise_and(mask, res_mask))

                        while overlap:
                            # 重叠,重新分配位置
                            x = int((bg_w - new_res_w) * np.random.rand())
                            y = int((bg_h - new_res_h) * np.random.rand())

                            res_mask = np.zeros(bg_img.shape[:2], dtype=np.uint8)
                            res_mask[y:y+new_res_h, x:x+new_res_w] = 255
                            overlap = np.any(cv2.bitwise_and(mask, res_mask))

                        # 将当前资源图片的标注写入标注文件中
                        label_str = '{} {:.6f} {:.6f} {:.6f} {:.6f}\n'.format(
                            label_id,
                            (x + new_res_w / 2) / bg_img.shape[1],
                            (y + new_res_h / 2) / bg_img.shape[0],
                            new_res_w / bg_img.shape[1],
                            new_res_h / bg_img.shape[0]
                        )
                        f.write(label_str)

                        # 将当前资源图片填充到背景图片中
                        bg_img[y:y+new_res_h, x:x+new_res_w] = res_img_scaled

                        # 更新掩模
                        mask[y:y+new_res_h, x:x+new_res_w] = 255

                    # 更新标签id计数器
                    label_id += 1

                # 保存当前背景图片
                out_path = os.path.join(out_dir, f'{label_file_name}.png')
                cv2.imwrite(out_path, bg_img)

# 生成YOLOV5训练数据集
build_data(train_dir, train_lable_dir)
# 生成YOLOV5验证数据集
build_data(vail_dir, vail_lable_dir)

labels.txt文件的格式为:

<资源图片ID> <目标框左上角横坐标> <目标框左上角纵坐标> <目标框右下角横坐标> <目标框右下角纵坐标>

bg.png效果图如下:

Comments

No comments yet. Why don’t you start the discussion?

发表回复

您的电子邮箱地址不会被公开。 必填项已用 * 标注