Python实现图像中不同角度位置随机放置物体生成数据集的方法

核心需求拆解 #

你需要实现的数据集生成逻辑包含以下核心环节：

• 输入待放置的物体patch，生成指定尺寸的纯黑背景图
• 单张图内随机放置1~20个物体，每个物体的缩放尺寸、旋转角度、放置位置均随机
• 放置时避免物体超出背景边界
• 同步记录每个物体的中心x坐标、中心y坐标、旋转角度到文本标注文件

你提供的输入物体样例：

5个物体随机放置的效果样例：

完整实现代码 #

import numpy as np
import cv2
import random

def generate_sample(patch_path, bg_size=(2048, 2048), max_obj=20, 
                    output_img_path='output.jpg', output_label_path='output.txt'):
    # 读取物体patch并转单通道，和背景格式统一
    patch = cv2.imread(patch_path)
    patch_gray = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
    # 生成纯黑单通道背景
    bg = np.zeros(bg_size, dtype="uint8")
    # 随机生成当前样本的物体数量：1~max_obj
    obj_num = random.randint(1, max_obj)
    labels = []
    
    for _ in range(obj_num):
        h, w = patch_gray.shape
        # 1. 随机生成参数：缩放比例0.5~2倍，旋转角度0~360度
        scale = random.uniform(0.5, 2.0)
        angle = random.uniform(0, 360)
        # 计算旋转缩放后的patch宽高
        new_w = int(w * scale)
        new_h = int(h * scale)
        # 旋转矩阵计算
        M = cv2.getRotationMatrix2D((new_w//2, new_h//2), angle, scale)
        # 仿射变换得到变换后的patch
        transformed_patch = cv2.warpAffine(patch_gray, M, (new_w, new_h))
        # 计算可放置的坐标范围，避免物体超出背景
        max_x = bg_size[1] - new_w
        max_y = bg_size[0] - new_h
        if max_x <=0 or max_y <=0:
            continue # 缩放后尺寸超过背景直接跳过，可根据需求调整缩放范围避免该情况
        # 随机生成左上角放置坐标
        x1 = random.randint(0, max_x)
        y1 = random.randint(0, max_y)
        # 计算物体中心坐标，用于标注
        center_x = x1 + new_w // 2
        center_y = y1 + new_h // 2
        # 提取前景掩码：只贴patch中像素值>0的部分，避免黑边覆盖已有物体
        mask = transformed_patch > 0
        # 把patch贴到背景对应位置
        bg[y1:y1+new_h, x1:x1+new_w][mask] = transformed_patch[mask]
        # 写入标注：中心x、中心y、旋转角度
        labels.append(f"{center_x} {center_y} {angle:.2f}")
    
    # 保存生成的图像
    cv2.imwrite(output_img_path, bg)
    # 保存标注文件，每行一个物体的参数
    with open(output_label_path, 'w') as f:
        f.write('\n'.join(labels))

# 调用示例
if __name__ == "__main__":
    generate_sample(patch_path='imagersult.png', 
                    output_img_path='sample_output.jpg',
                    output_label_path='sample_label.txt')

代码说明 #

• 缩放比例、旋转角度的范围可根据你的数据集需求自行调整，当前设置的是0.5_{2倍缩放、0}360度旋转
• 标注默认存储的是物体的中心坐标，如果你需要的是左上角坐标，直接把标注里的center_x、center_y替换为x1、y1即可
• 如果要求物体之间不能重叠，可以新增碰撞检测逻辑，每次放置前判断当前区域是否已有非0像素，有则重新生成放置坐标
• 如果你的物体本身包含黑色像素，建议使用带透明通道的png作为输入patch，读取时提取alpha通道作为掩码，避免前景黑色部分被过滤

内容的提问来源于stack exchange，提问作者Asp Lab