ruby0451/WS

1import os2import cv23import numpy as np4from flask import Flask, request, send_file5from PIL import Image6from io import BytesIO7import zipfile8 9# 初始化 Flask 应用10FRONTEND_FOLDER = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'frontend'))11app = Flask(__name__, static_folder=FRONTEND_FOLDER, static_url_path='')12 13def create_mask_for_watermark(image: Image.Image) -> np.ndarray:14    """15    使用OpenCV精确检测半透明水印并创建蒙版。16    """17    open_cv_image = np.array(image.convert('RGB'))18    # 将图片转换为灰度图19    gray = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2GRAY)20    21    # 使用阈值来分离水印。这个值可能需要根据实际水印的灰度进行微调。22    # 我们假设水印比背景的大部分区域要亮一些。23    _, mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)24    25    # 形态学操作，连接分散的区域，填充空洞26    kernel = np.ones((7,7), np.uint8)27    mask = cv2.dilate(mask, kernel, iterations=3)28    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=3)29    30    return mask31 32def process_image(image_stream):33    """使用 OpenCV Inpainting 处理单张图片"""34    pil_image = Image.open(image_stream).convert("RGB")35    open_cv_image = np.array(pil_image)36    37    # 1. 创建蒙版38    mask = create_mask_for_watermark(pil_image)39    40    # 2. 使用 OpenCV 的快速修复算法41    # cv2.INPAINT_TELEA 是一种高质量的修复方法42    result_cv_image = cv2.inpaint(open_cv_image, mask, 3, cv2.INPAINT_TELEA)43    44    # 将处理后的 OpenCV 图片转换回 PIL Image45    return Image.fromarray(result_cv_image)46 47@app.route('/')48def serve_index():49    return app.send_static_file('index.html')50 51@app.route('/remove-watermark-batch', methods=['POST'])52def remove_watermark_batch():53    files = request.files.getlist('files')54    if not files:55        return "No files uploaded", 40056 57    zip_buffer = BytesIO()58    with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zip_file:59        for file in files:60            if file and file.filename:61                print(f"Processing {file.filename}...")62                try:63                    processed_image = process_image(file.stream)64                    img_buffer = BytesIO()65                    processed_image.save(img_buffer, format='PNG')66                    img_buffer.seek(0)67                    new_filename = f"processed_{os.path.splitext(file.filename)[0]}.png"68                    zip_file.writestr(new_filename, img_buffer.read())69                except Exception as e:70                    print(f"Error processing {file.filename}: {e}")71                    error_info = f"Failed to process {file.filename}. Error: {e}"72                    zip_file.writestr(f"error_{file.filename}.txt", error_info)73 74    zip_buffer.seek(0)75    return send_file(76        zip_buffer,77        mimetype='application/zip',78        as_attachment=True,79        download_name='processed_images.zip'80    )81 82if __name__ == '__main__':83    app.run(host='0.0.0.0', port=5000, debug=True)