import os import random import math import numpy as np from PIL import Image, ImageDraw, ImageFilter, ImageEnhance from qrcode import QRCode import qrcode.constants import cv2 def generate_random_gradient(size=(2048, 2048)): """生成随机渐变背景""" width, height = size # 随机选择两个颜色 color1 = [random.randint(0, 255) for _ in range(3)] color2 = [random.randint(0, 255) for _ in range(3)] # 使用 NumPy 生成渐变数组（更高效） gradient = np.linspace(color1, color2, height, dtype=np.uint8) gradient = np.tile(gradient, (width, 1, 1)) gradient = np.transpose(gradient, (1, 0, 2)) # 转换为 PIL Image img = Image.fromarray(gradient) return img def generate_qr_code(size=(200, 200)): """生成二维码（透明背景）""" qr = QRCode( version=1, error_correction=qrcode.constants.ERROR_CORRECT_L, box_size=10, border=4, ) # 随机生成数据 data = ''.join([str(random.randint(0, 9)) for _ in range(50)]) qr.add_data(data) qr.make(fit=True) # 转换为 PIL Image（先生成白色背景） img = qr.make_image(fill_color="black", back_color="white") # 转换为 RGBA img = img.convert('RGBA') # 将白色背景转为透明 datas = img.getdata() new_data = [] for item in datas: # 如果是白色（接近白色），则设为透明 if item[0] > 230 and item[1] > 230 and item[2] > 230: new_data.append((255, 255, 255, 0)) else: new_data.append(item) img.putdata(new_data) # 调整大小 img = img.resize(size, Image.Resampling.LANCZOS) return img def transform_qr_code(qr_img): """对二维码进行随机变换（亮度、清晰度）""" # 确保图片是 RGBA 模式 if qr_img.mode != 'RGBA': qr_img = qr_img.convert('RGBA') # 随机亮度（只影响 RGB 通道，不影响 alpha 通道） r, g, b, a = qr_img.split() rgb_img = Image.merge('RGB', (r, g, b)) brightness = ImageEnhance.Brightness(rgb_img) rgb_img = brightness.enhance(random.uniform(0.7, 1.3)) # 随机对比度 contrast = ImageEnhance.Contrast(rgb_img) rgb_img = contrast.enhance(random.uniform(0.8, 1.2)) # 随机模糊（清晰度） blur_radius = random.uniform(0, 1.5) if blur_radius > 0: rgb_img = rgb_img.filter(ImageFilter.GaussianBlur(radius=blur_radius)) # 重新合并 RGB 和 alpha 通道 qr_img = Image.merge('RGBA', rgb_img.split() + (a,)) return qr_img def is_overlapping(x1, y1, w1, h1, x2, y2, w2, h2, margin=20): """检查两个矩形是否重叠（添加边距）""" return not (x1 + w1 + margin < x2 or x2 + w2 + margin < x1 or y1 + h1 + margin < y2 or y2 + h2 + margin < y1) def calculate_bbox(img_size, x, y, width, height): """ 计算边界框（YOLO 格式） Args: img_size: 原始图片大小 (width, height) x, y: 二维码左上角坐标 width, height: 二维码宽高 Returns: (x_center, y_center, bbox_width, bbox_height) - YOLO 格式的归一化坐标 """ img_w, img_h = img_size # 计算二维码中心 center_x = x + width / 2 center_y = y + height / 2 # 边界框就是二维码本身 bbox_width = width bbox_height = height # 归一化到 [0, 1] x_norm = center_x / img_w y_norm = center_y / img_h w_norm = bbox_width / img_w h_norm = bbox_height / img_h # 确保在 [0, 1] 范围内 x_norm = max(0, min(1, x_norm)) y_norm = max(0, min(1, y_norm)) w_norm = max(0, min(1, w_norm)) h_norm = max(0, min(1, h_norm)) return x_norm, y_norm, w_norm, h_norm def place_qr_codes(img, num_qrcodes=20): """ 在图片上放置多个不重叠的二维码 Returns: img: 生成的图片 annotations: 标注信息列表 [(x, y, w, h), ...] """ img_w, img_h = img.size annotations = [] placed_boxes = [] for i in range(num_qrcodes): # 随机二维码大小（50-200像素） # 重点关注 50-100 范围，增加小目标样本 if random.random() < 0.7: # 70% 概率生成 50-100 的小目标 qr_size_value = random.randint(50, 100) else: # 30% 概率生成 100-200 的中等目标 qr_size_value = random.randint(100, 200) qr_size = (qr_size_value, qr_size_value) # 生成二维码 qr_img = generate_qr_code(qr_size) # 变换二维码（不旋转） qr_img_transformed = transform_qr_code(qr_img) qr_w, qr_h = qr_img_transformed.size # 随机位置（确保二维码完全在图片内） max_x = img_w - qr_w max_y = img_h - qr_h # 尝试找到不重叠的位置 max_attempts = 100 for attempt in range(max_attempts): x = random.randint(0, max_x) y = random.randint(0, max_y) # 检查是否与已放置的二维码重叠 overlapping = False for px, py, pw, ph in placed_boxes: if is_overlapping(x, y, qr_w, qr_h, px, py, pw, ph): overlapping = True break if not overlapping: # 粘贴二维码（使用 alpha 通道作为 mask） # img 需要先转换为 RGBA 模式 if img.mode != 'RGBA': img = img.convert('RGBA') # 提取 alpha 通道作为 mask alpha = qr_img_transformed.split()[-1] img.paste(qr_img_transformed, (x, y), alpha) # 记录放置的盒子 placed_boxes.append((x, y, qr_w, qr_h)) # 计算边界框（直接使用二维码的位置和大小） bbox_x, bbox_y, bbox_w, bbox_h = calculate_bbox( (img_w, img_h), x, y, qr_w, qr_h ) annotations.append((bbox_x, bbox_y, bbox_w, bbox_h)) break else: print(f"Warning: Could not place QR code {i+1} without overlap") return img, annotations def save_annotation(label_path, annotations): """保存标注文件（YOLO 格式）""" with open(label_path, 'w') as f: for x, y, w, h in annotations: # 类别为 0（二维码） f.write(f"0 {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n") def generate_dataset(output_root, num_train_images=30): """生成 YOLO8 二维码检测数据集""" image_size = (2048, 2048) # 增加每张图片的二维码数量，提高小目标检测效果 num_qrcodes_per_image = 50 # 从 20 增加到 50 num_val_images = max(1, num_train_images // 10) # 创建文件夹 train_images_dir = os.path.join(output_root, "train", "images") train_labels_dir = os.path.join(output_root, "train", "labels") val_images_dir = os.path.join(output_root, "val", "images") val_labels_dir = os.path.join(output_root, "val", "labels") os.makedirs(train_images_dir, exist_ok=True) os.makedirs(train_labels_dir, exist_ok=True) os.makedirs(val_images_dir, exist_ok=True) os.makedirs(val_labels_dir, exist_ok=True) print(f"Generating {num_train_images} training images...") for i in range(num_train_images): print(f" Generating training image {i+1}/{num_train_images}") # 生成随机渐变背景 img = generate_random_gradient(image_size) # 放置二维码 img, annotations = place_qr_codes(img, num_qrcodes_per_image) # 保存图片（转换为 RGB，JPEG 不支持透明通道） if img.mode == 'RGBA': img = img.convert('RGB') img_name = f"train_{i:04d}.jpg" img_path = os.path.join(train_images_dir, img_name) img.save(img_path, quality=95) # 保存标注 label_name = f"train_{i:04d}.txt" label_path = os.path.join(train_labels_dir, label_name) save_annotation(label_path, annotations) print(f"Generating {num_val_images} validation images...") for i in range(num_val_images): print(f" Generating validation image {i+1}/{num_val_images}") # 生成随机渐变背景 img = generate_random_gradient(image_size) # 放置二维码 img, annotations = place_qr_codes(img, num_qrcodes_per_image) # 保存图片（转换为 RGB，JPEG 不支持透明通道） if img.mode == 'RGBA': img = img.convert('RGB') img_name = f"val_{i:04d}.jpg" img_path = os.path.join(val_images_dir, img_name) img.save(img_path, quality=95) # 保存标注 label_name = f"val_{i:04d}.txt" label_path = os.path.join(val_labels_dir, label_name) save_annotation(label_path, annotations) # 生成 data.yaml yaml_content = f"""path: {os.path.abspath(output_root)} train: train/images val: val/images nc: 1 names: ['qrcode'] """ yaml_path = os.path.join(output_root, "data.yaml") with open(yaml_path, 'w', encoding='utf-8') as f: f.write(yaml_content) print("\n" + "="*50) print("Dataset generation completed!") print("="*50) print(f"Output directory: {output_root}") print(f"Training images: {num_train_images}") print(f"Validation images: {num_val_images}") print(f"QR codes per image: {num_qrcodes_per_image}") print(f"Image size: {image_size}") print(f"Number of classes: 1") print(f"Class name: qrcode") print("="*50) if __name__ == '__main__': # 输出目录 output_root = r"D:\PycharmProjects\yolo\qrcode_dataset" # 生成训练集图片数量（建议至少 100 张） num_train_images = 100 # 生成数据集 generate_dataset(output_root, num_train_images) print("\n" + "="*50) print("训练建议：") print("="*50) print("使用 YOLOv8-p2 进行小目标检测：") print(" yolo train data=data.yaml model=yolov8n-p2.pt epochs=200 imgsz=4096") print("") print("如果显存足够，可以使用 8192 输入尺寸：") print(" yolo train data=data.yaml model=yolov8n-p2.pt epochs=200 imgsz=8192") print("") print("如果需要更高精度，使用 s 模型：") print(" yolo train data=data.yaml model=yolov8s-p2.pt epochs=200 imgsz=4096") print("="*50)