Video Watermarker Remover 
import threading
import queue
import os

import customtkinter


from watermark_remover import WatermarkRemover

# customtkinter.set_ctk_parent_class(tkinterDnD.Tk)

customtkinter.set_appearance_mode("System")  # Modes: "System" (standard), "Dark", "Light"
customtkinter.set_default_color_theme("blue")  # Themes: "blue" (standard), "green", "dark-blue"


# tkinter GUI example with threading and queue
# https://gist.github.com/roosnic1/f1d1d17c663476af3279ab6ae3e80206

"""
The following code is for non-UI function.

"""

def remove_watermark(threshold, kernel_size, sourceFolder, outputFolder):
    remover = WatermarkRemover(threshold, kernel_size, sourceFolder, outputFolder)
    outputVideoPath = remover.remove_video_watermark()

    # print("Add outputVideoPath to the Queue")
    app.queue_UI.put(outputVideoPath, block=False, timeout=None)
    # queue_UI.put_nowait(outputVideoPath) # same as queue_UI.put(outputVideoPath, block=False, timeout=None)


# get each files in the folder
def get_files_in_folder(folder_path):
    
    temp_files = []

    for file in os.listdir(folder_path):
        temp_files.append(os.path.join(folder_path, file))
        # print(f"Source path: {folder_path} , file: {file}", os.path.join(folder_path, file))
    return temp_files

# check if the folder is exist
def check_folder_exist(folder_path):
    if os.path.exists(folder_path):
        return True
    else:
        return False


    


"""
The following code is for sub-thread function.

"""
def start_new_thread(targetSubThread, *args):
    newThreadTemp = threading.Thread(target=targetSubThread, args=args)
    newThreadTemp.start()
    return newThreadTemp



class App(customtkinter.CTk):
    def __init__(self):
        super().__init__()


        self.title("Video Watermark Removal")
        self.minsize(900, 700)
        self.geometry("900x780")

        self.queue_UI = queue.Queue()



        """
        The following code is for layout components

        """

        self.grid_rowconfigure(0, weight=1)  # configure grid system
        self.grid_columnconfigure((0, 1), weight=1)
    
        frame1 = customtkinter.CTkFrame(master=self)
        frame1.grid(row=0, column=0, columnspan=4, pady=6, padx=6, sticky="nsew")
        self.frame1 = frame1


        self.label_threshold = customtkinter.CTkLabel(master=frame1, justify=customtkinter.LEFT, text="灰度阈值 threshold", font=("Arial", 16))
        self.label_threshold.grid(row=0, column=0, pady=6, padx=6, sticky="ens")

        self.entry_threshold = customtkinter.CTkEntry(master=frame1, placeholder_text="80")
        self.entry_threshold.grid(row=0, column=1, pady=6, padx=6, sticky="w")
        self.entry_threshold.insert(0, int(80))

        self.slider_threshold = customtkinter.CTkSlider(master=frame1, command=self.slider_threshold_callback, from_=0, to=255)
        self.slider_threshold.grid(row=0, column=2, columnspan=2, pady=6, padx=0, sticky="ew")
        self.slider_threshold.set(80)



        self.label_kernel_size = customtkinter.CTkLabel(master=frame1, justify=customtkinter.LEFT, text="边缘膨胀系数 kernel_size", font=("Arial", 16))
        self.label_kernel_size.grid(row=1, column=0, pady=6, padx=6, sticky="ens")

        self.entry_kernel_size = customtkinter.CTkEntry(master=frame1, placeholder_text="6")
        self.entry_kernel_size.grid(row=1, column=1, pady=6, padx=6, sticky="w")
        self.entry_kernel_size.insert(0, int(6))

        self.slider_kernel_size = customtkinter.CTkSlider(master=frame1, command=self.slider_kernel_size_callback, from_=0, to=50)
        self.slider_kernel_size.grid(row=1, column=2, columnspan=2, pady=6, padx=0, sticky="ew")
        self.slider_kernel_size.set(6)



        self.label_source_video = customtkinter.CTkLabel(master=frame1, justify=customtkinter.LEFT, text="源目录 Video source folder:", font=("Arial", 16))
        self.label_source_video.grid(row=2, column=0, pady=6, padx=6, sticky="ens")

        self.entry_source_video = customtkinter.CTkEntry(master=frame1, placeholder_text="")
        self.entry_source_video.grid(row=2, column=1, columnspan=2, pady=3, padx=1, sticky="ew")
        self.entry_source_video.insert(0, "video")

        self.button_source_video = customtkinter.CTkButton(master=frame1, command=self.button_video_source_callback, text="Select Folder")
        self.button_source_video.grid(row=2, column=3, pady=3, padx=6, sticky="w")



        self.label_output_video = customtkinter.CTkLabel(master=frame1, justify=customtkinter.LEFT, text="输出目录 Video output folder:", font=("Arial", 16))
        self.label_output_video.grid(row=3, column=0, pady=6, padx=6, sticky="ens")

        self.entry_output_video = customtkinter.CTkEntry(master=frame1, placeholder_text="")
        self.entry_output_video.grid(row=3, column=1, columnspan=2, pady=3, padx=1, sticky="ew")
        self.entry_output_video.insert(0, "output")

        self.button_output_video = customtkinter.CTkButton(master=frame1, command=self.button_video_output_callback, text="Select Folder")
        self.button_output_video.grid(row=3, column=3, pady=3, padx=6, sticky="w")




        self.button_start = customtkinter.CTkButton(master=frame1, command=self.button_start_callback, text="Remove Watermark")
        self.button_start.grid(row=4, column=1, pady=3, padx=10, sticky="ew")

        self.button_log = customtkinter.CTkButton(master=frame1, command=self.button_log_callback, text="Clear log")
        self.button_log.grid(row=4, column=2, pady=3, padx=10, sticky="ew")


        self.progressbar_1 = customtkinter.CTkProgressBar(master=frame1, orientation="horizontal", mode="determinate", determinate_speed=0.01, width=600)
        self.progressbar_1.grid(row=5, column=0, columnspan=4,  pady=10, padx=10)
        self.progressbar_1.set(0)  


        self.text_log = customtkinter.CTkTextbox(master=frame1, width=850, height=500)
        self.text_log.grid(row=6, column=0, columnspan=4, pady=10, padx=10)

        self.text_log.insert("0.0", "kernel_size 膨胀运算核尺寸, 范围为所有正整数, 用于处理水印或字幕边缘, 数值越大最后生成的蒙版覆盖面积边缘越大 \n")
        self.text_log.insert("0.0", "threshold 阈值分割灰度值，范围0~255, 大于该灰度值的水印将被处理, 数字越小被去掉的水印越多, 如果水印为亮白色该值数字尽量大一些,这样去掉的画面更少画质更好, 如果水印较暗该值就要小些才去掉水印,同时画面也会被去掉更多信息 \n")
        self.text_log.insert("0.0", "参数说明: \n")
        self.text_log.insert("0.0", " \n")

        self.update()


    """
    The following code is for UI event action callback 
    Example: button click, slider change, etc.

    """

    def button_start_callback(self):
        temp_threshold = int(self.entry_threshold.get())
        temp_kernel_size = int(self.entry_kernel_size.get())

        temp_sourceFolder = self.entry_source_video.get()
        temp_outputFolder = self.entry_output_video.get()

        if not temp_sourceFolder:
            temp_sourceFolder = "video"

        if not temp_outputFolder:
            temp_outputFolder = "output"
        
        if not check_folder_exist(temp_sourceFolder):
            self.text_log.insert("0.0", "Source folder not exist ! 源目录不存在 \n")
            return

        if not check_folder_exist(temp_outputFolder):
            self.text_log.insert("0.0", "Output folder not exist ! 输出目录不存在 \n")
            return
        
        filelist = get_files_in_folder(temp_sourceFolder)


        # check filelist is empty or not
        if not filelist:
            self.text_log.insert("0.0", "Source folder is empty ! 源目录为空 \n")
            return

        self.progressbar_1.set(0)
        self.progressbar_1.start() 

        self.button_start.configure(state="disabled")
        print("\n\n")
        # print(f"Start to remove watermark, threshold: {temp_threshold},  kernel_size: {temp_kernel_size}")
        self.text_log.insert("0.0", "Initiating watermark removal. Please wait... (threshold: "+str(temp_threshold)+", kernel_size: "+str(temp_kernel_size)+") \n")

        # Start a new thread to run the watermark removal process
        newThreadRemoveWatermark = start_new_thread(remove_watermark, temp_threshold, temp_kernel_size, temp_sourceFolder, temp_outputFolder)

        self.after(1000, self.updateUI_WatermarkRemovalFinished, newThreadRemoveWatermark)


    def button_log_callback(self):
        self.text_log.delete("0.0", customtkinter.END)


    def slider_threshold_callback(self, value):
        self.entry_threshold.delete(0, customtkinter.END)
        self.entry_threshold.insert(0, int(value))

    def slider_kernel_size_callback(self, value):
        self.entry_kernel_size.delete(0, customtkinter.END)
        self.entry_kernel_size.insert(0, int(value))

    def button_video_source_callback(self):
        sourceFolder = customtkinter.filedialog.askdirectory()
        self.entry_source_video.delete(0, customtkinter.END)
        self.entry_source_video.insert(0, sourceFolder)

    def button_video_output_callback(self):
        outputFolder = customtkinter.filedialog.askdirectory()
        self.entry_output_video.delete(0, customtkinter.END)
        self.entry_output_video.insert(0, outputFolder)



    """
    The following code is for after loop callback 
    Which is used to update UI after thread finished

    """

    def updateUI_WatermarkRemovalFinished(self, subThread):

        # if subThread.is_alive():
        #     print("subThread is still running. Thread Alive Status: ", subThread.is_alive())
        #     self.after(1000, self.updateUI_WatermarkRemovalFinished, subThread)
        # else:
        #     print("subThread is finished. Thread Alive Status: ", subThread.is_alive())
        #     self.button_start.configure(state="normal")
        #     self.progressbar_1.stop()
        #     self.progressbar_1.set(1)


        try:
            # print("Try to get output_path from queue")
            # print("===== SubThread is still running. Thread Alive Status: ", subThread.is_alive())
            # print("Quene 队列是否为空: {}, 队列还有 {} 个项目, ".format(self.queue_UI.empty(), self.queue_UI.qsize()))
            output_path = self.queue_UI.get(block=False)

        except queue.Empty:
            # print("Queue is empty, let's try again later")
            self.after(2000, self.updateUI_WatermarkRemovalFinished, subThread)
            return
        
        if output_path is not None:
            self.text_log.insert("0.0", " \n")
            self.text_log.insert("0.0", "Watermark removal successful! File: "+output_path+" \n")
            self.text_log.insert("0.0", " \n")

            self.button_start.configure(state="normal")
            self.progressbar_1.stop()
            self.progressbar_1.set(1)
            output_path = ""


if __name__ == "__main__":
    app = App()
    # app1 = customtkinter.CTk()
    print("Customtkinter APP Start!  Type: ", type(app))

    app.mainloop()




gui.py
jinwyp/VideoWatermarkerRemover
import os
import sys
import argparse

import cv2
import numpy
from moviepy import editor

VIDEO_SOURCE_PATH = 'video'
VIDEO_OUTPUT_PATH = 'output'
VIDEO_OUTPUT_PATH_2 = os.path.join(os.path.pardir, 'video_bilibili')
VIDEO_OUTPUT_TEMP_VIDEO = 'temp.mp4'


class WatermarkRemover():
  
  def __init__(self, threshold: int, kernel_size: int, source_path: str, output_path: str):
    self.threshold = threshold #阈值分割所用阈值
    self.kernel_size = kernel_size #膨胀运算核尺寸
    self.source_path = source_path #视频源文件路径
    self.output_path = output_path #输出文件路径


  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

  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(roi[1]), int(roi[1] + roi[3])
    start_y, end_y = int(roi[0]), int(roi[0] + roi[2])
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    roi_img[start_x:end_x, start_y:end_y] = gray[start_x:end_x, start_y:end_y]

    #阈值分割
    _, mask = cv2.threshold(roi_img, threshold, 255, cv2.THRESH_BINARY)
    print("=== Step 2: finish to generate_single_mask of frame ")
    return mask
  

  def generate_watermark_mask(self, video_path: str) -> numpy.ndarray:
    '''
    截取视频中多帧图像生成多张水印蒙版，通过逻辑与计算生成最终水印蒙版
    :param video_path: 视频文件路径
    :return: 水印蒙版
    '''
    video = cv2.VideoCapture(video_path)
    success, frame = video.read()
    print("=== Please select the region of the watermark. Then press SPACE or ENTER to continue ")
    print("=== 请选择要去掉的水印的区域 然后按空格或回车继续 \n")
    roi = self.select_roi(frame, 'Select watermark ROI')
    generateAllFrameMask = numpy.ones((frame.shape[0], frame.shape[1]), numpy.uint8)
    generateAllFrameMask.fill(255)

    step = video.get(cv2.CAP_PROP_FRAME_COUNT) // 5
    index = 0
    print(f"=== Step 1: start  generate_watermark_mask. Total frame: {video.get(cv2.CAP_PROP_FRAME_COUNT)}, Step: {step}, CAP_PROP_FRAME_COUNT: {cv2.CAP_PROP_FRAME_COUNT} " )

    while success:
      if index % step == 0:
        generateAllFrameMask = cv2.bitwise_and(generateAllFrameMask, self.generate_single_mask(frame, roi, self.threshold))
      success, frame = video.read()
      index += 1
    video.release()
    print("=== Step 1: finish to generate_watermark_mask ")

    return self.dilate_mask(generateAllFrameMask)

  def generate_subtitle_mask(self, frame: numpy.ndarray, roi: list) -> numpy.ndarray:
    '''
    通过手动选择ROI区域生成单帧图像字幕蒙版
    :param frame: 单帧图像
    :param roi: 手动选择区域坐标
    :return: 字幕蒙版
    '''
    mask = self.generate_single_mask(frame, [0, roi[1], frame.shape[1], roi[3]], self.threshold) #仅使用ROI横坐标区域
    return self.dilate_mask(mask)

  def inpaint_image(self, img: numpy.ndarray, mask: numpy.ndarray) -> numpy.ndarray:
    '''
    修复图像
    :param img: 单帧图像
    :parma mask: 蒙版
    :return: 修复后图像
    '''

    telea = cv2.inpaint(img, mask, 1, cv2.INPAINT_TELEA)
    return telea
  
  def merge_audio(self, input_path: str, merge_output_path: str, temp_path: str):
    '''
    合并音频与处理后视频
    :param input_path: 原视频文件路径
    :param merge_output_path: 封装音视频后文件路径
    :param temp_path: 无声视频文件路径 
    '''
    with editor.VideoFileClip(input_path) as video:
      audio = video.audio
      with editor.VideoFileClip(temp_path) as opencv_video:
        clip = opencv_video.set_audio(audio)
        clip.to_videofile(merge_output_path)

  def remove_video_watermark(self):
    '''
    去除视频水印
    '''
    if not os.path.exists(self.output_path):
      os.makedirs(self.output_path)

    filenames = []

    for singleFile in os.listdir(self.source_path):
        filenames.append(os.path.join(self.source_path, singleFile))
        print(f"Source file path: {self.source_path}, filename: {singleFile},  full path: {os.path.join(self.source_path, singleFile)} \n")

    mask = None

    for i, singleSourceVideoFileName in enumerate(filenames):

      #生成水印蒙版
      mask = self.generate_watermark_mask(singleSourceVideoFileName)
      
      #创建待写入文件对象
      video = cv2.VideoCapture(singleSourceVideoFileName)
      fps = video.get(cv2.CAP_PROP_FPS)
      size = (int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)))
      video_writer = cv2.VideoWriter(VIDEO_OUTPUT_TEMP_VIDEO, cv2.VideoWriter_fourcc(*'mp4v'), fps, size)
    
      print(f"=== Step 3: Start to inpaint all frame. FPS: {fps}. Please wait... " )

      #逐帧处理图像
      success, frame = video.read()
      while success:
        frame = self.inpaint_image(frame, mask)
        video_writer.write(frame)
        success, frame = video.read()

      print("=== Step 3: finish to inpaint all frame")
      video.release()
      video_writer.release()

      #封装视频
      (_, filename) = os.path.split(singleSourceVideoFileName)
      output_path = os.path.join(self.output_path, filename.split('.')[0] + '_no_watermark.mp4')#输出文件路径
      self.merge_audio(singleSourceVideoFileName, output_path, VIDEO_OUTPUT_TEMP_VIDEO)

      print(" \n")

    if os.path.exists(VIDEO_OUTPUT_TEMP_VIDEO):
      os.remove(VIDEO_OUTPUT_TEMP_VIDEO)

    return output_path
  
  

  def remove_video_subtitle(self):
    '''
    去除视频字幕
    '''
    if not os.path.exists(self.output_path):
      os.makedirs(self.output_path)


    filenames = [os.path.join(self.source_path, i) for i in os.listdir(self.source_path)]
    roi = []

    for i, name in enumerate(filenames):
      #创建待写入文件对象
      video = cv2.VideoCapture(name)
      fps = video.get(cv2.CAP_PROP_FPS)
      size = (int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)))
      video_writer = cv2.VideoWriter(VIDEO_OUTPUT_TEMP_VIDEO, cv2.VideoWriter_fourcc(*'mp4v'), fps, size)
    
      #逐帧处理图像
      success, frame = video.read()
      if i == 0:
        roi = self.select_roi(frame, 'select subtitle ROI')

      while success:
        mask = self.generate_subtitle_mask(frame, roi)
        frame = self.inpaint_image(frame, mask)
        video_writer.write(frame)
        success, frame = video.read()

      video.release()
      video_writer.release()

      #封装视频
      (_, filename) = os.path.split(name)
      output_path = os.path.join(self.output_path, filename.split('.')[0] + '_no_sub.mp4')#输出文件路径
      self.merge_audio(name, output_path, VIDEO_OUTPUT_TEMP_VIDEO)

    if os.path.exists(VIDEO_OUTPUT_TEMP_VIDEO):
      os.remove(VIDEO_OUTPUT_TEMP_VIDEO)

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument('--threshold', type=int, default=80, help='Threshold value for watermark removal')
  parser.add_argument('--kernel_size', type=int, default=5, help='Kernel size for watermark removal')
  parser.add_argument('--subtitle', action='count', default=0, help='Subtitle watermark removal')

  parser.add_argument('--output', action='count', default=0, help='Output file path')
  
  args = parser.parse_args()
  print("====================")
  print(f"APP start! arguments: [subtitle: {args.subtitle} , threshold: {args.threshold}, output: {args.output} ] ")
  print("====================")

  if args.output > 0:
    VIDEO_OUTPUT_PATH = VIDEO_OUTPUT_PATH_2
    
  if args.subtitle > 0:
    remover = WatermarkRemover(threshold=80, kernel_size=10, source_path=VIDEO_SOURCE_PATH, output_path=VIDEO_OUTPUT_PATH)
    remover.remove_video_subtitle()
  else:
    remover = WatermarkRemover(threshold=80, kernel_size=6, source_path=VIDEO_SOURCE_PATH, output_path=VIDEO_OUTPUT_PATH)
    remover.remove_video_watermark()



  #去除视频水印
  # remover = WatermarkRemover(threshold=80, kernel_size=5)
  # remover.remove_video_watermark()

  #去除视频字幕
  # remover = WatermarkRemover(threshold=80, kernel_size=10)
  # remover.remove_video_subtitle()


watermark_remover.py
jinwyp/VideoWatermarkerRemover