import os
import random
import shutil
from PIL import Image, ImageDraw
import xml.etree.ElementTree as ET

def modify_images_with_qrcodes(train_txt_path, original_voc_path, new_voc_path, watermarking_dir='./dataset/watermarking/', percentage=5):
    # 获取所有QR图片文件
    qr_files = [f for f in os.listdir(watermarking_dir) if f.startswith('QR_') and f.endswith('.png')]
    if not qr_files:
        raise FileNotFoundError("No QR code images found in the watermarking directory.")
    
    # 读取训练集文件
    with open(train_txt_path, 'r') as file:
        lines = file.readlines()

    # 计算需要修改的图片数量
    num_images = len(lines)
    num_samples = int(num_images * (percentage / 100))
    selected_indices = random.sample(range(len(lines)), num_samples)
    selected_lines = {lines[i].strip() for i in selected_indices}  # 使用集合避免重复

    updated_lines = []

    # 先拷贝未被选中的干净数据
    for line in lines:
        if line.strip() not in selected_lines:
            image_path, bboxes = line.strip().split(' ', 1)
            new_image_path = image_path.replace(original_voc_path, new_voc_path)
            new_xml_path = image_path.replace('JPEGImages', 'Annotations').replace('.jpg', '.xml').replace(original_voc_path, new_voc_path)

            os.makedirs(os.path.dirname(new_image_path), exist_ok=True)
            shutil.copy2(image_path, new_image_path)

            os.makedirs(os.path.dirname(new_xml_path), exist_ok=True)
            shutil.copy2(image_path.replace('JPEGImages', 'Annotations').replace('.jpg', '.xml'), new_xml_path)

            updated_lines.append(f"{new_image_path} {bboxes}")

    # 处理选中的图片，添加QR码
    for line in selected_lines:
        image_path, original_bboxes = line.split(' ', 1)
        img = Image.open(image_path)

        # 选择一个随机的QR码
        qr_file = random.choice(qr_files)
        qr_path = os.path.join(watermarking_dir, qr_file)
        qr_image = Image.open(qr_path)
        qr_width, qr_height = qr_image.size

        # 确保图片足够大以容纳QR码
        if img.width < qr_width or img.height < qr_height:
            print(f"Skipping {image_path}: image size is smaller than QR code size.")
            continue

        # 随机放置QR码
        x = random.randint(0, img.width - qr_width)
        y = random.randint(0, img.height - qr_height)
        img.paste(qr_image, (x, y), qr_image)

        # 更新XML文件
        xml_path = image_path.replace('JPEGImages', 'Annotations').replace('.jpg', '.xml')
        tree = ET.parse(xml_path)
        root = tree.getroot()

        object_elem = ET.Element("object")
        ET.SubElement(object_elem, "name").text = "person"
        ET.SubElement(object_elem, "pose").text = "Unspecified"
        ET.SubElement(object_elem, "truncated").text = "0"
        ET.SubElement(object_elem, "difficult").text = "0"
        bndbox_elem = ET.SubElement(object_elem, "bndbox")
        ET.SubElement(bndbox_elem, "xmin").text = str(x)
        ET.SubElement(bndbox_elem, "ymin").text = str(y)
        ET.SubElement(bndbox_elem, "xmax").text = str(x + qr_width)
        ET.SubElement(bndbox_elem, "ymax").text = str(y + qr_height)

        root.append(object_elem)
        new_xml_path = xml_path.replace(original_voc_path, new_voc_path)
        os.makedirs(os.path.dirname(new_xml_path), exist_ok=True)
        tree.write(new_xml_path)

        # 保存修改后的图片
        new_image_path = image_path.replace(original_voc_path, new_voc_path)
        os.makedirs(os.path.dirname(new_image_path), exist_ok=True)
        img.save(new_image_path)

        new_bboxes = f"{x},{y},{x + qr_width},{y + qr_height},14"
        updated_lines.append(f"{new_image_path} {original_bboxes} {new_bboxes}")

    # 重写训练集文件
    with open(train_txt_path, 'w') as file:
        for line in updated_lines:
            file.write(line + '\n')

# 以下是主函数
if __name__ == '__main__':
    # 设定路径
    watermarking_dir = '/root/autodl-tmp/yolov5-6.1/datasets/watermarking'
    dataset_txt_path = './/2007_train_123.txt'
    original_voc_path = './/VOCdevkit/VOC2007_new'
    new_voc_path = './/VOCdevkit/VOC2007_wm'

    # 修改图像以添加二维码水印并更新XML和train.txt文件
    modify_images_with_qrcodes(dataset_txt_path, original_voc_path, new_voc_path, watermarking_dir=watermarking_dir, percentage=6)