使用cv为批量图片自定义添加水印logo

一、使用cv为批量图片自定义添加文字/水印logo

需求：需要大量带有各式水印、二维码、自定义文字的图片，水印、二维码 大小不固定、位置不固定、模糊程度不固定

功能：自定义为图片添加 图片logo/文字
    1、添加水印
        ① 随机选取水印logo 添加的位置可以随机插入【“四角”，“去心域”】，也可以指定固定位置插入水印logo
        ② 可以自定义logo水印插入 “四角”、“去心域”的大小占比
        ③ logo水印大小、模糊程度自定义
    2、添加文本
        。。。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cv2
import os
import numpy as np
import time
import argparse
from tqdm import tqdm


class AddLogoText():
    def __init__(self,source_imgs_dir,
                 wm_imgs_dir=None,
                 text_path=None,
                 report_dir=None,
                 alpha=0.3,
                 position=None,
                 random=True,
                 insert_type='move_heart',
                 type_ = 'img',*args,**kwargs):

        self.source_imgs_dir = source_imgs_dir # 原图片的目录路径，str
        self.wm_imgs_dir = wm_imgs_dir # logo或者二维码的目录路径，logo/二维码图片不可是灰度图，str
        self.text_path = text_path # 图片写文字的文件路径
        self.report_dir = report_dir # 图片合成后输入的文件夹
        self.alpha = alpha # logo 插入的模糊度，1.0 表示完全不模糊
        self.position = position # loge、二维码插入的像素位置元祖,使用时需要将 random 置为 False，tuple
        self.random = random # 是否随机插入，开启时position无效 ，bool
        if not random:
            assert position is not None
        self.insert_type = insert_type # random = True时有效，随机插入的类型；'move_heart' 随机插入至原图片去心域；'four_corn' 随机插入至原图四个角的位置
        assert insert_type in ['move_heart','four_corn']
        self.type_ = type_ # "img" ：在图片中 插入 loge ， "text"：在图片中插入文本
        assert type_ in ['img','text']

        # 可变参数
        self.size_height = 70  # logo/二维码 resize 插入原图片时的高度的像素大小，int
        self.percent=0.33  # logo/二维码插入图片时 不能越界【左上 左下 右上 右下】 percent占比的位置 【仅仅在 随机插入 即random=True 有效 】，float
        assert 0.0 < self.percent < 0.5

    def read_image(self,fpath):
        """
        cv 加载图片，返回图片的数组，如果图片为 灰度图，flag = False，否则返回 true
        :return:
        """
        flag = True
        img_bgr = cv2.imread(fpath)
        
        if img_bgr is None:
            raise Exception("图片读取失败，可能图片路径中存在中文字符或者文件本身非cv可以读取为图片的数据类型，图片路径为{}".format(fpath))
            
        dim = img_bgr.ndim
        if dim == 2:
            flag = False
        return img_bgr,flag

    def load_images(self,source_dir):
        """
        给定文件夹，加载数据
        :return:
        """
        assert source_dir is not None

        father_dir_paths = []
        file_paths = []
        for root, dirs, files in os.walk(source_dir):
            for file in files:
                # 获取文件所属目录
                father_dir_paths.append(root)
                # 获取文件路径
                file_paths.append(os.path.join(root, file))
        return file_paths, father_dir_paths

    def time_stamp_str(self,str_):
        """
        为了保证字符串的唯一性，给字符串末尾添加时间戳
        :param str_:
        :return:
        """
        stamp_ = str(time.time()).replace('.','')
        str_ = str(str_) + stamp_
        return str_

    def cv_save_img(self,array,img_name,report_dir):
        """
        用 cv2 将 三通道 BGR array 保证为 img_name，保存在 report_dir中
        :param array: bgr 图片的数组
        :param img_name: 图片的名字，CV2文件中不可以有中文
        :param report_dir: 图片保存到的目录中, CV2 路径中不可以有中文
        :return:
        """
        assert report_dir is not None
        # 判断目录是否存在
        if not os.path.exists(report_dir):
            os.makedirs(report_dir)

        img_path = os.path.join(report_dir,img_name)
        cv2.imwrite(img_path,array)

    # TODO 核心函数
    def add_logo_single_img(self,source_img_fpath,
                              wm_img_fpath,
                              alpha=0.6,
                              size_height=100,
                              percent=0.3,
                              position=None,
                              random=True,
                              insert_type='move_heart'):

        """
        将一张图片 与 loge 通过cv2.addWeighted 的方式进行融合，本函数不处理原图为灰度图的情况
        tip：融合的特点，loge/二维码浅色的位置会随着深色的位置一起淡化或加深【如果想要loge/二维码深色或浅色的部分，其他舍弃掉需要用到 cv位运算 】

        :param source_img_path: 原图片的文件路径，str
        :param wm_img_path: logo或者二维码的文件路径，logo/二维码图片不可是灰度图，str
        :param alpha: logo 加入图片的模糊状态，最大 1.0，即不做模糊处理，float
        :param size_height: logo/二维码 resize 插入原图片时的高度的像素大小，int
        :param percent: logo/二维码插入图片时 不能越界【左上 左下 右上 右下】 percent占比的位置 【仅仅在 随机插入 即random=True 有效 】，float
        :param position: logo 插入的 左上角的像素位置，tuple
        :param random: 插入位置是否随机，bool
        :param insert_type: # random = True时有效，随机插入的类型；'move_heart' 随机插入至原图片去心域；'four_corn' 随机插入至原图四个角的位置
        :return: s_img,position; s_img:插入logo后图片的 bgr数组，position 插入logo的 左上角定点的 位置
        """

        if not random:
            assert position is not None
        # 加载 img logo图
        s_img,s_flag = self.read_image(source_img_fpath)
        s_img_h,s_img_w = s_img.shape[:2]
        wm_img,wm_flag = self.read_image(wm_img_fpath)
        tmp1_h, tmp1_w = wm_img.shape[:2]

        # bgr 通道执行
        if s_flag and wm_flag:
            # 将 logo/二维码 resize
            wm_logo = cv2.resize(wm_img,dsize=(int(tmp1_w/(tmp1_h*1.0/size_height)),size_height),interpolation=cv2.INTER_LINEAR)
            wm_logo_h,wm_logo_w = wm_logo.shape[:2]
            # print(s_img_h,s_img_w,wm_logo_h,wm_logo_w)

            # 判断 logo、二维码的尺寸小于图片尺寸
            if (wm_logo_h < s_img_h) and (wm_logo_w < s_img_w):

                # 如果是随机插入图片,需要生成随机插入点 ， 否则 使用position的地址作为图片的插入点
                if random:
                    if insert_type == 'four_corn':
                        # TODO 生成随机的 '四角'的顶点
                        percent_h_top = int(s_img_h * 1.0 * percent)
                        percent_h_tail = int(s_img_h * 1.0 * (1-percent))
                        percent_w_l = int(s_img_w * 1.0 * percent)
                        percent_w_r = int(s_img_w * 1.0 * (1-percent))

                        if np.random.uniform() < 0.5: # 如果随机数小于 0.5 判断顶点在上方，否则在下方
                            ran_h = np.random.randint(low=0,high=percent_h_top)
                        else:
                            ran_h = np.random.randint(low=percent_h_tail, high=s_img_h)
                        if np.random.uniform() < 0.5: # 如果随机数小于 0.5 判断顶点在左方，否则在右方
                            ran_w = np.random.randint(low=0,high=percent_w_l)
                        else:
                            ran_w = np.random.randint(low=percent_w_r, high=s_img_w)
                        position = (ran_h,ran_w)
                    elif insert_type == 'move_heart':
                        # TODO 生成随机的 去除'心域'的顶点
                        percent_h_top = int(s_img_h * 1.0 * percent)
                        percent_h_tail = int(s_img_h * 1.0 * (1 - percent))
                        percent_w_l = int(s_img_w * 1.0 * percent)
                        percent_w_r = int(s_img_w * 1.0 * (1 - percent))

                        ran_h = np.random.randint(low=0,high=s_img_h)
                        if percent_h_tail > ran_h >= percent_h_top: # 当 高度h 落在中间的 区域时，去除 '心域'
                            if np.random.uniform() < 0.5:  # 如果改值小于 0.5，选择 “左边”，否则选择 “右边”
                                ran_w = np.random.randint(low=0,high=percent_w_l)
                            else:
                                ran_w = np.random.randint(low=percent_w_r,high=s_img_w)
                        else:
                            ran_w = np.random.randint(low=0,high=s_img_w)
                        position = (ran_h,ran_w)

                # 判断 position 插入后是否越界，如果越界，将logo、二维码左移/上移 至极限位置
                start_h,start_w = position
                if start_h + wm_logo_h > s_img_h:
                    start_h = s_img_h - wm_logo_h
                if start_w + wm_logo_w > s_img_w:
                    start_w = s_img_w - wm_logo_w
                position = (start_h,start_w)

                # 将图片的数据 深拷贝一份，用于 cv2.addWeighted
                s_img_copy = s_img.copy()
                # 将loge、二维码数据嵌入
                s_img_copy[start_h:start_h + wm_logo_h , start_w:start_w + wm_logo_w ,...] = wm_logo
                # 图片融合
                s_img = cv2.addWeighted(s_img,1.0-alpha,s_img_copy,alpha,gamma=0)

                return s_img,position

            else:
                print('当前logo/二维码尺寸为h：{} w：{}，原图尺寸太小，不做处理'.format(wm_logo_h,wm_logo_w))
                return s_img,False

        else: # 原图为灰度图
            print('原图为灰度图，不做处理')
            return s_img,False

    def add_logo_allImgs(self):
        """
        通过 cv2.addWeighted 的方式将 所有的源图片 与 随机选择的logo 进行融合
        :return:
        """
        source_img_paths,_ = self.load_images(self.source_imgs_dir)
        wm_img_paths,_ = self.load_images(self.wm_imgs_dir)

        print('logo 嵌入中，一共{}张源图片，{}张logo图片'.format(len(source_img_paths),len(wm_img_paths)))
        for idx,img in tqdm(enumerate(source_img_paths)):
            # 随机取一个 二维码/loge图片
            wm_img_fpath = np.random.choice(wm_img_paths)
            img_bgr,_ = self.add_logo_single_img(source_img_fpath=img,
                                       wm_img_fpath=wm_img_fpath,
                                       alpha=self.alpha,
                                       size_height=self.size_height,
                                       percent=self.percent,
                                       position=self.position,
                                       random=self.random,
                                       insert_type=self.insert_type)

            img_name = self.time_stamp_str(str( idx+1) + '_' ) + '.jpg'
            # 保存图片
            self.cv_save_img(img_bgr,img_name,self.report_dir)

    def add_text_single_img(self):
        """
        给一张图片添加文本
        :return:
        """

    def add_text_allImgs(self):
        """
        给所有的图片添加文本
        :return:
        """
        pass

    def run(self):
        if self.type_ == 'img':
            print('当前模式为嵌入logo的模式！')
            self.add_logo_allImgs()
        elif self.type_ == "text":
            print('当前模式为嵌入文本的模式！')
            self.add_text_allImgs()


if __name__ == '__main__':

    """
    # 回调函数
    def str2bool(v):
        if isinstance(v, bool):
            return v
        if v.lower() in ('yes', 'true', 't', 'y', '1'):
            return True
        elif v.lower() in ('no', 'false', 'f', 'n', '0'):
            return False
        else:
            raise argparse.ArgumentTypeError('Boolean value expected.')


    def str2tuple(v):
        str_tuple = ''.join(v)
        return eval(str_tuple)


    parser = argparse.ArgumentParser()
    parser.add_argument('--source_dir', help='source imgs', type=str)
    parser.add_argument('--wm_dir', help='logo imgs', type=str)
    parser.add_argument('--text_path', help='text file path', type=str)
    parser.add_argument('--report_dir', help='report imgs path', type=str)
    parser.add_argument('--type_', help='insert img or text', type=str, default='img')
    parser.add_argument('--alpha', help='alpha', type=float, default=0.3)
    parser.add_argument('--position', help='position tuple logo or text, eg: (20,20)', type=str2tuple, default=(0, 0))
    parser.add_argument('--insert_type', help='insert logo or text type', type=str, default='move_heart')
    parser.add_argument('--random', help='random YES or NO', type=str2bool, default='1')

    args = parser.parse_args()
    source_dir = args.source_dir
    wm_dir = args.wm_dir
    text_path = args.text_path
    report_dir = args.report_dir
    type_ = args.type_
    alpha = args.alpha
    position = args.position
    print(position, type(position))
    random = args.random
    insert_type = args.insert_type


    执行脚本：
	python add_logo_text.py --source_dir C:\Users\******\Desktop\LanJunTools\datas\source_imgs --wm_dir C:\Users\******\Desktop\LanJunTools\datas\QR_logos --text_path NO --report_dir ./report_datas  --type_ img  --alpha 0.3  --position (20,20)  --insert_type move_heart --random 1

    python add_logo_text.py --source_dir C:\Users\******\Desktop\LanJunTools\datas\source_imgs --wm_dir C:\Users\******\Desktop\LanJunTools\datas\QR_logos --text_path NO --report_dir ./report_datas  --type_ img  --alpha 0.3  --position (20,20)  --insert_type move_heart --random 0
    """

    source_dir = './datas/source_imgs'
    wm_dir = './datas/QR_logos'
    text_path = None
    report_dir = './report_datas'
    type_ = 'img'
    alpha = 0.3
    position = None
    random = True
    insert_type = 'move_heart'

    obj = AddLogoText(source_imgs_dir=source_dir,
                 wm_imgs_dir=wm_dir,
                 text_path=text_path,
                 report_dir= report_dir,
                 alpha=alpha,
                 position=position,
                 random=random,
                 insert_type=insert_type,
                 type_ = type_)

    obj.run() 

本文地址：https://blog.csdn.net/qq_16555103/article/details/108856674