library/util/convert2img.py

#!/usr/bin/python3.6
# -*- coding: utf-8 -*-
# @Time    : 2021/3/17 13:13
# @Author  : lijunliang
# @Email   : lijunliang@topnet.net.cn
# @File    : convert2img.py
# @Software: PyCharm
import fitz
import os
from loguru import logger
import cv2 as cv
import numpy as np


def create_border(image):
    """
    修复尺寸
    :param image:
    :return:
    """
    img_height, img_width, _ = image.shape
    add_w_l, add_h_u = int((248 - img_height) / 2), int((200 - img_width) / 2)
    image = cv.copyMakeBorder(image, add_w_l, add_w_l, add_h_u, add_h_u, cv.BORDER_CONSTANT, value=[255, 255, 255, 255])
    return image


def convert_img(pdf_path: str) -> str:
    """
    pdf首页转图片
    :param pdf_path:pdf文件路径
    :return:
    """
    try:
        doc = fitz.open(pdf_path)
        pdf_name = os.path.splitext(pdf_path)[0]
        if doc.page_count < 1:
            return ""
        page = doc[0]
        rotate = int(0)
        # 每个尺寸的缩放系数为2，这将为我们生成分辨率提高四倍的图像。
        zoom_x = 2.0
        zoom_y = 2.0
        trans = fitz.Matrix(zoom_x, zoom_y).preRotate(rotate)
        pm = page.getPixmap(matrix=trans, alpha=False)
        image_path = '%s.png' % pdf_name
        pm.writePNG(image_path)
        image_path = image_zoom(image_path)
        return image_path
    except Exception as e:
        logger.warning("生成首页失败---》%s" % e)
        return ""


def image_zoom(img_path: str):
    """
    图片缩放
    :param img_path: 图片路径
    :return:
    """
    image = cv.imread(img_path)
    img_height, img_width, _ = image.shape
    if (img_height / img_width) > (248 / 200):
        rate = 248 / img_height
    else:
        rate = 200 / img_width

    resize_image = cv.resize(image, (int(img_width * rate), int(img_height * rate)),
                             interpolation=cv.INTER_AREA)
    new_image = create_border(resize_image)  # 修复尺寸
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)  # 定义一个核
    dst = cv.filter2D(new_image, -1, kernel=kernel)
    cv.imwrite(img_path, dst)  # 保存
    return img_path


if __name__ == '__main__':
    file_path="../folder/071efc233bac58e9b655f00961a3c63728d15a71.pdf"
    convert_img(file_path)