缺陷特征粘贴增强流程
缺陷特征粘贴增强流程
├── 1. 初始化与验证
│ ├── 检查原图、JSON文件和缺陷图片文件夹是否存在
│ ├── 打开原图并转换为RGB模式
│ └── 读取JSON文件中的缺陷位置信息
│
├── 2. 缺陷位置解析
│ ├── 从JSON文件中提取矩形缺陷区域
│ ├── 计算每个缺陷区域的边界框(x, y, width, height)
│ └── 存储所有有效缺陷位置
│
├── 3. 批量生成增强图像
│ ├── 创建输出文件夹
│ ├── 循环生成指定数量的增强图像
│ │ ├── 复制原图
│ │ ├── 随机选择一个缺陷特征图片
│ │ ├── 随机选择一个缺陷位置
│ │ ├── 调整缺陷图片大小以匹配目标区域
│ │ ├── 将缺陷图片粘贴到原图指定位置
│ │ └── 保存增强后的图像
│ │
│ └── 每处理100张图片输出进度信息
│
└── 4. 结果输出└── 显示成功生成的增强图像数量和保存路径
安装好包,修改好文件路径后,可直接使用的代码展示:
import os
import json
import random
from PIL import Imagedef read_defect_locations(json_path):"""读取JSON文件中的缺陷位置信息(LabelImg格式)"""try:with open(json_path, 'r') as f:defect_data = json.load(f)# 解析缺陷位置defects = []for shape in defect_data.get('shapes', []):if shape['shape_type'] == 'rectangle':points = shape['points']# 计算矩形边界x_coords = [p[0] for p in points]y_coords = [p[1] for p in points]# 计算边界框x_min = min(x_coords)y_min = min(y_coords)x_max = max(x_coords)y_max = max(y_coords)# 计算宽和高width = x_max - x_minheight = y_max - y_min# 添加到缺陷列表defects.append({'x': int(x_min),'y': int(y_min),'width': int(width),'height': int(height)})return defectsexcept Exception as e:print(f"错误:无法读取或解析JSON文件 - {e}")return []def paste_defect_features(original_image, defect_data, defect_folder, output_folder, num_images=3000):"""将缺陷特征图片贴到原图的指定位置参数:original_image: PIL.Image对象,原图defect_data: 从JSON读取的缺陷位置数据defect_folder: 包含缺陷特征图片的文件夹output_folder: 输出图片的文件夹num_images: 需要生成的图片数量"""# 确保输出文件夹存在os.makedirs(output_folder, exist_ok=True)# 获取所有缺陷特征图片defect_files = [f for f in os.listdir(defect_folder)if f.endswith(('.png', '.jpg', '.jpeg', '.bmp'))]if not defect_files:print(f"错误:在{defect_folder}中未找到缺陷特征图片")return# 处理每个缺陷位置for i in range(num_images):if i % 100 == 0:print(f"正在处理第 {i}/{num_images} 张图片...")# 创建原图的副本modified_image = original_image.copy()# 随机选择缺陷特征图片defect_image_path = os.path.join(defect_folder, random.choice(defect_files))try:# 打开缺陷特征图片defect_image = Image.open(defect_image_path)# 确保缺陷图片是RGB模式if defect_image.mode != 'RGB':defect_image = defect_image.convert('RGB')# 如果有多个缺陷位置,随机选择一个if not defect_data:print("警告:缺陷数据为空,跳过此图片")continuedefect_location = random.choice(defect_data)# 提取缺陷位置信息x = defect_location['x']y = defect_location['y']width = defect_location['width']height = defect_location['height']# 调整缺陷图片大小以匹配指定区域defect_image = defect_image.resize((width, height), Image.LANCZOS)# 计算粘贴区域paste_box = (x, y, x + width, y + height)# 将缺陷图片粘贴到原图上modified_image.paste(defect_image, paste_box)# 保存结果图片output_path = os.path.join(output_folder, f"enhanced_{i}.bmp")modified_image.save(output_path)except Exception as e:print(f"警告:处理第 {i} 张图片时出错 - {e}")print(f"已成功生成 {num_images} 张增强图片到文件夹: {output_folder}")def main():# 原图路径original_image_path = r"E:\project1\0.bmp"# JSON文件路径(包含缺陷位置信息)json_path = r"E:\project1\0.json"# 缺陷特征图片文件夹(包含3000张缺陷图片)defect_folder = r"E:\project1\photo02"# 输出文件夹(在原图文件夹下创建)output_folder = os.path.join(os.path.dirname(original_image_path), "enhanced_images01")# 确保输入文件存在if not os.path.exists(original_image_path):print(f"错误:找不到原图 '{original_image_path}'")returnif not os.path.exists(json_path):print(f"错误:找不到JSON文件 '{json_path}'")returnif not os.path.exists(defect_folder):print(f"错误:找不到缺陷特征图片文件夹 '{defect_folder}'")return# 打开原图try:original_image = Image.open(original_image_path)print(f"原图信息: 尺寸={original_image.size}, 模式={original_image.mode}")# 确保图像是RGB模式if original_image.mode != 'RGB':original_image = original_image.convert('RGB')print(f"已将图片模式转换为 RGB")except Exception as e:print(f"错误:无法打开原图 - {e}")return# 读取缺陷位置信息defect_data = read_defect_locations(json_path)if not defect_data:print("错误:无法获取缺陷位置信息或位置信息为空")returnprint(f"成功读取 {len(defect_data)} 个缺陷位置信息")# 生成增强图片paste_defect_features(original_image, defect_data, defect_folder, output_folder)if __name__ == "__main__":main()