ext/image_watermark_handler.py

import io
import os
import shutil

import cv2 as cv
import numpy as np
import requests
from PIL import Image


class ImageWatermarkHandler:

    def __init__(self):
        pass

    def index_watermark(self, img_path, block_size, step_len, threshold, temp_output_path):
        """
        该代码首先读取图像，并定义了块的大小和阈值。然后，它遍历图像的每个块，并计算每个块的标准差。
        如果标准差小于阈值，则认为当前块变化微小。然后将图片保存到本地（可选）并返回 list
        :param img_path: img路径
        :param watermark_size: 定义块的大小
        :param step_len: 遍历的步长
        :param threshold: 定义阈值
        :param temp_output_path: sub_img临时保存路径
        :return:
        """
        # read img
        img = cv.imread(img_path, cv.IMREAD_GRAYSCALE)  # 把图像转成单通道的灰度图输出

        # get width and height
        height, width = img.shape
        print("cur gray img width: %s,height: %s,block_size:%s" % (width, height, block_size))
        block_num = int(height * width // block_size)
        print("total split block num :  %s" % (block_num))

        # remove last res dir
        # 不保存图片就不用创建文件夹
        # if (os.path.exists(temp_output_path)):
        #     shutil.rmtree(temp_output_path)
        # os.mkdir(temp_output_path)

        # save pixel index  to memory and lcoal file
        list = []
        # foreach block
        for i in range(0, height, step_len):
            for j in range(0, width, step_len):
                # get pixel value
                block = img[i:i + block_size, j:j + block_size]
                # print("cur idx [%s,%s], block : %s " %(i,j,block))

                # calculate std_dev
                std_dev = np.std(block)
                # print(" cur std_dev :{} ,cur threshold : {} ".format(std_dev, threshold))  # 测试的像素区域,w:45-65--->com

                # 如果标准差小于阈值，则认为当前块变化微小
                if std_dev < threshold and std_dev > 0:
                    # save memory
                    dict = {}
                    dict['w'] = j
                    dict['h'] = i
                    list.append(dict)
                    # save local file
                    f = temp_output_path + "{}-{}.png".format(j, i)
                    print("save split img =====> w : %s ,h : %s ,cur std_dev : %s,cur threshold : %s ".format(j, i,
                                                                                                              std_dev,
                                                                                                              threshold))  # 测试的像素区域,w:45-65--->com
                    # 可以不保存图片
                    # cv.imwrite(f, block)
        return list

    def get_mask(self, img_path, list, block_size, mask_img_path):
        """
        获取mask
        :param img_path:
        :param list:
        :param block_size:
        :param mask_img_path:
        :return:
        """
        img = cv.imread(img_path, cv.IMREAD_COLOR)
        # black color
        img[:] = 0
        for item in list:
            w = int(item.get("w"))
            h = int(item.get("h"))
            x1, y1 = w, h  # 左上角坐标
            x2, y2 = w + block_size, h + block_size  # 右下角坐标
            # white color
            img[y1:y2, x1:x2] = 255, 255, 255
            # print(img[y1,x1])
        # save
        cv.imwrite(mask_img_path, img)
        return img