import cv2, os, time
import numpy as np

IMG_DIR = "C:\\Users\\95711\\Desktop\\test\\"
SAVE_DIR = "C:\\Users\\95711\\Desktop\\save\\"
TEXT = "cyberbrain.top"
MAXSCALE = 5  # 字体放大倍数
MINSCALE = 0  # Can not reach
SCALESTEP = -0.2
SIZERATIO = 1.3 # 文件大小限制比例

#####################################
# 1、得到当前图片的h&w
# 2、从最大号开始计算，textSize
# 3、如果符合要求，则4；不符合，则减小一号
# 4、判断选定矩形的灰度均值，考虑使用白色or黑色
# 5、产生新图片，保存
# 6、如果新图片体积过大，则减小quality，直到符合要求倍数SIZERATIO
# ###################################
# 写个打水印程序自己用
# 2021.8.9
# by 无星之夜的咏叹调
#####################################
files = os.listdir(IMG_DIR)

# 遍历每一个图片文件
for index,file in enumerate(files):
    thickness = 2  # 默认字体粗2
    myFont = cv2.FONT_HERSHEY_SIMPLEX # 默认字体无衬线普通大小

    img = cv2.imread(IMG_DIR+file)
    filesize = os.path.getsize(IMG_DIR+file) #得到磁盘文件大小
    h, w, c = img.shape

    # 字体倍数从max到min，从大到小测试
    for fontscale in np.arange(MAXSCALE,MINSCALE,SCALESTEP):
        if fontscale < 1: thickness = 1 # 如果fontscale小于1了都还没有匹配，尝试减小字体粗细（并且会好看些）
        textSize, baseline = cv2.getTextSize(TEXT, myFont, fontscale, thickness)
        ratio = textSize[0] / w
        # print("ratio",ratio)
        # 如果图片过大，则直接使用允许的最大倍数
        if fontscale == MAXSCALE and ratio < 1/3 :
            break
        # 找到合适大小的缩放倍数了
        if 1/3 < ratio < 1/2:
            break
    # 如果图片过小，没办法，换个小字体强行加吧 QAQ
    else:
        myFont = cv2.FONT_HERSHEY_PLAIN
        textSize, baseline = cv2.getTextSize(TEXT, myFont, fontscale, thickness)
        ratio = textSize[0] / w
        
    # 切割得到原图中即将写字的部分
    putx, puty = w - textSize[0], h - textSize[1] - baseline 
    timg = img[puty:h,putx:w]     # (0,0)是左上角
    gray=cv2.cvtColor(timg,cv2.COLOR_BGR2GRAY)

    # 如果图片较暗则使用白色；如果图片较亮则使用黑色
    if gray.mean() < 127:
        colorfg, colorbg = (255,255,255),(0,0,0)
    else:
        colorbg, colorfg = (255,255,255),(0,0,0)

    # 打印水印，字体粗的为阴影，增加对比度
    cv2.putText(img,TEXT,(putx,h-baseline),myFont,fontscale,colorbg,thickness+2)
    cv2.putText(img,TEXT,(putx,h-baseline),myFont,fontscale,colorfg,thickness)

    # 判断图片类型，仅支持png和jpg
    if file.endswith(".jpg"):
        quaType = cv2.IMWRITE_JPEG_QUALITY
        scale = (100,0,-1)
    elif file.endswith(".png"):
        quaType = cv2.IMWRITE_PNG_COMPRESSION
        scale = (0,10,1)

    # 保存图片，并限制文件大小
    for quality in range(*scale):
        time.sleep(0.001)
        if not cv2.imwrite(SAVE_DIR+file,img,[quaType,quality]):
            print("%d:%s   imwrite Error! quality:%d" % (index+1,file,quality)) 
            break
        time.sleep(0.005) # 好像延时一下可以减少错误发生，因为windows文件写入较慢
        newfilesize = os.path.getsize(SAVE_DIR+file)
        sizeRatio = float(newfilesize/filesize)
        # print("#TEMP:%d:%s  \t quality:%d, sizeratio:%f" % (index+1,file,quality,sizeRatio))  # DEBUG
        # 最大允许被放大SIZERATIO倍
        if sizeRatio < SIZERATIO:
            print("%d:%s  \tFinish! quality:%d, sizeRatio:%f" % (index+1,file,quality,sizeRatio))
            break