import glob
import os
import time
import numpy as np
import easyocr
import cv2


class WatermarksRemover:
    def __init__(self, words='', image_save_folder='data/results', show_process=False):
        self.ocr = easyocr.Reader(['ch_sim', 'en'])  # Initialize EasyOCR reader with Chinese and English languages
        self.words = words
        self.image_save_folder = image_save_folder
        self.show_process = show_process

    @staticmethod
    def read_image(image_path):
        return cv2.imread(image_path)

    def repair(self, image_path, expand=True, expand_l=10, expand_r=10, expand_t=10, expand_b=10):
        st = time.time()
        image = self.read_image(image_path)
        mask = np.zeros(image.shape[:2], np.uint8)

        results = self.ocr.readtext(image)

        for result in results:
            text_box_position = np.array(result[0], dtype=np.int32).reshape((-1, 1, 2))
            if expand:
                self.adjust_text_box(text_box_position, expand_l, expand_r, expand_t, expand_b)
            self.add_to_mask(mask, text_box_position)

        radius = 5
        result_image = cv2.inpaint(image, mask, radius, cv2.INPAINT_TELEA)
        # image save
        self.save_images(image_path, result_image)
        et = time.time()
        return round(et - st, 4)

    @staticmethod
    def adjust_text_box(text_box_position, expand_l, expand_r, expand_t, expand_b):
        # 水平方向调整
        text_box_position[1][0][0] += expand_r  # 右上角
        text_box_position[2][0][0] += expand_r  # 右下角
        text_box_position[0][0][0] -= expand_l  # 左上角
        text_box_position[3][0][0] -= expand_l  # 左下角
        # 垂直方向调整
        text_box_position[0][0][1] -= expand_t  # 左上角
        text_box_position[1][0][1] -= expand_t  # 右上角
        text_box_position[2][0][1] += expand_b  # 右下角
        text_box_position[3][0][1] += expand_b  # 左下角

    @staticmethod
    def add_to_mask(mask, text_box_position):
        cv2.fillPoly(mask, [text_box_position], (255, 255, 255))

    def save_images(self, original_image_path, result_image):
        folder = self.image_save_folder

        if not os.path.exists(folder):
            os.makedirs(folder)

        base_name = os.path.splitext(os.path.basename(original_image_path))[0]
        result_image_filename = f"{base_name}_after.jpg"
        cv2.imwrite(os.path.join(folder, result_image_filename), result_image)

    def get_image_filenames(self):
        image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.gif', '*.bmp']
        image_files = []
        for extension in image_extensions:
            image_files.extend(glob.glob(f"{self.image_save_folder}/{extension}"))
        return sorted(image_files)


# 确保结果保存文件夹存在
output_folder = 'data/results'
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 创建WatermarksRemover实例
watermarks_remover = WatermarksRemover(words='上饶旅行', image_save_folder=output_folder, show_process=False)

# 提供图片地址
image_path = 'y.jpg'

# 去除图片上的文字
time_taken = watermarks_remover.repair(image_path)
print(f"Time taken to remove watermark: {time_taken} seconds")