# # -- encoding: utf-8 -- # # Description: # Create: 2025/01/06 # """ import time from typing import Dict, Any import cv2 import numpy as np from loguru import logger from datamate.common.utils import bytes_to_numpy from datamate.common.utils import numpy_to_bytes from datamate.core.base_op import Mapper from .watermark_ocr_model import WatermarkOcrModel DEFAULT_MAX_CHARACTERS = 10 DEFAULT_BINARY_THRESHOLD_LOW = 200 class ImgWatermarkRemove(Mapper): use_model = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.remove_str = kwargs.get("watermarkStr", "知乎,抖音") self.ocr_model = self.get_model(*args, **kwargs) @staticmethod def _has_kw(result_list, kw_list): """ 图片是否包含目标水印，返回匹配到的文字列表 """ result_str_list = [] for line in result_list: for kw in kw_list: if kw in line[1][0]: result_str_list.append(line[1][0]) break return result_str_list @staticmethod def _overlay_mask(background_img, prospect_img, img_over_x, img_over_y): back_r, back_c, _ = background_img.shape # 背景图像行数、列数 is_x_direction_failed = img_over_x > back_c or img_over_x < 0 is_y_direction_failed = img_over_y > back_r or img_over_y < 0 if is_x_direction_failed or is_y_direction_failed: # 前景图不在背景图范围内, 直接返回原图 return background_img pro_r, pro_c, _ = prospect_img.shape # 前景图像行数、列数 if img_over_x + pro_c > back_c: # 如果水平方向展示不全 pro_c = back_c - img_over_x # 截取前景图的列数 prospect_img = prospect_img[:, 0:pro_c, :] # 截取前景图 if img_over_y + pro_r > back_r: # 如果垂直方向展示不全 pro_r = back_r - img_over_y # 截取前景图的行数 prospect_img = prospect_img[0:pro_r, :, :] # 截取前景图 prospect_img = cv2.cvtColor(prospect_img, cv2.COLOR_BGR2BGRA) # 前景图转为4通道图像 prospect_tmp = np.zeros((back_r, back_c, 4), np.uint8) # 与背景图像等大的临时前景图层 # 前景图像放到前景图层里 prospect_tmp[img_over_y:img_over_y + pro_r, img_over_x: img_over_x + pro_c, :] = prospect_img _, binary = cv2.threshold(prospect_img, 254, 255, cv2.THRESH_BINARY) # 前景图阈值处理 prospect_mask = np.zeros((pro_r, pro_c, 1), np.uint8) # 单通道前景图像掩模 prospect_mask[:, :, 0] = binary[:, :, 3] # 不透明像素的值作为掩模的值 mask = np.zeros((back_r, back_c, 1), np.uint8) mask[img_over_y:img_over_y + prospect_mask.shape[0], img_over_x: img_over_x + prospect_mask.shape[1]] = prospect_mask mask_not = cv2.bitwise_not(mask) prospect_tmp = cv2.bitwise_and(prospect_tmp, prospect_tmp, mask=mask) background_img = cv2.bitwise_and(background_img, background_img, mask=mask_not) prospect_tmp = cv2.cvtColor(prospect_tmp, cv2.COLOR_BGRA2BGR) # 前景图层转为三通道图像 return prospect_tmp + background_img # 前景图层与背景图像相加合并 def execute(self, sample: Dict[str, Any]): start = time.time() file_name = sample[self.filename_key] file_type = "." + sample[self.filetype_key] img_bytes = sample[self.data_key] if img_bytes: data = bytes_to_numpy(img_bytes) correct_data = self._watermark_remove(data, file_name, self.ocr_model) sample[self.data_key] = numpy_to_bytes(correct_data, file_type) logger.info(f"fileName: {file_name}, method: ImgWatermarkRemove costs {time.time() - start:6f} s") return sample def delete_watermark(self, result_list, kw_list, data): """ 将符合目标的水印，模糊化处理 """ # 获取所有符合目标的文本框位置 text_axes_list = [] for line in result_list: for kw in kw_list: if kw in line[1][0]: min_width = int(min(line[0][0][0], line[0][3][0])) max_width = int(max(line[0][1][0], line[0][2][0])) min_hight = int(min(line[0][0][1], line[0][1][1])) max_hight = int(max(line[0][2][1], line[0][3][1])) text_axes_list.append([min_width, min_hight, max_width, max_hight]) break # 去除水印 delt = DEFAULT_MAX_CHARACTERS # 文本框范围扩大 img = data for text_axes in text_axes_list: hight, width = img.shape[0:2] # 截取图片 min_width = text_axes[0] - delt if text_axes[0] - delt >= 0 else 0 min_hight = text_axes[1] - delt if text_axes[1] - delt >= 0 else 0 max_width = text_axes[2] + delt if text_axes[2] + delt <= width else width max_hight = text_axes[3] + delt if text_axes[3] + delt <= hight else hight cropped = img[min_hight:max_hight, min_width:max_width] # 裁剪坐标为[y0:y1, x0:x1] # 图片二值化处理，把[200,200,200]-[250,250,250]以外的颜色变成0 start_rgb = DEFAULT_BINARY_THRESHOLD_LOW thresh = cv2.inRange(cropped, np.array([start_rgb, start_rgb, start_rgb]), np.array([250, 250, 250])) # 创建形状和尺寸的结构元素 kernel = np.ones((3, 3), np.uint8) # 设置卷积核3*3全是1；将当前的数组作为图像类型来进⾏各种操作，就要转换到uint8类型 # 扩展待修复区域 hi_mask = cv2.dilate(thresh, kernel, iterations=10) # 膨胀操作，白色区域增大，iterations迭代次数 specular = cv2.inpaint(cropped, hi_mask, 5, flags=cv2.INPAINT_TELEA) # imgSY：输入8位1通道或3通道图像。 # hi_mask：修复掩码，8位1通道图像。非零像素表示需要修复的区域。 # specular：输出与imgSY具有相同大小和类型的图像。 # 5：算法考虑的每个点的圆形邻域的半径。 # flags：NPAINT_NS基于Navier-Stokes的方法、Alexandru Telea的INPAINT_TELEA方法 result = self._overlay_mask(img, specular, min_width, min_hight) img = result return img def init_model(self, *args, **kwargs): return WatermarkOcrModel(*args, **kwargs).ocr_model def _watermark_remove(self, data, file_name, model): """ 去除水印的方法 """ remove_str = self.remove_str # 勾选去水印的信息为空，则直接返回原图 if remove_str == "": return data kw_list = remove_str.split(',') # 加载模型 ocr_model = model try: result = ocr_model.ocr(data, cls=True) except RuntimeError as e: logger.error(f"fileName: {file_name}, method: ocr predict error {e}") return data if result and result[0]: logger.info(f"fileName: {file_name}, method: ocrModel detect watermark info {str(result)}") return self.delete_watermark(result[0], kw_list, data) else: logger.info(f"fileName: {file_name}, method: ImgWatermarkRemove not need remove target ocr") return data