You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

90 lines
2.9 KiB
Python

import os
import re
import subprocess
from table_ocr.util import get_logger, working_dir
logger = get_logger(__name__)
# Wrapper around the Poppler command line utility "pdfimages" and helpers for
# finding the output files of that command.
def pdf_to_images(pdf_filepath):
"""
Turn a pdf into images
Returns the filenames of the created images sorted lexicographically.
"""
directory, filename = os.path.split(pdf_filepath)
image_filenames = pdfimages(pdf_filepath)
# Since pdfimages creates a number of files named each for there page number
# and doesn't return us the list that it created
return sorted([os.path.join(directory, f) for f in image_filenames])
def pdfimages(pdf_filepath):
"""
Uses the `pdfimages` utility from Poppler
(https://poppler.freedesktop.org/). Creates images out of each page. Images
are prefixed by their name sans extension and suffixed by their page number.
This should work up to pdfs with 999 pages since find matching files in dir
uses 3 digits in its regex.
"""
directory, filename = os.path.split(pdf_filepath)
if not os.path.isabs(directory):
directory = os.path.abspath(directory)
filename_sans_ext = filename.split(".pdf")[0]
# pdfimages outputs results to the current working directory
with working_dir(directory):
subprocess.run(["pdfimages", "-png", filename, filename.split(".pdf")[0]])
image_filenames = find_matching_files_in_dir(filename_sans_ext, directory)
logger.debug(
"Converted {} into files:\n{}".format(pdf_filepath, "\n".join(image_filenames))
)
return image_filenames
def find_matching_files_in_dir(file_prefix, directory):
files = [
filename
for filename in os.listdir(directory)
if re.match(r"{}-\d{{3}}.*\.png".format(re.escape(file_prefix)), filename)
]
return files
def preprocess_img(filepath, tess_params=None):
"""Processing that involves running shell executables,
like mogrify to rotate.
Uses tesseract to detect rotation.
Orientation and script detection is only available for legacy tesseract
(--oem 0). Some versions of tesseract will segfault if you let it run OSD
with the default oem (3).
"""
if tess_params is None:
tess_params = ["--psm", "0", "--oem", "0"]
rotate = get_rotate(filepath, tess_params)
logger.debug("Rotating {} by {}.".format(filepath, rotate))
mogrify(filepath, rotate)
def get_rotate(image_filepath, tess_params):
"""
"""
tess_command = ["tesseract"] + tess_params + [image_filepath, "-"]
output = (
subprocess.check_output(tess_command)
.decode("utf-8")
.split("\n")
)
output = next(l for l in output if "Rotate: " in l)
output = output.split(": ")[1]
return output
def mogrify(image_filepath, rotate):
subprocess.run(["mogrify", "-rotate", rotate, image_filepath])