有资源图片在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效果图如下:
