# -*- coding: utf-8 -*-
# @Author: wangkun
# @Time: 2023/7/21
import cv2


class MaskWatermark:
    @classmethod
    def find_watermark(cls, image_path):
        """
        基于OpenCV自动识别水印并获取其位置
        :param image_path:水印
        :return:watermark_area
        """
        # 读取图像
        image = cv2.imread(image_path)

        # 将图像转换为灰度图像
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 使用边缘检测算法如Canny来检测图像中的边缘信息
        edges = cv2.Canny(gray, 100, 200)

        # 使用霍夫变换检测直线，以获得边缘中的直线段
        lines = cv2.HoughLinesP(edges, 1, cv2.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)

        # 对检测到的直线进行筛选，以识别可能表示水印的直线
        watermark_lines = []
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                # 根据实际情况确定水印直线的特征，例如长度、斜率等
                # 这里只是一个示例，您需要根据具体情况进行调整
                if abs(y2 - y1) < 5 and abs(x2 - x1) > 50:
                    watermark_lines.append(line)

        # 根据检测到的水印直线，计算水印区域的位置和大小
        if len(watermark_lines) > 1:
            x_coords = [line[0][0] for line in watermark_lines] + [line[0][2] for line in watermark_lines]
            y_coords = [line[0][1] for line in watermark_lines] + [line[0][3] for line in watermark_lines]
            min_x = min(x_coords)
            max_x = max(x_coords)
            min_y = min(y_coords)
            max_y = max(y_coords)
            watermark_area = (min_x, min_y, max_x - min_x, max_y - min_y)  # 水印区域的位置和大小
        else:
            watermark_area = None

        return watermark_area

    @classmethod
    def mask_watermark(cls, input_path, output_path, watermark_area):
        # 读取视频
        video = cv2.VideoCapture(input_path)

        # 获取视频的宽度和高度
        width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 创建输出视频对象
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 可根据需要更换视频编码器
        output = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

        while True:
            ret, frame = video.read()

            if not ret:
                break

            # 在水印区域替换成其他像素或进行遮挡处理
            x, y, w, h = watermark_area
            frame[y:y + h, x:x + w] = 0  # 这里将水印区域像素设为0，可根据需要进行更复杂的像素替换或遮挡处理

            # 将处理后的帧写入输出视频
            output.write(frame)

        # 释放资源
        video.release()
        output.release()

        print("成功去除水印，并保存为", output_path)

    @classmethod
    def remove_watermark(cls, video_path, output_path):
        # 读取视频
        video = cv2.VideoCapture(video_path)

        # 获取视频的宽度和高度
        width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 创建输出视频对象
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 可根据需要更换视频编码器
        output = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

        # 读取第一帧作为背景帧
        ret, background = video.read()
        if not ret:
            print("无法读取背景帧")
            return

        while True:
            ret, frame = video.read()

            if not ret:
                break

            # 计算当前帧与背景帧的差值
            diff = cv2.absdiff(frame, background)

            # 将差值转换为灰度图像
            gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)

            # 应用阈值二值化，通过调整阈值以过滤差异
            _, threshold = cv2.threshold(gray, 25, 255, cv2.THRESH_BINARY)

            # 进行形态学操作，填充小区域，平滑边缘
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
            dilated = cv2.dilate(threshold, kernel, iterations=3)

            # 寻找轮廓
            contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            # 去除检测到的水印轮廓区域
            for contour in contours:
                # 这里只是一个示例，您可以根据具体情况进行调整，例如根据轮廓面积、形状等进行过滤
                if cv2.contourArea(contour) > threshold_area:
                    # 在当前帧上用背景帧进行填充
                    cv2.drawContours(frame, [contour], -1, (0, 0, 0), cv2.FILLED)

            # 将处理后的帧写入输出视频
            output.write(frame)

        # 释放资源
        video.release()
        output.release()


if __name__ == "__main__":
    # 示例调用
    image_path = 'image.jpg'  # 替换为待识别水印的图像路径
    watermark_area = MaskWatermark.find_watermark(image_path)
    print("水印区域的位置和大小：", watermark_area)

    # 示例调用
    input_path = 'input.mp4'  # 替换为输入视频文件路径
    output_path = 'output.mp4'  # 替换为输出视频文件路径
    watermark_area = (100, 100, 200, 200)  # 替换为水印区域的位置和大小，表示为 (x, y, width, height)
    MaskWatermark.mask_watermark(input_path, output_path, watermark_area)

    # 示例调用
    video_path = 'video.mp4'  # 替换为视频文件路径
    output_path = 'output.mp4'  # 替换为输出视频文件路径
    threshold_area = 1000  # 轮廓区域的阈值，根据具体情况进行调整
    MaskWatermark.remove_watermark(video_path, output_path)
    print("成功去除水印，并保存为", output_path)

    pass