63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
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 |