去除PDF中的水印

from itertools import product
import fitz
import os


# 先找出pdf水印的rgb值，然后把该值装换成空白，把每一页pdf转换成图片，再把图片转成pdf
# 会在当前目录下生成两个pdf文件，NEW1.pdf和NEW2.pdf,那个清晰用那个

# 去除pdf的水印，并把每一页pdf转化成图片
def remove_pdfwatermark():
    # 打开源pfd文件，修改为需要去除水印的pdf文件路径
    pdf_path = 'D:/test.pdf'
    pdf_file = fitz.open(pdf_path)

    # page_no 设置为0
    page_no = 0
    # page在pdf文件中遍历
    for page in pdf_file:
        # 获取每一页对应的图片pix (pix对象类似于我们上面看到的img对象，可以读取、修改它的 RGB)
        # page.get_pixmap() 这个操作是不可逆的，即能够实现从 PDF 到图片的转换，但修改图片 RGB 后无法应用到 PDF 上，只能输出为图片
        pix = page.get_pixmap()

        # 遍历图片中的宽和高，如果像素的rgb值总和大于711，就认为是水印，转换成255，255,255-->即白色
        for pos in product(range(pix.width), range(pix.height)):
            # 711是该pdf文档中水印的像素rgb值总和，不同的水印该值也不同
            if sum(pix.pixel(pos[0], pos[1])) >= 711:
                pix.set_pixel(pos[0], pos[1], (255, 255, 255))
        # 截图路径
        page_path = './images'
        if not os.path.exists(page_path):  # 判断存放图片的文件夹是否存在
            os.makedirs(page_path)  # 若图片文件夹不存在就创建
        # 保存去掉水印的截图
        pix.pil_save(f"./images/{page_no}.png", dpi=(30000, 30000))
        # 打印结果
        print(f'第 {page_no} 页去除完成')
        page_no += 1


# 去除的pdf水印的pdf页面图片，添加到pdf文件中
def pictooldpdf():
    # 水印截图所在的文件夹
    # pic_dir = input("请输入图片文件夹路径：")
    pic_dir = './images'
    if not os.path.exists(pic_dir):  # 判断存放图片的文件夹是否存在
        os.makedirs(pic_dir)  # 若图片文件夹不存在就创建

    pdf = fitz.open()
    # 图片数字文件先转换成int类型进行排序
    img_files = sorted(os.listdir(pic_dir), key=lambda x: int(str(x).split('.')[0]))
    for img in img_files:
        print(img)
        imgdoc = fitz.open(pic_dir + '/' + img)
        # 将打开后的图片转成单页pdf
        pdfbytes = imgdoc.convert_to_pdf()
        imgpdf = fitz.open("pdf", pdfbytes)
        # 将单页pdf插入到新的pdf文档中
        pdf.insert_pdf(imgpdf)
    pdf.save("NEW1.pdf")
    pdf.close()


# 去除水印后的图片，增强清晰度
def pyPDF_fitz():
    pdfPath = './NEW1.pdf'
    imagePath = './newimages'
    if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
        os.makedirs(imagePath)  # 若图片文件夹不存在就创建
    print("imagePath=" + imagePath)
    pdfDoc = fitz.open(pdfPath)
    for pg in range(pdfDoc.page_count):
        page = pdfDoc[pg]
        rotate = int(0)
        # 清晰度在这改(设为1.33333333是1056x816像素;设为2是>1584x1224像素)
        zoom_x = 2
        zoom_y = 2
        mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
        pix = page.get_pixmap(matrix=mat, alpha=False)
        if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
            os.makedirs(imagePath)  # 若图片文件夹不存在就创建
        pix.pil_save(imagePath + '/' + '%s.png' % pg)  # 将图片写入指定的文件夹内


# 去除的pdf水印的pdf页面图片，添加到pdf文件中
def pictopdf():
    # 水印截图所在的文件夹
    # pic_dir = input("请输入图片文件夹路径：")
    pic_dir = './newimages'
    if not os.path.exists(pic_dir):  # 判断存放图片的文件夹是否存在
        os.makedirs(pic_dir)  # 若图片文件夹不存在就创建

    pdf = fitz.open()
    # 图片数字文件先转换成int类型进行排序
    img_files = sorted(os.listdir(pic_dir), key=lambda x: int(str(x).split('.')[0]))
    for img in img_files:
        print(img)
        imgdoc = fitz.open(pic_dir + '/' + img)
        # 将打开后的图片转成单页pdf
        pdfbytes = imgdoc.convert_to_pdf()
        imgpdf = fitz.open("pdf", pdfbytes)
        # 将单页pdf插入到新的pdf文档中
        pdf.insert_pdf(imgpdf)
    pdf.save("NEW2.pdf")
    pdf.close()


if __name__ == '__main__':
    # 去除水印
    remove_pdfwatermark()
    # 生成初始去除水印后的pdf
    pictooldpdf()
    # pdf页面转成的图片，增强清晰度
    pyPDF_fitz()
    # 生成增强清晰度后的pdf
    pictopdf()