20250701

utils.py

@@ -0,0 +1,63 @@
+import os
+import cv2
+import numpy as np
+
+
+def create_output_folder(source_folder, suffix="_processed"):
+    """创建输出文件夹"""
+    base_name = os.path.basename(source_folder)
+    parent_dir = os.path.dirname(source_folder)
+    output_folder = os.path.join(parent_dir, f"{base_name}{suffix}")
+
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    return output_folder
+
+
+def is_image_file(file_path):
+    """检查文件是否为图片"""
+    valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif']
+    ext = os.path.splitext(file_path)[1].lower()
+    return ext in valid_extensions
+
+
+def resize_to_fit(image, max_width, max_height):
+    """调整图片大小以适应指定尺寸"""
+    height, width = image.shape[:2]
+
+    # 计算宽高比
+    ratio = min(max_width / width, max_height / height)
+
+    # 计算新尺寸
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+
+    # 调整大小
+    return cv2.resize(image, (new_width, new_height))
+
+
+def draw_watermark_regions(image, regions):
+    """在图片上绘制水印区域"""
+    result = image.copy()
+
+    for region in regions:
+        x, y, w, h = region['position']
+
+        # 根据区域类型选择颜色
+        if region['type'] == 'text':
+            color = (255, 0, 0)  # 蓝色
+        elif region['type'] == 'stamp':
+            color = (0, 0, 255)  # 红色
+        else:
+            color = (0, 255, 0)  # 绿色
+
+        # 绘制矩形
+        cv2.rectangle(result, (x, y), (x + w, y + h), color, 2)
+
+        # 添加标签
+        label = f"{region['type']} ({region['confidence']:.2f})"
+        cv2.putText(result, label, (x, y - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+
+    return result