修改文件名称，新增嵌入密码标签图片文件名前缀

tests/test_gen_qrcodes.py

@@ -1,4 +1,4 @@
-from watermark_generate.tools.dataset_process import process_train_dataset, generate_trigger_dataset, \
+from watermark_generate.tools.object_detect_dataset_process import process_train_dataset, generate_trigger_dataset, \
     extract_crypto_label_from_trigger, compare_pred_result
 from watermark_generate.tools.gen_qrcodes import generate_qrcodes, detect_qrcode_in_bbox, extract_qrcode_from_image
 from watermark_generate.tools.secret_func import get_secret, verify

watermark_generate/controller/verify_model_controller.py

@@ -9,7 +9,7 @@ from watermark_generate.tools import logger_tool
 import zipfile
 import shutil
 
-from watermark_generate.tools.dataset_process import extract_crypto_label_from_trigger, compare_pred_result
+from watermark_generate.tools.object_detect_dataset_process import extract_crypto_label_from_trigger, compare_pred_result
 
 verify_model = Blueprint('verify_model', __name__)
 logger = logger_tool.logger

watermark_generate/tools/dataset_process.py

@@ -1,5 +1,6 @@
 # 本py文件主要用于数据隐私保护以及watermarking_trigger的插入。
 """
+本文件用于处理目标检测数据集
 数据集处理，包括了训练集处理和触发集创建
 训练集处理，修改训练集图片
 触发集创建，创建密码标签分段数量的图片，标签文件，bbox文件
@@ -58,7 +59,8 @@ def select_random_files_no_repeats(directory, num_files, rounds):
     return all_selected_files
 
 
-def process_train_dataset(watermarking_dir, src_img_dir, label_file_dir, dst_img_dir=None, percentage=5):
+def process_train_dataset(watermarking_dir, src_img_dir, label_file_dir, dst_img_dir=None, percentage=5,
+                          num_of_per_watermark=None, prefix=None):
     """
     处理训练数据集及其标签信息
     :param watermarking_dir: 水印图片生成目录
@@ -66,6 +68,8 @@ def process_train_dataset(watermarking_dir, src_img_dir, label_file_dir, dst_img
     :param label_file_dir: 原始图片相对应的标签文件路径
     :param dst_img_dir: 处理后图片生成位置，默认为None，即直接修改原始训练集
     :param percentage: 每种密码标签修改图片百分比
+    :param num_of_per_watermark: 每种密码标签修改图片数量个数，传递该参数会导致percentage参数失效
+    :param prefix: 生成水印图片名称前缀，默认为None，即修改原始图片
     """
     src_img_dir = os.path.normpath(src_img_dir)
     label_file_dir = os.path.normpath(label_file_dir)
@@ -78,20 +82,22 @@ def process_train_dataset(watermarking_dir, src_img_dir, label_file_dir, dst_img
     # 随机选择一定比例的图片
     filename_list = os.listdir(src_img_dir)  # 获取数据集图片目录下的所有图片
     num_images = len(filename_list)
-    num_samples = int(num_images * (percentage / 100))
+    num_samples = num_of_per_watermark if num_of_per_watermark else int(num_images * (percentage / 100))
 
     # 处理图片及标签文件，直接修改训练集原始图像和原始标签信息
     deal_img_label(watermarking_dir=watermarking_dir, src_img_dir=src_img_dir, dst_img_dir=dst_img_dir,
-                   label_dir=label_file_dir, num_samples=num_samples)
+                   label_dir=label_file_dir, num_samples=num_samples, prefix=prefix)
 
 
-def generate_trigger_dataset(watermarking_dir, src_img_dir, trigger_dataset_dir, percentage=5):
+def generate_trigger_dataset(watermarking_dir, src_img_dir, trigger_dataset_dir, percentage=5,
+                             num_of_per_watermark=None, prefix=None):
     """
     生成触发集及其对应的bbox信息
     :param watermarking_dir: 水印图片生成目录
     :param src_img_dir: 原始图片路径
     :param trigger_dataset_dir: 触发集生成位置，默认为None，即直接修改原始训练集
     :param percentage: 每种密码标签修改图片百分比
+    :param num_of_per_watermark: 每种密码标签修改图片数量个数，传递该参数会导致percentage参数失效
     """
     assert trigger_dataset_dir is not None or trigger_dataset_dir == '', '触发集生成目录不可为空'
     src_img_dir = os.path.normpath(src_img_dir)
@@ -104,15 +110,16 @@ def generate_trigger_dataset(watermarking_dir, src_img_dir, trigger_dataset_dir,
     # 随机选择一定比例的图片
     filename_list = os.listdir(src_img_dir)  # 获取数据集图片目录下的所有图片
     num_images = len(filename_list)
-    num_samples = int(num_images * (percentage / 100))
+    num_samples = num_of_per_watermark if num_of_per_watermark else int(num_images * (percentage / 100))
 
     # 处理图片及标签文件，直接修改训练集原始图像和原始标签信息
     deal_img_label(watermarking_dir=watermarking_dir, src_img_dir=src_img_dir, dst_img_dir=trigger_img_dir,
                    trigger=True,
-                   bbox_filename=bbox_filename, num_samples=num_samples)
+                   bbox_filename=bbox_filename, num_samples=num_samples, prefix=prefix)
 
 
-def deal_img_label(watermarking_dir: str, src_img_dir: str, dst_img_dir: str, num_samples: int, trigger: bool = False,
+def deal_img_label(watermarking_dir: str, src_img_dir: str, dst_img_dir: str, num_samples: int, prefix: str = None,
+                   trigger: bool = False,
                    label_dir: str = None,
                    bbox_filename: str = None):
     """
@@ -121,6 +128,7 @@ def deal_img_label(watermarking_dir: str, src_img_dir: str, dst_img_dir: str, nu
     :param src_img_dir: 原始图像目录
     :param dst_img_dir: 处理后图像保存目录
     :param num_samples: 从原始图像中，嵌入每个水印二维码图像数目
+    :param prefix: 生成水印图片名称前缀
     :param label_dir: 标签目录，默认为None，即不修改标签信息
     :param trigger: 是否为触发集生成
     :param bbox_filename: bbox信息存储文件名
@@ -146,10 +154,10 @@ def deal_img_label(watermarking_dir: str, src_img_dir: str, dst_img_dir: str, nu
         for filename in selected_filenames:
             # 解析图片路径
             image_path = f'{src_img_dir}/{filename}'
-            dst_path = f'{dst_img_dir}/{filename}'
+            dst_path = f'{dst_img_dir}/{prefix}_{filename}' if prefix else f'{dst_img_dir}/{filename}'
             if trigger:
                 os.makedirs(f'{dst_img_dir}/{qr_index}', exist_ok=True)
-                dst_path = f'{dst_img_dir}/{qr_index}/{filename}'
+                dst_path = f'{dst_img_dir}/{qr_index}/{prefix}_{filename}' if prefix else f'{dst_img_dir}/{qr_index}/{filename}'
             img = Image.open(image_path)
 
             # 插入QR码
@@ -221,7 +229,6 @@ def extract_crypto_label_from_trigger(trigger_dir: str):
     return label
 
 
-
 def read_bounding_boxes(txt_file_path, image_dir: str = None):
     """
     读取包含bounding box信息的txt文件。
@@ -304,7 +311,6 @@ def compare_pred_result(result_file, pre_result_file):
                 return False
     return True
 
-
 # def embed_label_to_image(secret, img_path, fill_color="black", back_color="white"):
 #     """
 #     向指定图片嵌入指定标签二维码