pdf/pic.py at main · caimaoy/pdf · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""pic.py

"""

import os
from itertools import product

import fitz


class RemoveWatermark:
    """RemoveWatermark remove watermark

    Args:
        object (_type_): _description_
    """

    @staticmethod
    def remove_pdfwatermark():
        """remove pdf watermark"""
        if not os.path.exists(r"PDFimages"):
            os.mkdir("PDFimages")
        path = os.path.join("PDF", os.listdir("PDF")[0])
        print(path)
        pdf_file = fitz.open(path)
        page_no = 0
        for page in pdf_file:
            zoom_x = 4.0  # horizontal zoom
            zomm_y = 4.0  # vertical zoom
            mat = fitz.Matrix(zoom_x, zomm_y)  # zoom factor 2 in each dimension
            pix = page.get_pixmap(matrix=mat)
            # pix = page.get_pixmap()
            for pos in product(range(pix.width), range(pix.height)):
                if sum(pix.pixel(pos[0], pos[1])) >= 600:  # 将这里的600替换成要去除的水印rgb值
                    print(pix.pixel(pos[0], pos[1]))
                    # pix.set_pixel(pos[0], pos[1], 255, 255, 255)
                    pix.set_pixel(pos[0], pos[1], (255, 255, 255, 255))
            # 将去除水印后的图片保存到根文件夹下的PDFimages文件夹中，并以页码命名
            pix.pil_save(f"./PDFimages/{page_no}.png", dpi=(30000, 30000))
            print(f"第 {page_no} 页去除完成")
            page_no += 1

    @staticmethod
    def pictopdf():
        """设置路径为去除水印后的图片保存路径"""
        pic_dir = "./PDFimages"
        pdf = fitz.open()
        img_files = sorted(os.listdir(pic_dir), key=lambda x: int(str(x).split(".", maxsplit=1)[0]))
        for img in img_files:
            print("正在合成", img)
            imgdoc = fitz.open(pic_dir + "/" + img)
            pdfbytes = imgdoc.convert_to_pdf()
            imgpdf = fitz.open("pdf", pdfbytes)
            pdf.insert_pdf(imgpdf)
        pdf.save("result.pdf")  # 保存的pdf名
        pdf.close()
        for img in os.listdir("PDFimages"):
            img = os.path.join("PDFimages", img)
            os.remove(img)


def main():
    """main"""
    if not os.path.exists(r"PDFimages"):
        os.mkdir("PDFimages")
    list_ = os.listdir("PDFimages")
    if list_:
        for img in list_:
            img = os.path.join("PDFimages", img)
            os.remove(img)
    if not os.path.exists(r"PDF"):
        os.mkdir("PDF")
    print("请将一个PDF文件放在PDF文件夹下")
    y = input("若完成上述步骤，请输入y即可去水印，去水印的结果为result.pdf文件。。。")
    if y == "y":
        obj = RemoveWatermark()
        obj.remove_pdfwatermark()
        obj.pictopdf()


if __name__ == "__main__":
    main()