python实战之去除视频水印&；字幕_python 去除视频水印

import os
import sys
import cv2
import numpy
from moviepy import editor

VIDEO_PATH = ‘video’
OUTPUT_PATH = ‘output’
TEMP_VIDEO = ‘temp.mp4’

class WatermarkRemover():

def __init__(self, threshold: int, kernel_size: int):
    self.threshold = threshold  # 阈值分割所用阈值
    self.kernel_size = kernel_size  # 膨胀运算核尺寸


#根据用户手动选择的ROI（Region of Interest，感兴趣区域）框选水印或字幕位置。
def select_roi(self, img: numpy.ndarray, hint: str) -> list:
    '''
框选水印或字幕位置，SPACE或ENTER键退出
:param img: 显示图片
:return: 框选区域坐标
'''
    COFF = 0.7
    w, h = int(COFF * img.shape[1]), int(COFF * img.shape[0])
    resize_img = cv2.resize(img, (w, h))
    roi = cv2.selectROI(hint, resize_img, False, False)
    cv2.destroyAllWindows()
    watermark_roi = [int(roi[0] / COFF), int(roi[1] / COFF), int(roi[2] / COFF), int(roi[3] / COFF)]
    return watermark_roi


#对输入的蒙版进行膨胀运算，扩大蒙版的范围
def dilate_mask(self, mask: numpy.ndarray) -> numpy.ndarray:

    '''
对蒙版进行膨胀运算
:param mask: 蒙版图片
:return: 膨胀处理后蒙版
'''
    kernel = numpy.ones((self.kernel_size, self.kernel_size), numpy.uint8)
    mask = cv2.dilate(mask, kernel)
    return mask

#根据手动选择的ROI区域，在单帧图像中生成水印或字幕的蒙版。
def generate_single_mask(self, img: numpy.ndarray, roi: list, threshold: int) -> numpy.ndarray:
    '''
通过手动选择的ROI区域生成单帧图像的水印蒙版
:param img: 单帧图像
:param roi: 手动选择区域坐标
:param threshold: 二值化阈值
:return: 水印蒙版
'''
    # 区域无效，程序退出
    if len(roi) != 4:
        print('NULL ROI!')
        sys.exit()

    # 复制单帧灰度图像ROI内像素点
    roi_img = numpy.zeros((img.shape[0], img.shape[1]), numpy.uint8)
    start_x, end_x = in