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 = int