Python实现图像空域随机水印加入与提取

本文要点在于Python扩展库pillow中Image类的运用。图像空域随机水印的主要思路在于：在原始图像中随机选取一些位置替换为水印图片中的非背景像素，同时生成日志文件记录替换的位置和水印中像素位置的对应关系，然后可以根据加入水印的图片和日志文件来提取和验证水印。

from os import remove

from os.path import isfile 

from random import sample, choice

from PIL import Image

def mergeWaterMark(originPic, watermarkPic, logTxt):

    #原始图片和水印文件必须为图片格式

    if (not originPic.endswith(('.jpg', '.bmp', '.png'))) or (not watermarkPic.endswith(('.jpg', '.bmp', '.png'))):

        return 'Error format.'

    

    #打开原图和水印图片，并获取大小

    imOrigin = Image.open(originPic)

    originWidth, originHeight = imOrigin.size

    imWaterMark = Image.open(watermarkPic)

    watermarkWidth, watermarkHeight = imWaterMark.size

    #随机生成水印位置

    allPositions = [(w,h) for w in range(originWidth) for h in range(originHeight)]

    positions = sample(allPositions, watermarkWidth*watermarkHeight)

    fpLog = open(logTxt, 'w')

    #写入水印文件大小

    fpLog.write(str((watermarkWidth,watermarkHeight))+'\n')

    

    for w in range(watermarkWidth):

        for h in range(watermarkHeight):

            c = imWaterMark.getpixel((w,h))

            c = c[:3]

            #只写入不是白色的像素

            if c != (255,255,255):

                p = choice(positions)

                #写入像素值

                imOrigin.putpixel(p, c)

                #避免重复修改同一个像素

                positions.remove(p)

                #生成日志文件，用来提取水印

                fpLog.write(str(p+(w,h))+'\n')

    fpLog.close()

    #生成加入水印的新图片

    imOrigin.save(originPic[:-4]+'_new'+originPic[-4:])

def restoreWaterMark(mergedPic, logTxt, watermarkPic):

    #删除原来提取过的水印文件

    if isfile(watermarkPic):

        remove(watermarkPic)

    imMerged = Image.open(mergedPic)

    with open(logTxt) as fp:

        for line in fp:

            #读取每一行并还原为元组

            line = eval(line.strip())

            #第一行是水印图片尺寸，先创建水印文件

            if len(line)==2:

                imWaterMark = Image.new('RGB', line, (255,255,255))

            else:

                #提取水印像素并写入水印文件

                c = imMerged.getpixel((line[0],line[1]))

                c = c[:3]

                imWaterMark.putpixel((line[2],line[3]), c)

    #保存提取的水印

    imWaterMark.save(watermarkPic)

#测试

mergeWaterMark('origin.bmp', 'watermark.png', 'logg.txt')

restoreWaterMark('origin_new.bmp', 'logg.txt', 'restoredWaterMark.png')

测试用的原始图片：

测试用的水印图片：

加入水印以后的图片：

提取出来的水印：